lxy-9602 commented on code in PR #194:
URL: https://github.com/apache/paimon-cpp/pull/194#discussion_r3776030203
##########
src/paimon/core/operation/raw_file_split_read.h:
##########
@@ -65,16 +64,26 @@ class RawFileSplitRead : public AbstractSplitRead {
const std::shared_ptr<MemoryPool>& memory_pool,
const std::shared_ptr<Executor>& executor);
+ /// Also accepts an `IndexedSplit` over a single-file data split, in which
case its row
+ /// ranges narrow the read to the given file-local physical positions.
Result<std::unique_ptr<BatchReader>> CreateReader(const
std::shared_ptr<Split>& split) override;
Result<std::unique_ptr<BatchReader>> CreateReader(
const BinaryRow& partition, int32_t bucket,
const std::vector<std::shared_ptr<DataFileMeta>>& files,
const std::vector<std::optional<DeletionFile>>& deletion_files);
+ /// Reads with an optional selection of file-local row positions. The
ranges apply to
+ /// every file of the split, so callers pass them only for single-file
splits.
Result<std::unique_ptr<BatchReader>> CreateReader(
const BinaryRow& partition, int32_t bucket,
const std::vector<std::shared_ptr<DataFileMeta>>& files,
- DeletionVector::Factory dv_factory);
+ const std::vector<std::optional<DeletionFile>>& deletion_files,
+ const std::optional<std::vector<Range>>& local_row_ranges);
+
+ Result<std::unique_ptr<BatchReader>> CreateReader(
+ const BinaryRow& partition, int32_t bucket,
+ const std::vector<std::shared_ptr<DataFileMeta>>& files,
DeletionVector::Factory dv_factory,
+ const std::optional<std::vector<Range>>& local_row_ranges =
std::nullopt);
Review Comment:
Please avoid using default arguments in production code.
##########
src/paimon/core/index/pksorted/pk_sorted_index_file.cpp:
##########
@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/index/pksorted/pk_sorted_index_file.h"
+
+#include <cstddef>
+#include <optional>
+#include <utility>
+
+#include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
+#include "fmt/format.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/path_util.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/index/pk/primary_key_index_source_meta.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_index_writer.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+
+namespace paimon {
+Result<std::shared_ptr<IndexFileMeta>> PkSortedIndexFile::Build(
+ const DataField& field, const std::string& index_type,
+ const std::map<std::string, std::string>& options, int32_t data_level,
+ const std::vector<PrimaryKeyIndexSourceFile>& source_files,
+ const std::shared_ptr<arrow::Array>& sorted_values, std::vector<int64_t>
sorted_ordinals,
+ const std::shared_ptr<GlobalIndexFileWriter>& file_writer, bool
is_external_path,
+ const std::shared_ptr<MemoryPool>& pool) {
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta,
Review Comment:
The current implementation puts all planned sorted data for an entire index
file into a single `sorted_values`, which could lead to excessive memory usage.
While Java uses an external sort buffer here.
Given the scope of the current PR, I’d suggest adding a TODO to clearly
document this as a known issue and plan to fix it in a follow-up PR.
##########
src/paimon/core/table/source/primary_key_sorted_index_result.cpp:
##########
@@ -0,0 +1,146 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+
+#include <limits>
+#include <map>
+#include <set>
+#include <utility>
+
+#include "paimon/core/global_index/indexed_split_impl.h"
+#include "paimon/core/table/source/deletion_file.h"
+
+namespace paimon {
+namespace {
+struct RangeConversion {
+ bool use_index;
+ std::vector<Range> ranges;
+};
+
+Result<std::shared_ptr<DataSplitImpl>> ToSingleFileSplit(
+ const PrimaryKeySortedIndexScan::FilePlan& file) {
+ const std::shared_ptr<DataSplitImpl>& source = file.SourceSplit();
+ std::vector<std::shared_ptr<DataFileMeta>> data_files{file.DataFile()};
+ DataSplitImpl::Builder builder(source->Partition(), source->Bucket(),
source->BucketPath(),
+ std::move(data_files));
+ builder.WithSnapshot(source->SnapshotId())
+ .WithTotalBuckets(source->TotalBuckets())
+ .IsStreaming(false)
+ .RawConvertible(source->RawConvertible());
Review Comment:
It seems in Java, `rawConvertible` is always `false` here. The current
implementation would make the scan results inconsistent with Java.
##########
src/paimon/core/table/source/primary_key_sorted_index_scan.h:
##########
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "paimon/core/index/pk/primary_key_index_definition.h"
+#include "paimon/core/index/pksorted/pk_sorted_index_group.h"
+#include "paimon/core/manifest/index_manifest_entry.h"
+#include "paimon/core/schema/table_schema.h"
+#include "paimon/core/table/source/data_split_impl.h"
+#include "paimon/core/utils/index_file_path_factories.h"
+#include "paimon/fs/file_system.h"
+#include "paimon/global_index/global_index_reader.h"
+#include "paimon/global_index/global_index_result.h"
+#include "paimon/memory/memory_pool.h"
+#include "paimon/predicate/predicate.h"
+#include "paimon/result.h"
+
+namespace paimon {
+class Executor;
+
+/// Plans and evaluates source-backed primary-key scalar index groups in
file-local
+/// row-position space.
+///
+/// The scan works on one captured snapshot: data splits and index manifest
entries must
+/// come from the same snapshot. Every active data file is associated with the
validated
+/// payload group of its (bucket, field, data level); files without a valid
group keep an
+/// empty group map and later fall back to a normal scan. Evaluation runs the
predicate
+/// against the group payloads once per group and query, then localizes group
ordinals to
+/// per-file physical row positions using the ordered source row-count prefix.
+class PrimaryKeySortedIndexScan {
+ public:
+ PrimaryKeySortedIndexScan() = delete;
+ ~PrimaryKeySortedIndexScan() = delete;
+
+ /// One active data file and its complete field-local payload groups.
+ class FilePlan {
+ public:
+ FilePlan(std::shared_ptr<DataSplitImpl> source_split, int32_t
file_index,
+ std::map<int32_t, std::shared_ptr<PkSortedIndexGroup>> groups)
+ : source_split_(std::move(source_split)),
Review Comment:
`PkSortedIndexGroup::Create` returns `optional<PkSortedIndexGroup>`, but the
call sites seem to use `std::shared_ptr<PkSortedIndexGroup>`. Should we adjust
the return type of `Create` for consistency?
##########
src/paimon/core/table/source/primary_key_index_batch_scan.cpp:
##########
@@ -0,0 +1,298 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_index_batch_scan.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+#include <set>
+#include <string>
+#include <thread>
+
+#include "fmt/format.h"
+#include "paimon/core/index/index_file_handler.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/table/source/plan_impl.h"
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+#include "paimon/core/table/source/snapshot/snapshot_reader.h"
+#include "paimon/core/utils/index_file_path_factories.h"
+#include "paimon/core/utils/snapshot_manager.h"
+#include "paimon/executor.h"
+#include "paimon/predicate/compound_predicate.h"
+#include "paimon/predicate/leaf_predicate.h"
+#include "paimon/predicate/predicate_builder.h"
+
+namespace paimon {
+namespace {
+Result<std::shared_ptr<Executor>> CreateGlobalIndexExecutor(const CoreOptions&
core_options) {
+ uint32_t thread_num = std::thread::hardware_concurrency();
+ std::optional<int32_t> configured_thread_num =
core_options.GetGlobalIndexThreadNum();
+ if (configured_thread_num) {
+ if (configured_thread_num.value() <= 0) {
+ return Status::Invalid(fmt::format("invalid global index thread
number {}",
+ configured_thread_num.value()));
+ }
+ thread_num = static_cast<uint32_t>(configured_thread_num.value());
+ } else if (thread_num == 0) {
+ thread_num = 1;
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Executor> executor,
CreateDefaultExecutor(thread_num));
+ return executor;
+}
+
+/// Restricts a predicate to leaves over the indexed fields: an AND keeps its
convertible
+/// children, an OR is only kept when every child is convertible, and
everything else is
+/// dropped. A null return means no part of the predicate can use the index.
+Result<std::shared_ptr<Predicate>> ProjectToIndexedFields(
+ const std::shared_ptr<Predicate>& predicate, const std::set<std::string>&
indexed_fields) {
+ if (predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (auto leaf_predicate =
std::dynamic_pointer_cast<LeafPredicate>(predicate)) {
+ if (indexed_fields.count(leaf_predicate->FieldName()) > 0) {
+ return predicate;
+ }
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ auto compound_predicate =
std::dynamic_pointer_cast<CompoundPredicate>(predicate);
+ if (compound_predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ bool is_and = compound_predicate->GetFunction().GetType() ==
Function::Type::AND;
+ bool is_or = compound_predicate->GetFunction().GetType() ==
Function::Type::OR;
+ if (!is_and && !is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ std::vector<std::shared_ptr<Predicate>> converted_children;
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> converted_child,
+ ProjectToIndexedFields(child, indexed_fields));
+ if (converted_child != nullptr) {
+ converted_children.push_back(std::move(converted_child));
+ } else if (is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ }
+ if (converted_children.empty()) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (converted_children.size() == 1) {
+ return converted_children[0];
+ }
+ if (is_and) {
+ return PredicateBuilder::And(converted_children);
+ }
+ return PredicateBuilder::Or(converted_children);
+}
+
+void FlattenChildren(const std::shared_ptr<CompoundPredicate>&
compound_predicate,
+ std::vector<std::shared_ptr<Predicate>>* flattened) {
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ auto compound_child =
std::dynamic_pointer_cast<CompoundPredicate>(child);
+ if (compound_child != nullptr &&
compound_child->GetFunction().GetType() ==
+
compound_predicate->GetFunction().GetType()) {
+ FlattenChildren(compound_child, flattened);
+ } else {
+ flattened->push_back(child);
+ }
+ }
+}
+
+/// A predicate is null-rejecting when it cannot match a row whose tested
field is null.
+/// Under SQL three-valued logic every comparison and match predicate rejects
null; only
+/// IS NULL accepts it, and IS NOT NULL is the predicate being pruned.
+bool IsNullRejecting(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ if (leaf_predicate == nullptr) {
+ return false;
+ }
+ switch (leaf_predicate->GetFunction().GetType()) {
+ case Function::Type::EQUAL:
+ case Function::Type::NOT_EQUAL:
+ case Function::Type::GREATER_THAN:
+ case Function::Type::GREATER_OR_EQUAL:
+ case Function::Type::LESS_THAN:
+ case Function::Type::LESS_OR_EQUAL:
+ case Function::Type::IN:
+ case Function::Type::NOT_IN:
+ case Function::Type::STARTS_WITH:
+ case Function::Type::ENDS_WITH:
+ case Function::Type::CONTAINS:
+ case Function::Type::LIKE:
+ return true;
+ default:
+ return false;
+ }
+}
+
+bool IsIsNotNull(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ return leaf_predicate != nullptr &&
+ leaf_predicate->GetFunction().GetType() ==
Function::Type::IS_NOT_NULL;
+}
+
+/// Flattens nested same-function compounds and, inside an AND, removes `f IS
NOT NULL`
+/// leaves made redundant by a null-rejecting sibling on the same field.
Pruning must not
+/// consider `f IS NULL` as constraining: dropping IS NOT NULL from
+/// "f IS NULL AND f IS NOT NULL" would turn the empty result into the set of
null rows.
+Result<std::shared_ptr<Predicate>> NormalizePredicate(const
std::shared_ptr<Predicate>& predicate) {
Review Comment:
I’d suggest moving some of the complex func into `.h` as `static` methods so
it can be tested at a finer granularity.
##########
src/paimon/core/operation/raw_file_split_read.cpp:
##########
@@ -64,6 +67,21 @@ RawFileSplitRead::RawFileSplitRead(const
std::shared_ptr<FileStorePathFactory>&
Result<std::unique_ptr<BatchReader>> RawFileSplitRead::CreateReader(
const std::shared_ptr<Split>& split) {
+ if (auto indexed_split =
std::dynamic_pointer_cast<IndexedSplitImpl>(split)) {
+ PAIMON_RETURN_NOT_OK(indexed_split->Validate());
+ const std::shared_ptr<DataSplit>& inner_split =
indexed_split->GetDataSplit();
+ auto inner_split_impl =
std::dynamic_pointer_cast<DataSplitImpl>(inner_split);
+ if (!inner_split_impl) {
+ return Status::Invalid("cannot cast indexed inner split to
data_split");
+ }
+ if (inner_split_impl->DataFiles().size() != 1) {
+ return Status::Invalid(
+ "indexed splits with file-local row ranges must contain
exactly one file");
+ }
Review Comment:
If the indexed split contains scores, is read supported right now? I don’t
see any handling logic for that. If it’s not supported, please fail fast and
add a TODO to mark it clearly.
##########
src/paimon/core/index/pksorted/pk_sorted_index_file.cpp:
##########
@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/index/pksorted/pk_sorted_index_file.h"
+
+#include <cstddef>
+#include <optional>
+#include <utility>
+
+#include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
+#include "fmt/format.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/path_util.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/index/pk/primary_key_index_source_meta.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_index_writer.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+
+namespace paimon {
+Result<std::shared_ptr<IndexFileMeta>> PkSortedIndexFile::Build(
+ const DataField& field, const std::string& index_type,
+ const std::map<std::string, std::string>& options, int32_t data_level,
+ const std::vector<PrimaryKeyIndexSourceFile>& source_files,
+ const std::shared_ptr<arrow::Array>& sorted_values, std::vector<int64_t>
sorted_ordinals,
+ const std::shared_ptr<GlobalIndexFileWriter>& file_writer, bool
is_external_path,
+ const std::shared_ptr<MemoryPool>& pool) {
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta,
+ PrimaryKeyIndexSourceMeta::Create(data_level,
source_files));
+ int64_t source_row_count = 0;
+ for (const PrimaryKeyIndexSourceFile& source_file : source_files) {
+ if (__builtin_add_overflow(source_row_count, source_file.row_count,
&source_row_count)) {
+ return Status::Invalid("Source row count overflows in sorted index
build.");
+ }
+ }
+ if (source_row_count <= 0) {
+ return Status::Invalid("A sorted index group must reference at least
one source row.");
+ }
+ if (sorted_values == nullptr || sorted_values->length() !=
source_row_count ||
+ static_cast<int64_t>(sorted_ordinals.size()) != source_row_count) {
+ return Status::Invalid(
+ fmt::format("Sorted index input row count {} does not match source
row count {}.",
+ sorted_values == nullptr ? 0 :
sorted_values->length(), source_row_count));
+ }
+ std::vector<bool> seen_ordinals(static_cast<size_t>(source_row_count),
false);
+ for (int64_t ordinal : sorted_ordinals) {
+ if (ordinal < 0 || ordinal >= source_row_count) {
+ return Status::Invalid(
+ fmt::format("Row id {} is outside sorted index group row range
[0, {}).", ordinal,
+ source_row_count));
+ }
+ if (seen_ordinals[ordinal]) {
+ return Status::Invalid(fmt::format("Row id {} appears more than
once.", ordinal));
+ }
+ seen_ordinals[ordinal] = true;
+ }
+
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<GlobalIndexer> indexer,
+ GlobalIndexerFactory::Get(index_type, options));
+ if (indexer == nullptr) {
+ return Status::Invalid(fmt::format("Index type {} is not registered.",
index_type));
+ }
+ auto arrow_field = DataField::ConvertDataFieldToArrowField(field);
+ auto arrow_schema = arrow::schema({arrow_field});
+ ArrowSchema c_arrow_schema;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema,
&c_arrow_schema));
+ ScopeGuard schema_guard([&]() { ArrowSchemaRelease(&c_arrow_schema); });
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<GlobalIndexWriter> writer,
+ indexer->CreateWriter(field.Name(),
&c_arrow_schema, file_writer, pool));
+
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray>
struct_array,
+
arrow::StructArray::Make({sorted_values}, {field.Name()}));
+ ::ArrowArray c_array;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array,
&c_array));
+ ScopeGuard array_guard([&]() { ArrowArrayRelease(&c_array); });
+ PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array,
std::move(sorted_ordinals)));
+ PAIMON_ASSIGN_OR_RAISE(std::vector<GlobalIndexIOMeta> io_metas,
writer->Finish());
+ if (io_metas.size() != 1) {
+ return Status::Invalid(fmt::format(
+ "Sorted index build must produce exactly one payload file, but
produced {}.",
+ io_metas.size()));
+ }
+ const GlobalIndexIOMeta& io_meta = io_metas[0];
+
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Bytes> source_meta_bytes,
source_meta.Serialize(pool));
+ std::optional<std::string> external_path;
+ if (is_external_path) {
+ external_path = io_meta.file_path;
+ }
Review Comment:
Please change this to something like:
```cpp
if (is_external_path) {
PAIMON_ASSIGN_OR_RAISE(Path path, PathUtil::ToPath(io_meta.file_path));
external_path = path.ToString();
}
```
Could we normalize the path here? Regular global index writing does this in
`global_index_write_task.cpp`. In Java, `org.apache.paimon.fs.Path` parses and
normalizes the URI during construction, so the final stored value is the
normalized form from `Path.toString()`.
##########
src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp:
##########
@@ -0,0 +1,548 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "fmt/format.h"
+#include "gtest/gtest.h"
+#include "paimon/core/global_index/indexed_split_impl.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/index/pksorted/pk_sorted_index_file.h"
+#include "paimon/core/manifest/file_kind.h"
+#include "paimon/core/manifest/file_source.h"
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+#include "paimon/global_index/bitmap_global_index_result.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+#include "paimon/global_index/io/global_index_file_reader.h"
+#include "paimon/global_index/io/global_index_file_writer.h"
+#include "paimon/predicate/predicate_builder.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+namespace {
+constexpr int32_t kPriceFieldId = 1;
+constexpr int64_t kSnapshotId = 7;
+constexpr int64_t kFileARows = 100;
+constexpr int64_t kFileBRows = 200;
+constexpr int64_t kTotalRows = kFileARows + kFileBRows;
+
+class TestGlobalIndexFileWriter : public GlobalIndexFileWriter {
+ public:
+ TestGlobalIndexFileWriter(const std::shared_ptr<FileSystem>& fs, const
std::string& base_path)
+ : fs_(fs), base_path_(base_path) {}
+
+ Result<std::string> NewFileName(const std::string& prefix) const override {
+ return fmt::format("{}-index-{}", prefix, file_counter_++);
+ }
+
+ Result<std::unique_ptr<OutputStream>> NewOutputStream(
+ const std::string& file_name) const override {
+ return fs_->Create(base_path_ + "/" + file_name, true);
+ }
+
+ Result<int64_t> GetFileSize(const std::string& file_name) const override {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FileStatus> file_status,
+ fs_->GetFileStatus(base_path_ + "/" +
file_name));
+ return file_status->GetLen();
+ }
+
+ std::string ToPath(const std::string& file_name) const override {
+ return base_path_ + "/" + file_name;
+ }
+
+ private:
+ std::shared_ptr<FileSystem> fs_;
+ std::string base_path_;
+ mutable int64_t file_counter_ = 0;
+};
+
+class TestGlobalIndexFileReader : public GlobalIndexFileReader {
+ public:
+ explicit TestGlobalIndexFileReader(const std::shared_ptr<FileSystem>& fs)
: fs_(fs) {}
+
+ Result<std::unique_ptr<InputStream>> GetInputStream(
+ const std::string& file_path) const override {
+ return fs_->Open(file_path);
+ }
+
+ private:
+ std::shared_ptr<FileSystem> fs_;
+};
+
+/// A reader stub whose equality result is fully controlled by the test, used
to exercise
+/// the untrusted-position fallbacks.
+class StubGlobalIndexReader : public GlobalIndexReader {
+ public:
+ explicit StubGlobalIndexReader(RoaringBitmap64 equal_result)
+ : equal_result_(std::move(equal_result)) {}
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNotNull() override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNull() override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEqual(const Literal&
literal) override {
+ RoaringBitmap64 copy = equal_result_;
+ return std::make_shared<BitmapGlobalIndexResult>(
+ [bitmap = std::move(copy)]() -> Result<RoaringBitmap64> { return
bitmap; });
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotEqual(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessThan(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessOrEqual(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterThan(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterOrEqual(
+ const Literal& literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIn(
+ const std::vector<Literal>& literals) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotIn(
+ const std::vector<Literal>& literals) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitStartsWith(const Literal&
prefix) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEndsWith(const Literal&
suffix) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitContains(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLike(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<ScoredGlobalIndexResult>> VisitVectorSearch(
+ const std::shared_ptr<VectorSearch>& vector_search) override {
+ return Status::Invalid("not supported");
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitFullTextSearch(
+ const std::shared_ptr<FullTextSearch>& full_text_search) override {
+ return Status::Invalid("not supported");
+ }
+ bool IsThreadSafe() const override {
+ return false;
+ }
+ std::string GetIndexType() const override {
+ return "btree";
+ }
+
+ private:
+ static Result<std::shared_ptr<GlobalIndexResult>> NotEvaluable() {
+ return std::shared_ptr<GlobalIndexResult>(nullptr);
+ }
+
+ RoaringBitmap64 equal_result_;
+};
+} // namespace
+
+class PrimaryKeySortedIndexScanTest : public ::testing::Test {
+ protected:
+ void SetUp() override {
+ pool_ = GetDefaultPool();
+ test_dir_ = UniqueTestDirectory::Create("local");
+ fs_ = test_dir_->GetFileSystem();
+ base_path_ = test_dir_->Str();
+
+ std::vector<DataField> fields = {
+ DataField(0, arrow::field("id", arrow::int64())),
+ DataField(kPriceFieldId, arrow::field("price", arrow::int64())),
+ DataField(2, arrow::field("status", arrow::utf8())),
+ };
+ std::map<std::string, std::string> options =
{{"pk-btree.index.columns", "price"}};
+ table_schema_ = std::make_shared<TableSchema>(
+ /*version=*/3, /*id=*/0, fields, /*highest_field_id=*/2,
+ /*partition_keys=*/std::vector<std::string>(),
+ /*primary_keys=*/std::vector<std::string>{"id"}, options,
+ /*comment=*/std::nullopt, /*time_millis=*/0);
+ ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions,
+
PrimaryKeyIndexDefinitions::Create(*table_schema_));
+ definitions_ = definitions.ScalarDefinitions();
+ ASSERT_EQ(definitions_.size(), 1);
+ }
+
+ std::shared_ptr<DataFileMeta> MakeDataFile(const std::string& name,
int64_t row_count,
+ int32_t level, const
FileSource& file_source) {
+ return std::make_shared<DataFileMeta>(
+ name, /*file_size=*/1024, row_count,
+ /*min_key=*/BinaryRow::EmptyRow(),
/*max_key=*/BinaryRow::EmptyRow(),
+ /*key_stats=*/SimpleStats::EmptyStats(),
/*value_stats=*/SimpleStats::EmptyStats(),
+ /*min_sequence_number=*/0, /*max_sequence_number=*/row_count,
/*schema_id=*/0, level,
+ /*extra_files=*/std::vector<std::optional<std::string>>(),
+ /*creation_time=*/Timestamp(1721643142456LL, 0),
+ /*delete_row_count=*/0, /*embedded_index=*/nullptr, file_source,
+ /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt,
+ /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt);
+ }
+
+ Result<std::shared_ptr<IndexFileMeta>> BuildPayload(std::vector<int64_t>
ordinals) {
+ std::vector<PrimaryKeyIndexSourceFile> source_files = {{"a.parquet",
kFileARows},
+ {"b.parquet",
kFileBRows}};
+ arrow::Int64Builder values_builder;
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Append(2 * i));
+ }
+ std::shared_ptr<arrow::Array> sorted_values;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Finish(&sorted_values));
+ PAIMON_ASSIGN_OR_RAISE(DataField field,
table_schema_->GetField(kPriceFieldId));
+ auto file_writer = std::make_shared<TestGlobalIndexFileWriter>(fs_,
base_path_);
+ return PkSortedIndexFile::Build(field, "btree",
definitions_[0].Options(),
+ /*data_level=*/5, source_files,
sorted_values,
+ std::move(ordinals), file_writer,
+ /*is_external_path=*/false, pool_);
+ }
+
+ /// Builds the standard payload of this fixture: sources a.parquet(100) +
b.parquet(200)
+ /// on level 5, indexed value at group ordinal `i` is `2 * i`.
+ Result<std::shared_ptr<IndexFileMeta>> BuildPayload() {
+ std::vector<int64_t> ordinals;
+ ordinals.reserve(kTotalRows);
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ ordinals.push_back(i);
+ }
+ return BuildPayload(std::move(ordinals));
+ }
+
+ std::shared_ptr<DataSplitImpl> MakeSplit(
+ const std::vector<std::shared_ptr<DataFileMeta>>& files, bool
raw_convertible,
+ const std::vector<std::optional<DeletionFile>>& deletion_files = {}) {
+ std::vector<std::shared_ptr<DataFileMeta>> data_files = files;
+ DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0,
+ base_path_ + "/bucket-0",
std::move(data_files));
+
builder.WithSnapshot(kSnapshotId).IsStreaming(false).RawConvertible(raw_convertible);
+ if (!deletion_files.empty()) {
+ builder.WithDataDeletionFiles(deletion_files);
+ }
+ EXPECT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> split,
builder.Build());
+ return split;
+ }
+
+ std::vector<IndexManifestEntry> MakeEntries(const
std::shared_ptr<IndexFileMeta>& payload) {
+ return {IndexManifestEntry(FileKind::Add(), BinaryRow::EmptyRow(),
/*bucket=*/0, payload)};
+ }
+
+ PrimaryKeySortedIndexScan::ReaderFactory PayloadReaderFactory() {
+ std::shared_ptr<FileSystem> fs = fs_;
+ std::string base_path = base_path_;
+ std::shared_ptr<TableSchema> table_schema = table_schema_;
+ std::shared_ptr<MemoryPool> pool = pool_;
+ return [fs, base_path, table_schema, pool](
+ const PrimaryKeySortedIndexScan::FilePlan& file,
+ const PrimaryKeyIndexDefinition& definition,
+ const PkSortedIndexGroup& group) ->
Result<std::shared_ptr<GlobalIndexReader>> {
+ PAIMON_ASSIGN_OR_RAISE(
+ std::unique_ptr<GlobalIndexer> indexer,
+ GlobalIndexerFactory::Get(definition.IndexType(),
definition.Options()));
+ if (indexer == nullptr) {
+ return Status::Invalid("btree indexer is not registered");
+ }
+ const std::shared_ptr<IndexFileMeta>& payload = group.Payload();
+ std::vector<GlobalIndexIOMeta> io_metas;
+ io_metas.emplace_back(base_path + "/" + payload->FileName(),
payload->FileSize(),
+
payload->GetGlobalIndexMeta().value().index_meta);
+ PAIMON_ASSIGN_OR_RAISE(DataField field,
table_schema->GetField(definition.FieldId()));
+ auto arrow_field = DataField::ConvertDataFieldToArrowField(field);
+ auto arrow_schema = arrow::schema({arrow_field});
+ ArrowSchema c_arrow_schema;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema,
&c_arrow_schema));
+ auto file_reader = std::make_shared<TestGlobalIndexFileReader>(fs);
+ return indexer->CreateReader(&c_arrow_schema, file_reader,
io_metas, pool);
+ };
+ }
+
+ Result<std::vector<std::shared_ptr<Split>>> PlanEvaluateConvert(
+ const std::vector<std::shared_ptr<DataSplitImpl>>& splits,
+ const std::vector<IndexManifestEntry>& entries, const
std::shared_ptr<Predicate>& predicate,
+ const PrimaryKeySortedIndexScan::ReaderFactory& reader_factory) {
+ PAIMON_ASSIGN_OR_RAISE(
+ PrimaryKeySortedIndexScan::Plan plan,
+ PrimaryKeySortedIndexScan::CreatePlan(kSnapshotId, splits,
definitions_, entries));
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeySortedIndexScan::EvaluatedPlan
evaluated,
+ PrimaryKeySortedIndexScan::Evaluate(plan,
table_schema_, predicate,
+
definitions_, reader_factory));
+ return PrimaryKeySortedIndexResult::ToSplits(evaluated);
+ }
+
+ std::shared_ptr<Predicate> PriceEqual(int64_t value) {
+ return PredicateBuilder::Equal(/*field_index=*/1, "price",
FieldType::BIGINT,
+ Literal(value));
+ }
+
+ std::shared_ptr<MemoryPool> pool_;
+ std::shared_ptr<UniqueTestDirectory> test_dir_;
+ std::shared_ptr<FileSystem> fs_;
+ std::string base_path_;
+ std::shared_ptr<TableSchema> table_schema_;
+ std::vector<PrimaryKeyIndexDefinition> definitions_;
+};
+
+TEST_F(PrimaryKeySortedIndexScanTest, EqualNarrowsToSingleFileRange) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ // Value 10 sits at group ordinal 5, i.e. row 5 of a.parquet.
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
PayloadReaderFactory()));
+ ASSERT_EQ(splits.size(), 1);
+ auto indexed_split =
std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]);
+ ASSERT_TRUE(indexed_split != nullptr);
+ auto inner_split =
std::dynamic_pointer_cast<DataSplitImpl>(indexed_split->GetDataSplit());
+ ASSERT_TRUE(inner_split != nullptr);
+ ASSERT_EQ(inner_split->DataFiles().size(), 1);
+ ASSERT_EQ(inner_split->DataFiles()[0]->file_name, "a.parquet");
+ ASSERT_EQ(indexed_split->RowRanges().size(), 1);
+ ASSERT_EQ(indexed_split->RowRanges()[0].from, 5);
+ ASSERT_EQ(indexed_split->RowRanges()[0].to, 5);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, BuildRejectsDuplicateOrdinals) {
+ std::vector<int64_t> ordinals;
+ ordinals.reserve(kTotalRows);
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ ordinals.push_back(i);
+ }
+ ordinals[1] = 0;
+ ASSERT_NOK_WITH_MSG(BuildPayload(std::move(ordinals)).status(),
+ "Row id 0 appears more than once");
+}
Review Comment:
`ASSERT_NOK_WITH_MSG(Func(), error_message);` is OK here.
##########
src/paimon/core/table/source/primary_key_index_batch_scan.cpp:
##########
@@ -0,0 +1,298 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_index_batch_scan.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+#include <set>
+#include <string>
+#include <thread>
+
+#include "fmt/format.h"
+#include "paimon/core/index/index_file_handler.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/table/source/plan_impl.h"
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+#include "paimon/core/table/source/snapshot/snapshot_reader.h"
+#include "paimon/core/utils/index_file_path_factories.h"
+#include "paimon/core/utils/snapshot_manager.h"
+#include "paimon/executor.h"
+#include "paimon/predicate/compound_predicate.h"
+#include "paimon/predicate/leaf_predicate.h"
+#include "paimon/predicate/predicate_builder.h"
+
+namespace paimon {
+namespace {
+Result<std::shared_ptr<Executor>> CreateGlobalIndexExecutor(const CoreOptions&
core_options) {
+ uint32_t thread_num = std::thread::hardware_concurrency();
+ std::optional<int32_t> configured_thread_num =
core_options.GetGlobalIndexThreadNum();
+ if (configured_thread_num) {
+ if (configured_thread_num.value() <= 0) {
+ return Status::Invalid(fmt::format("invalid global index thread
number {}",
+ configured_thread_num.value()));
+ }
+ thread_num = static_cast<uint32_t>(configured_thread_num.value());
+ } else if (thread_num == 0) {
+ thread_num = 1;
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Executor> executor,
CreateDefaultExecutor(thread_num));
+ return executor;
+}
+
+/// Restricts a predicate to leaves over the indexed fields: an AND keeps its
convertible
+/// children, an OR is only kept when every child is convertible, and
everything else is
+/// dropped. A null return means no part of the predicate can use the index.
+Result<std::shared_ptr<Predicate>> ProjectToIndexedFields(
+ const std::shared_ptr<Predicate>& predicate, const std::set<std::string>&
indexed_fields) {
+ if (predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (auto leaf_predicate =
std::dynamic_pointer_cast<LeafPredicate>(predicate)) {
+ if (indexed_fields.count(leaf_predicate->FieldName()) > 0) {
+ return predicate;
+ }
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ auto compound_predicate =
std::dynamic_pointer_cast<CompoundPredicate>(predicate);
+ if (compound_predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ bool is_and = compound_predicate->GetFunction().GetType() ==
Function::Type::AND;
+ bool is_or = compound_predicate->GetFunction().GetType() ==
Function::Type::OR;
+ if (!is_and && !is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ std::vector<std::shared_ptr<Predicate>> converted_children;
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> converted_child,
+ ProjectToIndexedFields(child, indexed_fields));
+ if (converted_child != nullptr) {
+ converted_children.push_back(std::move(converted_child));
+ } else if (is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ }
+ if (converted_children.empty()) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (converted_children.size() == 1) {
+ return converted_children[0];
+ }
+ if (is_and) {
+ return PredicateBuilder::And(converted_children);
+ }
+ return PredicateBuilder::Or(converted_children);
+}
+
+void FlattenChildren(const std::shared_ptr<CompoundPredicate>&
compound_predicate,
+ std::vector<std::shared_ptr<Predicate>>* flattened) {
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ auto compound_child =
std::dynamic_pointer_cast<CompoundPredicate>(child);
+ if (compound_child != nullptr &&
compound_child->GetFunction().GetType() ==
+
compound_predicate->GetFunction().GetType()) {
+ FlattenChildren(compound_child, flattened);
+ } else {
+ flattened->push_back(child);
+ }
+ }
+}
+
+/// A predicate is null-rejecting when it cannot match a row whose tested
field is null.
+/// Under SQL three-valued logic every comparison and match predicate rejects
null; only
+/// IS NULL accepts it, and IS NOT NULL is the predicate being pruned.
+bool IsNullRejecting(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ if (leaf_predicate == nullptr) {
+ return false;
+ }
+ switch (leaf_predicate->GetFunction().GetType()) {
+ case Function::Type::EQUAL:
+ case Function::Type::NOT_EQUAL:
+ case Function::Type::GREATER_THAN:
+ case Function::Type::GREATER_OR_EQUAL:
+ case Function::Type::LESS_THAN:
+ case Function::Type::LESS_OR_EQUAL:
+ case Function::Type::IN:
+ case Function::Type::NOT_IN:
+ case Function::Type::STARTS_WITH:
+ case Function::Type::ENDS_WITH:
+ case Function::Type::CONTAINS:
+ case Function::Type::LIKE:
+ return true;
+ default:
+ return false;
+ }
+}
+
+bool IsIsNotNull(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ return leaf_predicate != nullptr &&
+ leaf_predicate->GetFunction().GetType() ==
Function::Type::IS_NOT_NULL;
+}
+
+/// Flattens nested same-function compounds and, inside an AND, removes `f IS
NOT NULL`
+/// leaves made redundant by a null-rejecting sibling on the same field.
Pruning must not
+/// consider `f IS NULL` as constraining: dropping IS NOT NULL from
+/// "f IS NULL AND f IS NOT NULL" would turn the empty result into the set of
null rows.
+Result<std::shared_ptr<Predicate>> NormalizePredicate(const
std::shared_ptr<Predicate>& predicate) {
+ auto compound_predicate =
std::dynamic_pointer_cast<CompoundPredicate>(predicate);
+ if (compound_predicate == nullptr) {
+ return predicate;
+ }
+ std::vector<std::shared_ptr<Predicate>> children;
+ FlattenChildren(compound_predicate, &children);
+
+ bool is_and = compound_predicate->GetFunction().GetType() ==
Function::Type::AND;
+ if (is_and) {
+ std::set<std::string> constrained_fields;
+ for (const std::shared_ptr<Predicate>& child : children) {
+ if (IsNullRejecting(child)) {
+ constrained_fields.insert(
+
std::dynamic_pointer_cast<LeafPredicate>(child)->FieldName());
+ }
+ }
+ if (!constrained_fields.empty()) {
+ std::vector<std::shared_ptr<Predicate>> pruned;
+ pruned.reserve(children.size());
+ for (const std::shared_ptr<Predicate>& child : children) {
+ if (IsIsNotNull(child) &&
+ constrained_fields.count(
+
std::dynamic_pointer_cast<LeafPredicate>(child)->FieldName()) > 0) {
+ continue;
+ }
+ pruned.push_back(child);
+ }
+ children = std::move(pruned);
+ }
+ }
+
+ std::vector<std::shared_ptr<Predicate>> normalized_children;
+ normalized_children.reserve(children.size());
+ for (const std::shared_ptr<Predicate>& child : children) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> normalized_child,
+ NormalizePredicate(child));
+ normalized_children.push_back(std::move(normalized_child));
+ }
+ if (normalized_children.size() == 1) {
+ return normalized_children[0];
+ }
+ if (is_and) {
+ return PredicateBuilder::And(normalized_children);
+ }
+ return PredicateBuilder::Or(normalized_children);
+}
+} // namespace
Review Comment:
In Java, this `NormalizePredicate` logic seems to live in
`GlobalIndexEvaluator`, which also applies to append table queries. I’d suggest
following the same approach as Java.
##########
src/paimon/core/table/source/primary_key_index_batch_scan.cpp:
##########
@@ -0,0 +1,298 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_index_batch_scan.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+#include <set>
+#include <string>
+#include <thread>
+
+#include "fmt/format.h"
+#include "paimon/core/index/index_file_handler.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/table/source/plan_impl.h"
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+#include "paimon/core/table/source/snapshot/snapshot_reader.h"
+#include "paimon/core/utils/index_file_path_factories.h"
+#include "paimon/core/utils/snapshot_manager.h"
+#include "paimon/executor.h"
+#include "paimon/predicate/compound_predicate.h"
+#include "paimon/predicate/leaf_predicate.h"
+#include "paimon/predicate/predicate_builder.h"
+
+namespace paimon {
+namespace {
+Result<std::shared_ptr<Executor>> CreateGlobalIndexExecutor(const CoreOptions&
core_options) {
+ uint32_t thread_num = std::thread::hardware_concurrency();
+ std::optional<int32_t> configured_thread_num =
core_options.GetGlobalIndexThreadNum();
+ if (configured_thread_num) {
+ if (configured_thread_num.value() <= 0) {
+ return Status::Invalid(fmt::format("invalid global index thread
number {}",
+ configured_thread_num.value()));
+ }
+ thread_num = static_cast<uint32_t>(configured_thread_num.value());
+ } else if (thread_num == 0) {
+ thread_num = 1;
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Executor> executor,
CreateDefaultExecutor(thread_num));
+ return executor;
+}
+
+/// Restricts a predicate to leaves over the indexed fields: an AND keeps its
convertible
+/// children, an OR is only kept when every child is convertible, and
everything else is
+/// dropped. A null return means no part of the predicate can use the index.
+Result<std::shared_ptr<Predicate>> ProjectToIndexedFields(
+ const std::shared_ptr<Predicate>& predicate, const std::set<std::string>&
indexed_fields) {
+ if (predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (auto leaf_predicate =
std::dynamic_pointer_cast<LeafPredicate>(predicate)) {
+ if (indexed_fields.count(leaf_predicate->FieldName()) > 0) {
+ return predicate;
+ }
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ auto compound_predicate =
std::dynamic_pointer_cast<CompoundPredicate>(predicate);
+ if (compound_predicate == nullptr) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ bool is_and = compound_predicate->GetFunction().GetType() ==
Function::Type::AND;
+ bool is_or = compound_predicate->GetFunction().GetType() ==
Function::Type::OR;
+ if (!is_and && !is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ std::vector<std::shared_ptr<Predicate>> converted_children;
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> converted_child,
+ ProjectToIndexedFields(child, indexed_fields));
+ if (converted_child != nullptr) {
+ converted_children.push_back(std::move(converted_child));
+ } else if (is_or) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ }
+ if (converted_children.empty()) {
+ return std::shared_ptr<Predicate>(nullptr);
+ }
+ if (converted_children.size() == 1) {
+ return converted_children[0];
+ }
+ if (is_and) {
+ return PredicateBuilder::And(converted_children);
+ }
+ return PredicateBuilder::Or(converted_children);
+}
+
+void FlattenChildren(const std::shared_ptr<CompoundPredicate>&
compound_predicate,
+ std::vector<std::shared_ptr<Predicate>>* flattened) {
+ for (const std::shared_ptr<Predicate>& child :
compound_predicate->Children()) {
+ auto compound_child =
std::dynamic_pointer_cast<CompoundPredicate>(child);
+ if (compound_child != nullptr &&
compound_child->GetFunction().GetType() ==
+
compound_predicate->GetFunction().GetType()) {
+ FlattenChildren(compound_child, flattened);
+ } else {
+ flattened->push_back(child);
+ }
+ }
+}
+
+/// A predicate is null-rejecting when it cannot match a row whose tested
field is null.
+/// Under SQL three-valued logic every comparison and match predicate rejects
null; only
+/// IS NULL accepts it, and IS NOT NULL is the predicate being pruned.
+bool IsNullRejecting(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ if (leaf_predicate == nullptr) {
+ return false;
+ }
+ switch (leaf_predicate->GetFunction().GetType()) {
+ case Function::Type::EQUAL:
+ case Function::Type::NOT_EQUAL:
+ case Function::Type::GREATER_THAN:
+ case Function::Type::GREATER_OR_EQUAL:
+ case Function::Type::LESS_THAN:
+ case Function::Type::LESS_OR_EQUAL:
+ case Function::Type::IN:
+ case Function::Type::NOT_IN:
+ case Function::Type::STARTS_WITH:
+ case Function::Type::ENDS_WITH:
+ case Function::Type::CONTAINS:
+ case Function::Type::LIKE:
+ return true;
+ default:
+ return false;
+ }
+}
+
+bool IsIsNotNull(const std::shared_ptr<Predicate>& predicate) {
+ auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicate>(predicate);
+ return leaf_predicate != nullptr &&
+ leaf_predicate->GetFunction().GetType() ==
Function::Type::IS_NOT_NULL;
+}
+
+/// Flattens nested same-function compounds and, inside an AND, removes `f IS
NOT NULL`
+/// leaves made redundant by a null-rejecting sibling on the same field.
Pruning must not
+/// consider `f IS NULL` as constraining: dropping IS NOT NULL from
+/// "f IS NULL AND f IS NOT NULL" would turn the empty result into the set of
null rows.
+Result<std::shared_ptr<Predicate>> NormalizePredicate(const
std::shared_ptr<Predicate>& predicate) {
Review Comment:
This class contains a lot of complex logic. Please add unit tests to cover
various predicate scenarios.
##########
src/paimon/core/operation/raw_file_split_read.h:
##########
@@ -65,16 +64,26 @@ class RawFileSplitRead : public AbstractSplitRead {
const std::shared_ptr<MemoryPool>& memory_pool,
const std::shared_ptr<Executor>& executor);
+ /// Also accepts an `IndexedSplit` over a single-file data split, in which
case its row
+ /// ranges narrow the read to the given file-local physical positions.
Result<std::unique_ptr<BatchReader>> CreateReader(const
std::shared_ptr<Split>& split) override;
Result<std::unique_ptr<BatchReader>> CreateReader(
const BinaryRow& partition, int32_t bucket,
const std::vector<std::shared_ptr<DataFileMeta>>& files,
const std::vector<std::optional<DeletionFile>>& deletion_files);
+ /// Reads with an optional selection of file-local row positions. The
ranges apply to
+ /// every file of the split, so callers pass them only for single-file
splits.
Result<std::unique_ptr<BatchReader>> CreateReader(
const BinaryRow& partition, int32_t bucket,
const std::vector<std::shared_ptr<DataFileMeta>>& files,
- DeletionVector::Factory dv_factory);
+ const std::vector<std::optional<DeletionFile>>& deletion_files,
+ const std::optional<std::vector<Range>>& local_row_ranges);
+
+ Result<std::unique_ptr<BatchReader>> CreateReader(
+ const BinaryRow& partition, int32_t bucket,
+ const std::vector<std::shared_ptr<DataFileMeta>>& files,
DeletionVector::Factory dv_factory,
+ const std::optional<std::vector<Range>>& local_row_ranges =
std::nullopt);
Review Comment:
Also, I’m not sure it’s necessary to add a new `CreateReader` function.
Could we just extend the existing one with a `local_row_ranges` parameter
instead?
##########
src/paimon/core/table/source/primary_key_sorted_index_scan.cpp:
##########
@@ -0,0 +1,568 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+
+#include <cassert>
+#include <set>
+#include <unordered_map>
+#include <utility>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
+#include "fmt/format.h"
+#include "paimon/common/types/data_field.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/global_index/global_index_evaluator_impl.h"
+#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h"
+#include "paimon/core/manifest/file_kind.h"
+#include "paimon/global_index/bitmap_global_index_result.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+#include "paimon/global_index/io/global_index_file_reader.h"
+#include "paimon/predicate/predicate_utils.h"
+
+namespace paimon {
+namespace {
+using BucketKey = std::pair<BinaryRow, int32_t>;
+
+enum class QueryOperation {
+ IS_NOT_NULL,
+ IS_NULL,
+ EQUAL,
+ NOT_EQUAL,
+ LESS_THAN,
+ LESS_OR_EQUAL,
+ GREATER_THAN,
+ GREATER_OR_EQUAL,
+ IN,
+ NOT_IN,
+ STARTS_WITH,
+ ENDS_WITH,
+ CONTAINS,
+ LIKE,
+};
+
Review Comment:
Could you help clarify the distinction between `QueryOperation` and
`Function`? They seem somewhat overlapping to me, so I’m wondering why both are
needed separately.
##########
src/paimon/core/table/source/primary_key_sorted_index_scan.cpp:
##########
@@ -0,0 +1,568 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+
+#include <cassert>
+#include <set>
+#include <unordered_map>
+#include <utility>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
+#include "fmt/format.h"
+#include "paimon/common/types/data_field.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/global_index/global_index_evaluator_impl.h"
+#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h"
+#include "paimon/core/manifest/file_kind.h"
+#include "paimon/global_index/bitmap_global_index_result.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+#include "paimon/global_index/io/global_index_file_reader.h"
+#include "paimon/predicate/predicate_utils.h"
+
+namespace paimon {
+namespace {
+using BucketKey = std::pair<BinaryRow, int32_t>;
+
+enum class QueryOperation {
+ IS_NOT_NULL,
+ IS_NULL,
+ EQUAL,
+ NOT_EQUAL,
+ LESS_THAN,
+ LESS_OR_EQUAL,
+ GREATER_THAN,
+ GREATER_OR_EQUAL,
+ IN,
+ NOT_IN,
+ STARTS_WITH,
+ ENDS_WITH,
+ CONTAINS,
+ LIKE,
+};
+
+struct QueryKey {
+ QueryOperation operation;
+ std::vector<Literal> literals;
+
+ bool operator==(const QueryKey& other) const {
+ return operation == other.operation && literals == other.literals;
+ }
+};
+
+/// Shares one group payload reader and its group-scope query results across
all source
+/// files of the group; localizes group ordinals to file-local physical
positions using the
+/// ordered source row-count prefix.
+class SharedGroupReader {
+ public:
+ using UnderlyingReaderFactory =
std::function<Result<std::shared_ptr<GlobalIndexReader>>()>;
+
+ SharedGroupReader(const std::shared_ptr<PkSortedIndexGroup>& group,
+ UnderlyingReaderFactory reader_factory)
+ : group_(group), reader_factory_(std::move(reader_factory)) {
+ const std::vector<PrimaryKeyIndexSourceFile>& source_files =
group->SourceFiles();
+ source_offsets_.reserve(source_files.size() + 1);
+ source_offsets_.push_back(0);
+ for (const PrimaryKeyIndexSourceFile& source_file : source_files) {
+ source_offsets_.push_back(source_offsets_.back() +
source_file.row_count);
+ }
+ }
+
+ const std::shared_ptr<PkSortedIndexGroup>& Group() const {
+ return group_;
+ }
+
+ /// Runs one group-scope query with caching; equal queries evaluate
exactly once.
+ Result<std::shared_ptr<GlobalIndexResult>> Query(
+ const QueryKey& key,
+ const
std::function<Result<std::shared_ptr<GlobalIndexResult>>(GlobalIndexReader*)>&
+ query) {
+ for (const auto& cached : query_cache_) {
+ if (cached.first == key) {
+ if (!cached.second.status.ok()) {
+ return cached.second.status;
+ }
+ return cached.second.result;
+ }
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> result = RunQuery(query);
+ CachedQuery cached_query;
+ if (result.ok()) {
+ cached_query.result = result.value();
+ } else {
+ cached_query.status = result.status();
+ }
+ query_cache_.emplace_back(key, cached_query);
+ return result;
+ }
+
+ /// Restricts one group-scope result to the local row positions of
`source_index`.
+ /// Any out-of-range group ordinal fails the localization so that every
covered file
+ /// of this group falls back to a normal scan; a poison marker would not
survive the
+ /// AND/OR combination of results from other indexes.
+ Result<std::shared_ptr<GlobalIndexResult>> Localize(
+ const std::shared_ptr<GlobalIndexResult>& result, size_t source_index)
{
+ if (result == nullptr) {
+ return std::shared_ptr<GlobalIndexResult>(nullptr);
+ }
+ assert(source_index + 1 < source_offsets_.size());
+ auto localized = localized_cache_.find(result.get());
+ if (localized == localized_cache_.end()) {
+
PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<GlobalIndexResult>>
partitions,
+ PartitionBySource(result));
+ localized = localized_cache_.emplace(result.get(),
std::move(partitions)).first;
+ }
+ return localized->second[source_index];
+ }
+
+ private:
+ struct CachedQuery {
+ Status status;
+ std::shared_ptr<GlobalIndexResult> result;
+ };
+
+ Result<std::shared_ptr<GlobalIndexResult>> RunQuery(
+ const
std::function<Result<std::shared_ptr<GlobalIndexResult>>(GlobalIndexReader*)>&
+ query) {
+ if (!reader_status_.ok()) {
+ return reader_status_;
+ }
+ if (reader_ == nullptr) {
+ Result<std::shared_ptr<GlobalIndexReader>> reader_result =
reader_factory_();
+ if (!reader_result.ok()) {
+ reader_status_ = reader_result.status();
+ return reader_status_;
+ }
+ reader_ = reader_result.value();
+ if (reader_ == nullptr) {
+ // The index type has no usable reader; keep normal scan
semantics.
+ return std::shared_ptr<GlobalIndexResult>(nullptr);
+ }
+ }
+ return query(reader_.get());
+ }
+
+ Result<std::vector<std::shared_ptr<GlobalIndexResult>>> PartitionBySource(
+ const std::shared_ptr<GlobalIndexResult>& result) {
+ size_t source_count = source_offsets_.size() - 1;
+ std::vector<RoaringBitmap64> partitions(source_count);
+ int64_t total_row_count = source_offsets_.back();
+ size_t source_index = 0;
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<GlobalIndexResult::Iterator>
iterator,
+ result->CreateIterator());
+ while (iterator->HasNext()) {
+ int64_t position = iterator->Next();
+ if (position < 0 || position >= total_row_count) {
+ return Status::Invalid(fmt::format(
+ "Sorted index returned group ordinal {} outside the source
row range "
+ "[0, {}).",
+ position, total_row_count));
+ }
+ while (position >= source_offsets_[source_index + 1]) {
+ source_index++;
+ }
+ partitions[source_index].Add(position -
source_offsets_[source_index]);
+ }
+ std::vector<std::shared_ptr<GlobalIndexResult>> localized;
+ localized.reserve(source_count);
+ for (RoaringBitmap64& partition : partitions) {
+ auto bitmap =
std::make_shared<RoaringBitmap64>(std::move(partition));
+ localized.push_back(std::make_shared<BitmapGlobalIndexResult>(
+ [bitmap]() -> Result<RoaringBitmap64> { return *bitmap; }));
+ }
+ return localized;
+ }
+
+ std::shared_ptr<PkSortedIndexGroup> group_;
+ UnderlyingReaderFactory reader_factory_;
+ std::vector<int64_t> source_offsets_;
+ std::vector<std::pair<QueryKey, CachedQuery>> query_cache_;
+ std::unordered_map<const GlobalIndexResult*,
std::vector<std::shared_ptr<GlobalIndexResult>>>
+ localized_cache_;
+ std::shared_ptr<GlobalIndexReader> reader_;
+ Status reader_status_;
+};
+
+/// Restricts merged source-group ordinals to one source file's local row
positions.
+class FileLocalGroupReader : public GlobalIndexReader {
+ public:
+ FileLocalGroupReader(std::shared_ptr<SharedGroupReader> shared_reader,
size_t source_index)
+ : shared_reader_(std::move(shared_reader)),
source_index_(source_index) {}
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNotNull() override {
+ return Query({QueryOperation::IS_NOT_NULL, {}},
+ [](GlobalIndexReader* reader) { return
reader->VisitIsNotNull(); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNull() override {
+ return Query({QueryOperation::IS_NULL, {}},
+ [](GlobalIndexReader* reader) { return
reader->VisitIsNull(); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEqual(const Literal&
literal) override {
+ return Query({QueryOperation::EQUAL, {literal}},
+ [&literal](GlobalIndexReader* reader) { return
reader->VisitEqual(literal); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotEqual(const Literal&
literal) override {
+ return Query({QueryOperation::NOT_EQUAL, {literal}},
[&literal](GlobalIndexReader* reader) {
+ return reader->VisitNotEqual(literal);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessThan(const Literal&
literal) override {
+ return Query({QueryOperation::LESS_THAN, {literal}},
[&literal](GlobalIndexReader* reader) {
+ return reader->VisitLessThan(literal);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessOrEqual(const Literal&
literal) override {
+ return Query(
+ {QueryOperation::LESS_OR_EQUAL, {literal}},
+ [&literal](GlobalIndexReader* reader) { return
reader->VisitLessOrEqual(literal); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterThan(const Literal&
literal) override {
+ return Query(
+ {QueryOperation::GREATER_THAN, {literal}},
+ [&literal](GlobalIndexReader* reader) { return
reader->VisitGreaterThan(literal); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterOrEqual(
+ const Literal& literal) override {
+ return Query(
+ {QueryOperation::GREATER_OR_EQUAL, {literal}},
+ [&literal](GlobalIndexReader* reader) { return
reader->VisitGreaterOrEqual(literal); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIn(
+ const std::vector<Literal>& literals) override {
+ return Query({QueryOperation::IN, literals},
+ [&literals](GlobalIndexReader* reader) { return
reader->VisitIn(literals); });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotIn(
+ const std::vector<Literal>& literals) override {
+ return Query({QueryOperation::NOT_IN, literals},
[&literals](GlobalIndexReader* reader) {
+ return reader->VisitNotIn(literals);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitStartsWith(const Literal&
prefix) override {
+ return Query({QueryOperation::STARTS_WITH, {prefix}},
[&prefix](GlobalIndexReader* reader) {
+ return reader->VisitStartsWith(prefix);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEndsWith(const Literal&
suffix) override {
+ return Query({QueryOperation::ENDS_WITH, {suffix}},
[&suffix](GlobalIndexReader* reader) {
+ return reader->VisitEndsWith(suffix);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitContains(const Literal&
literal) override {
+ return Query({QueryOperation::CONTAINS, {literal}},
[&literal](GlobalIndexReader* reader) {
+ return reader->VisitContains(literal);
+ });
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLike(const Literal&
literal) override {
+ return Query({QueryOperation::LIKE, {literal}},
+ [&literal](GlobalIndexReader* reader) { return
reader->VisitLike(literal); });
+ }
+
+ Result<std::shared_ptr<ScoredGlobalIndexResult>> VisitVectorSearch(
+ const std::shared_ptr<VectorSearch>& vector_search) override {
+ return Status::Invalid("Primary-key sorted index does not support
vector search.");
+ }
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitFullTextSearch(
+ const std::shared_ptr<FullTextSearch>& full_text_search) override {
+ return Status::Invalid("Primary-key sorted index does not support full
text search.");
+ }
+
+ bool IsThreadSafe() const override {
+ return false;
+ }
+
+ std::string GetIndexType() const override {
+ return shared_reader_->Group()->Payload()->IndexType();
+ }
+
+ private:
+ Result<std::shared_ptr<GlobalIndexResult>> Query(
+ QueryKey key,
+ const
std::function<Result<std::shared_ptr<GlobalIndexResult>>(GlobalIndexReader*)>&
+ query) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<GlobalIndexResult> group_result,
+ shared_reader_->Query(key, query));
+ return shared_reader_->Localize(group_result, source_index_);
+ }
+
+ std::shared_ptr<SharedGroupReader> shared_reader_;
+ size_t source_index_;
+};
+
+Result<size_t> FindSourceIndex(const PkSortedIndexGroup& group, const
DataFileMeta& data_file) {
+ const std::vector<PrimaryKeyIndexSourceFile>& source_files =
group.SourceFiles();
+ for (size_t i = 0; i < source_files.size(); i++) {
+ if (source_files[i].file_name == data_file.file_name &&
+ source_files[i].row_count == data_file.row_count) {
+ return i;
+ }
+ }
+ return Status::Invalid(fmt::format(
+ "Data file {} is not covered by its sorted-index source group.",
data_file.file_name));
+}
+} // namespace
+
+Result<PrimaryKeySortedIndexScan::Plan> PrimaryKeySortedIndexScan::CreatePlan(
+ int64_t snapshot_id, const std::vector<std::shared_ptr<DataSplitImpl>>&
data_splits,
+ const std::vector<PrimaryKeyIndexDefinition>& definitions,
+ const std::vector<IndexManifestEntry>& index_entries) {
+ std::unordered_map<BucketKey, std::vector<std::shared_ptr<IndexFileMeta>>>
payloads_by_bucket;
+ for (const IndexManifestEntry& entry : index_entries) {
+ const std::shared_ptr<IndexFileMeta>& payload = entry.index_file;
+ if (payload == nullptr || !(entry.kind == FileKind::Add())) {
+ continue;
+ }
+ const std::optional<GlobalIndexMeta>& meta =
payload->GetGlobalIndexMeta();
+ if (meta == std::nullopt || meta.value().source_meta == nullptr) {
+ continue;
+ }
+ payloads_by_bucket[BucketKey(entry.partition,
entry.bucket)].push_back(payload);
+ }
+
+ std::vector<PrimaryKeyIndexDefinition> scalar_definitions;
+ for (const PrimaryKeyIndexDefinition& definition : definitions) {
+ if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE
||
+ definition.GetFamily() ==
PrimaryKeyIndexDefinition::Family::BITMAP) {
+ scalar_definitions.push_back(definition);
+ }
+ }
+
Review Comment:
Could `ScalarDefinitions` be handling something similar here?
##########
src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp:
##########
@@ -0,0 +1,548 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "paimon/core/table/source/primary_key_sorted_index_scan.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "fmt/format.h"
+#include "gtest/gtest.h"
+#include "paimon/core/global_index/indexed_split_impl.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/index/pksorted/pk_sorted_index_file.h"
+#include "paimon/core/manifest/file_kind.h"
+#include "paimon/core/manifest/file_source.h"
+#include "paimon/core/table/source/primary_key_sorted_index_result.h"
+#include "paimon/global_index/bitmap_global_index_result.h"
+#include "paimon/global_index/global_index_io_meta.h"
+#include "paimon/global_index/global_indexer.h"
+#include "paimon/global_index/global_indexer_factory.h"
+#include "paimon/global_index/io/global_index_file_reader.h"
+#include "paimon/global_index/io/global_index_file_writer.h"
+#include "paimon/predicate/predicate_builder.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+namespace {
+constexpr int32_t kPriceFieldId = 1;
+constexpr int64_t kSnapshotId = 7;
+constexpr int64_t kFileARows = 100;
+constexpr int64_t kFileBRows = 200;
+constexpr int64_t kTotalRows = kFileARows + kFileBRows;
+
+class TestGlobalIndexFileWriter : public GlobalIndexFileWriter {
+ public:
+ TestGlobalIndexFileWriter(const std::shared_ptr<FileSystem>& fs, const
std::string& base_path)
+ : fs_(fs), base_path_(base_path) {}
+
+ Result<std::string> NewFileName(const std::string& prefix) const override {
+ return fmt::format("{}-index-{}", prefix, file_counter_++);
+ }
+
+ Result<std::unique_ptr<OutputStream>> NewOutputStream(
+ const std::string& file_name) const override {
+ return fs_->Create(base_path_ + "/" + file_name, true);
+ }
+
+ Result<int64_t> GetFileSize(const std::string& file_name) const override {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FileStatus> file_status,
+ fs_->GetFileStatus(base_path_ + "/" +
file_name));
+ return file_status->GetLen();
+ }
+
+ std::string ToPath(const std::string& file_name) const override {
+ return base_path_ + "/" + file_name;
+ }
+
+ private:
+ std::shared_ptr<FileSystem> fs_;
+ std::string base_path_;
+ mutable int64_t file_counter_ = 0;
+};
+
+class TestGlobalIndexFileReader : public GlobalIndexFileReader {
+ public:
+ explicit TestGlobalIndexFileReader(const std::shared_ptr<FileSystem>& fs)
: fs_(fs) {}
+
+ Result<std::unique_ptr<InputStream>> GetInputStream(
+ const std::string& file_path) const override {
+ return fs_->Open(file_path);
+ }
+
+ private:
+ std::shared_ptr<FileSystem> fs_;
+};
+
+/// A reader stub whose equality result is fully controlled by the test, used
to exercise
+/// the untrusted-position fallbacks.
+class StubGlobalIndexReader : public GlobalIndexReader {
+ public:
+ explicit StubGlobalIndexReader(RoaringBitmap64 equal_result)
+ : equal_result_(std::move(equal_result)) {}
+
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNotNull() override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIsNull() override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEqual(const Literal&
literal) override {
+ RoaringBitmap64 copy = equal_result_;
+ return std::make_shared<BitmapGlobalIndexResult>(
+ [bitmap = std::move(copy)]() -> Result<RoaringBitmap64> { return
bitmap; });
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotEqual(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessThan(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLessOrEqual(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterThan(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitGreaterOrEqual(
+ const Literal& literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitIn(
+ const std::vector<Literal>& literals) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitNotIn(
+ const std::vector<Literal>& literals) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitStartsWith(const Literal&
prefix) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitEndsWith(const Literal&
suffix) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitContains(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitLike(const Literal&
literal) override {
+ return NotEvaluable();
+ }
+ Result<std::shared_ptr<ScoredGlobalIndexResult>> VisitVectorSearch(
+ const std::shared_ptr<VectorSearch>& vector_search) override {
+ return Status::Invalid("not supported");
+ }
+ Result<std::shared_ptr<GlobalIndexResult>> VisitFullTextSearch(
+ const std::shared_ptr<FullTextSearch>& full_text_search) override {
+ return Status::Invalid("not supported");
+ }
+ bool IsThreadSafe() const override {
+ return false;
+ }
+ std::string GetIndexType() const override {
+ return "btree";
+ }
+
+ private:
+ static Result<std::shared_ptr<GlobalIndexResult>> NotEvaluable() {
+ return std::shared_ptr<GlobalIndexResult>(nullptr);
+ }
+
+ RoaringBitmap64 equal_result_;
+};
+} // namespace
+
+class PrimaryKeySortedIndexScanTest : public ::testing::Test {
+ protected:
+ void SetUp() override {
+ pool_ = GetDefaultPool();
+ test_dir_ = UniqueTestDirectory::Create("local");
+ fs_ = test_dir_->GetFileSystem();
+ base_path_ = test_dir_->Str();
+
+ std::vector<DataField> fields = {
+ DataField(0, arrow::field("id", arrow::int64())),
+ DataField(kPriceFieldId, arrow::field("price", arrow::int64())),
+ DataField(2, arrow::field("status", arrow::utf8())),
+ };
+ std::map<std::string, std::string> options =
{{"pk-btree.index.columns", "price"}};
+ table_schema_ = std::make_shared<TableSchema>(
+ /*version=*/3, /*id=*/0, fields, /*highest_field_id=*/2,
+ /*partition_keys=*/std::vector<std::string>(),
+ /*primary_keys=*/std::vector<std::string>{"id"}, options,
+ /*comment=*/std::nullopt, /*time_millis=*/0);
+ ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions,
+
PrimaryKeyIndexDefinitions::Create(*table_schema_));
+ definitions_ = definitions.ScalarDefinitions();
+ ASSERT_EQ(definitions_.size(), 1);
+ }
+
+ std::shared_ptr<DataFileMeta> MakeDataFile(const std::string& name,
int64_t row_count,
+ int32_t level, const
FileSource& file_source) {
+ return std::make_shared<DataFileMeta>(
+ name, /*file_size=*/1024, row_count,
+ /*min_key=*/BinaryRow::EmptyRow(),
/*max_key=*/BinaryRow::EmptyRow(),
+ /*key_stats=*/SimpleStats::EmptyStats(),
/*value_stats=*/SimpleStats::EmptyStats(),
+ /*min_sequence_number=*/0, /*max_sequence_number=*/row_count,
/*schema_id=*/0, level,
+ /*extra_files=*/std::vector<std::optional<std::string>>(),
+ /*creation_time=*/Timestamp(1721643142456LL, 0),
+ /*delete_row_count=*/0, /*embedded_index=*/nullptr, file_source,
+ /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt,
+ /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt);
+ }
+
+ Result<std::shared_ptr<IndexFileMeta>> BuildPayload(std::vector<int64_t>
ordinals) {
+ std::vector<PrimaryKeyIndexSourceFile> source_files = {{"a.parquet",
kFileARows},
+ {"b.parquet",
kFileBRows}};
+ arrow::Int64Builder values_builder;
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Append(2 * i));
+ }
+ std::shared_ptr<arrow::Array> sorted_values;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Finish(&sorted_values));
+ PAIMON_ASSIGN_OR_RAISE(DataField field,
table_schema_->GetField(kPriceFieldId));
+ auto file_writer = std::make_shared<TestGlobalIndexFileWriter>(fs_,
base_path_);
+ return PkSortedIndexFile::Build(field, "btree",
definitions_[0].Options(),
+ /*data_level=*/5, source_files,
sorted_values,
+ std::move(ordinals), file_writer,
+ /*is_external_path=*/false, pool_);
+ }
+
+ /// Builds the standard payload of this fixture: sources a.parquet(100) +
b.parquet(200)
+ /// on level 5, indexed value at group ordinal `i` is `2 * i`.
+ Result<std::shared_ptr<IndexFileMeta>> BuildPayload() {
+ std::vector<int64_t> ordinals;
+ ordinals.reserve(kTotalRows);
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ ordinals.push_back(i);
+ }
+ return BuildPayload(std::move(ordinals));
+ }
+
+ std::shared_ptr<DataSplitImpl> MakeSplit(
+ const std::vector<std::shared_ptr<DataFileMeta>>& files, bool
raw_convertible,
+ const std::vector<std::optional<DeletionFile>>& deletion_files = {}) {
+ std::vector<std::shared_ptr<DataFileMeta>> data_files = files;
+ DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0,
+ base_path_ + "/bucket-0",
std::move(data_files));
+
builder.WithSnapshot(kSnapshotId).IsStreaming(false).RawConvertible(raw_convertible);
+ if (!deletion_files.empty()) {
+ builder.WithDataDeletionFiles(deletion_files);
+ }
+ EXPECT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> split,
builder.Build());
+ return split;
+ }
+
+ std::vector<IndexManifestEntry> MakeEntries(const
std::shared_ptr<IndexFileMeta>& payload) {
+ return {IndexManifestEntry(FileKind::Add(), BinaryRow::EmptyRow(),
/*bucket=*/0, payload)};
+ }
+
+ PrimaryKeySortedIndexScan::ReaderFactory PayloadReaderFactory() {
+ std::shared_ptr<FileSystem> fs = fs_;
+ std::string base_path = base_path_;
+ std::shared_ptr<TableSchema> table_schema = table_schema_;
+ std::shared_ptr<MemoryPool> pool = pool_;
+ return [fs, base_path, table_schema, pool](
+ const PrimaryKeySortedIndexScan::FilePlan& file,
+ const PrimaryKeyIndexDefinition& definition,
+ const PkSortedIndexGroup& group) ->
Result<std::shared_ptr<GlobalIndexReader>> {
+ PAIMON_ASSIGN_OR_RAISE(
+ std::unique_ptr<GlobalIndexer> indexer,
+ GlobalIndexerFactory::Get(definition.IndexType(),
definition.Options()));
+ if (indexer == nullptr) {
+ return Status::Invalid("btree indexer is not registered");
+ }
+ const std::shared_ptr<IndexFileMeta>& payload = group.Payload();
+ std::vector<GlobalIndexIOMeta> io_metas;
+ io_metas.emplace_back(base_path + "/" + payload->FileName(),
payload->FileSize(),
+
payload->GetGlobalIndexMeta().value().index_meta);
+ PAIMON_ASSIGN_OR_RAISE(DataField field,
table_schema->GetField(definition.FieldId()));
+ auto arrow_field = DataField::ConvertDataFieldToArrowField(field);
+ auto arrow_schema = arrow::schema({arrow_field});
+ ArrowSchema c_arrow_schema;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema,
&c_arrow_schema));
+ auto file_reader = std::make_shared<TestGlobalIndexFileReader>(fs);
+ return indexer->CreateReader(&c_arrow_schema, file_reader,
io_metas, pool);
+ };
+ }
+
+ Result<std::vector<std::shared_ptr<Split>>> PlanEvaluateConvert(
+ const std::vector<std::shared_ptr<DataSplitImpl>>& splits,
+ const std::vector<IndexManifestEntry>& entries, const
std::shared_ptr<Predicate>& predicate,
+ const PrimaryKeySortedIndexScan::ReaderFactory& reader_factory) {
+ PAIMON_ASSIGN_OR_RAISE(
+ PrimaryKeySortedIndexScan::Plan plan,
+ PrimaryKeySortedIndexScan::CreatePlan(kSnapshotId, splits,
definitions_, entries));
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeySortedIndexScan::EvaluatedPlan
evaluated,
+ PrimaryKeySortedIndexScan::Evaluate(plan,
table_schema_, predicate,
+
definitions_, reader_factory));
+ return PrimaryKeySortedIndexResult::ToSplits(evaluated);
+ }
+
+ std::shared_ptr<Predicate> PriceEqual(int64_t value) {
+ return PredicateBuilder::Equal(/*field_index=*/1, "price",
FieldType::BIGINT,
+ Literal(value));
+ }
+
+ std::shared_ptr<MemoryPool> pool_;
+ std::shared_ptr<UniqueTestDirectory> test_dir_;
+ std::shared_ptr<FileSystem> fs_;
+ std::string base_path_;
+ std::shared_ptr<TableSchema> table_schema_;
+ std::vector<PrimaryKeyIndexDefinition> definitions_;
+};
+
+TEST_F(PrimaryKeySortedIndexScanTest, EqualNarrowsToSingleFileRange) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ // Value 10 sits at group ordinal 5, i.e. row 5 of a.parquet.
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
PayloadReaderFactory()));
+ ASSERT_EQ(splits.size(), 1);
+ auto indexed_split =
std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]);
+ ASSERT_TRUE(indexed_split != nullptr);
+ auto inner_split =
std::dynamic_pointer_cast<DataSplitImpl>(indexed_split->GetDataSplit());
+ ASSERT_TRUE(inner_split != nullptr);
+ ASSERT_EQ(inner_split->DataFiles().size(), 1);
+ ASSERT_EQ(inner_split->DataFiles()[0]->file_name, "a.parquet");
+ ASSERT_EQ(indexed_split->RowRanges().size(), 1);
+ ASSERT_EQ(indexed_split->RowRanges()[0].from, 5);
+ ASSERT_EQ(indexed_split->RowRanges()[0].to, 5);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, BuildRejectsDuplicateOrdinals) {
+ std::vector<int64_t> ordinals;
+ ordinals.reserve(kTotalRows);
+ for (int64_t i = 0; i < kTotalRows; i++) {
+ ordinals.push_back(i);
+ }
+ ordinals[1] = 0;
+ ASSERT_NOK_WITH_MSG(BuildPayload(std::move(ordinals)).status(),
+ "Row id 0 appears more than once");
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, RangeSpansFileBoundary) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ // Values in [190, 210] sit at group ordinals 95..105: rows 95..99 of
a.parquet and
+ // rows 0..5 of b.parquet.
+ std::shared_ptr<Predicate> lower = PredicateBuilder::GreaterOrEqual(
+ /*field_index=*/1, "price", FieldType::BIGINT,
Literal(static_cast<int64_t>(190)));
+ std::shared_ptr<Predicate> upper = PredicateBuilder::LessOrEqual(
+ /*field_index=*/1, "price", FieldType::BIGINT,
Literal(static_cast<int64_t>(210)));
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> predicate,
+ PredicateBuilder::And({lower, upper}));
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), predicate,
PayloadReaderFactory()));
+ ASSERT_EQ(splits.size(), 2);
+ auto indexed_a = std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]);
+ auto indexed_b = std::dynamic_pointer_cast<IndexedSplitImpl>(splits[1]);
+ ASSERT_TRUE(indexed_a != nullptr);
+ ASSERT_TRUE(indexed_b != nullptr);
+ ASSERT_EQ(indexed_a->RowRanges().size(), 1);
+ ASSERT_EQ(indexed_a->RowRanges()[0].from, 95);
+ ASSERT_EQ(indexed_a->RowRanges()[0].to, 99);
+ ASSERT_EQ(indexed_b->RowRanges().size(), 1);
+ ASSERT_EQ(indexed_b->RowRanges()[0].from, 0);
+ ASSERT_EQ(indexed_b->RowRanges()[0].to, 5);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, EmptyResultOmitsAllFiles) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ // All indexed values are even, so 11 matches nothing.
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(11),
PayloadReaderFactory()));
+ ASSERT_TRUE(splits.empty());
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, UnindexedFieldPredicateFallsBack) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ std::shared_ptr<Predicate> predicate = PredicateBuilder::Equal(
+ /*field_index=*/2, "status", FieldType::STRING,
Literal(FieldType::STRING, "hit", 3));
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), predicate,
PayloadReaderFactory()));
+ ASSERT_EQ(1, splits.size());
+ ASSERT_EQ(split, splits[0]);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, UncoveredFileFallsBackOthersNarrow) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact()),
+ MakeDataFile("c.parquet", 50, 0, FileSource::Append())},
+ /*raw_convertible=*/true);
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
PayloadReaderFactory()));
+ // a.parquet narrows to an indexed split, b.parquet is omitted, c.parquet
has no
+ // coverage and keeps a normal single-file scan.
+ ASSERT_EQ(splits.size(), 2);
+ auto indexed_split =
std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]);
+ ASSERT_TRUE(indexed_split != nullptr);
+ auto fallback_split = std::dynamic_pointer_cast<DataSplitImpl>(splits[1]);
+ ASSERT_TRUE(fallback_split != nullptr);
+ ASSERT_EQ(fallback_split->DataFiles().size(), 1);
+ ASSERT_EQ(fallback_split->DataFiles()[0]->file_name, "c.parquet");
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, NonRawConvertibleSplitPreserved) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/false);
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
PayloadReaderFactory()));
+ ASSERT_EQ(splits.size(), 1);
+ ASSERT_EQ(splits[0].get(), split.get());
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, InvalidRowRangePayloadFallsBack) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ // Rebuild the payload metadata with a row range end beyond the source
rows: the group
+ // validation must reject it and every file keeps a normal scan.
+ const GlobalIndexMeta& meta = payload->GetGlobalIndexMeta().value();
+ auto broken_payload = std::make_shared<IndexFileMeta>(
+ payload->IndexType(), payload->FileName(), payload->FileSize(),
payload->RowCount(),
+ std::nullopt, std::nullopt,
+ GlobalIndexMeta(meta.row_range_start, meta.row_range_end + 1,
meta.index_field_id,
+ meta.extra_field_ids, meta.index_meta,
meta.source_meta));
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split},
MakeEntries(broken_payload), PriceEqual(10),
+ PayloadReaderFactory()));
+ ASSERT_EQ(1, splits.size());
+ ASSERT_EQ(split, splits[0]);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, OutOfRangePositionsFailAllCoveredFiles) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::shared_ptr<DataSplitImpl> split =
+ MakeSplit({MakeDataFile("a.parquet", kFileARows, 5,
FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5,
FileSource::Compact())},
+ /*raw_convertible=*/true);
+ RoaringBitmap64 poisoned;
+ poisoned.Add(5);
+ poisoned.Add(kTotalRows + 10);
+ PrimaryKeySortedIndexScan::ReaderFactory stub_factory =
+ [&poisoned](const PrimaryKeySortedIndexScan::FilePlan& file,
+ const PrimaryKeyIndexDefinition& definition,
+ const PkSortedIndexGroup& group) ->
Result<std::shared_ptr<GlobalIndexReader>> {
+ return std::make_shared<StubGlobalIndexReader>(poisoned);
+ };
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
stub_factory));
+ // Both covered files fall back together, preserving the planner's
original bin packing.
+ ASSERT_EQ(1, splits.size());
+ ASSERT_EQ(split, splits[0]);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, OverFragmentedResultFallsBack) {
+ // One data file, 20000 rows; every second row selected produces > 4096
ranges.
+ std::vector<PrimaryKeyIndexSourceFile> source_files = {{"big.parquet",
20000}};
+ std::shared_ptr<DataSplitImpl> split = MakeSplit(
+ {MakeDataFile("big.parquet", 20000, 5, FileSource::Compact())},
/*raw_convertible=*/true);
+ ASSERT_OK_AND_ASSIGN(
+ std::shared_ptr<Bytes> source_meta_bytes, ([&]() ->
Result<std::shared_ptr<Bytes>> {
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta,
+ PrimaryKeyIndexSourceMeta::Create(5,
source_files));
+ return source_meta.Serialize(pool_);
+ }()));
+ auto big_payload = std::make_shared<IndexFileMeta>(
+ "btree", "big-index-file", /*file_size=*/1, /*row_count=*/20000,
std::nullopt, std::nullopt,
+ GlobalIndexMeta(0, 19999, kPriceFieldId, std::nullopt, nullptr,
source_meta_bytes));
+ RoaringBitmap64 fragmented;
+ for (int64_t i = 0; i < 20000; i += 2) {
+ fragmented.Add(i);
+ }
+ PrimaryKeySortedIndexScan::ReaderFactory stub_factory =
+ [&fragmented](
+ const PrimaryKeySortedIndexScan::FilePlan& file,
+ const PrimaryKeyIndexDefinition& definition,
+ const PkSortedIndexGroup& group) ->
Result<std::shared_ptr<GlobalIndexReader>> {
+ return std::make_shared<StubGlobalIndexReader>(fragmented);
+ };
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(big_payload), PriceEqual(10),
stub_factory));
+ ASSERT_EQ(splits.size(), 1);
+ ASSERT_TRUE(std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]) ==
nullptr);
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, DeletionFileStaysAlignedWithIndexedFile)
{
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ DeletionFile deletion_file("dv-a", /*offset=*/0, /*length=*/16,
/*cardinality=*/1);
+ std::shared_ptr<DataSplitImpl> split = MakeSplit(
+ {MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()),
+ MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())},
+ /*raw_convertible=*/true, {std::optional<DeletionFile>(deletion_file),
std::nullopt});
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<std::shared_ptr<Split>> splits,
+ PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10),
PayloadReaderFactory()));
+ ASSERT_EQ(splits.size(), 1);
+ auto indexed_split =
std::dynamic_pointer_cast<IndexedSplitImpl>(splits[0]);
+ ASSERT_TRUE(indexed_split != nullptr);
+ auto inner_split =
std::dynamic_pointer_cast<DataSplitImpl>(indexed_split->GetDataSplit());
+ ASSERT_TRUE(inner_split != nullptr);
+ ASSERT_EQ(inner_split->DeletionFiles().size(), 1);
+ ASSERT_TRUE(inner_split->DeletionFiles()[0] != std::nullopt);
+ ASSERT_EQ(inner_split->DeletionFiles()[0].value().path, "dv-a");
+}
+
+TEST_F(PrimaryKeySortedIndexScanTest, SnapshotMismatchIsRejected) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<IndexFileMeta> payload,
BuildPayload());
+ std::vector<std::shared_ptr<DataFileMeta>> files = {
+ MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact())};
+ DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0,
base_path_,
+ std::move(files));
+ builder.WithSnapshot(kSnapshotId +
1).IsStreaming(false).RawConvertible(true);
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> split,
builder.Build());
+ Result<PrimaryKeySortedIndexScan::Plan> plan =
PrimaryKeySortedIndexScan::CreatePlan(
+ kSnapshotId, {split}, definitions_, MakeEntries(payload));
+ ASSERT_NOK(plan.status());
+}
Review Comment:
ASSERT_NOK(PrimaryKeySortedIndexScan::CreatePlan(
kSnapshotId, {split}, definitions_, MakeEntries(payload)));
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
