feat(scan): prune append buckets from equality predicates

Author: wangyong.alen

docs/source/api/scan.rst

@@ -21,6 +21,17 @@ Scan
 
 .. _cpp-api-scan:
 
+Bucket pruning
+==============
+
+For fixed-bucket append tables, an equality predicate on every bucket key lets
+the scan derive the target bucket using the table's bucket function. Other buckets
+are excluded from the scan plan without requiring an explicit bucket ID from the
+caller. An explicit bucket filter takes precedence. Queries that do not constrain
+all bucket keys with equality, and bucket-unaware tables, keep the existing scan
+behavior. Inferred pruning applies only to files matching the scan schema and
+bucket count; older layouts retain the existing filtering behavior.
+
 Interface
 =========
 

src/paimon/core/operation/append_only_file_store_scan.cpp

@@ -26,11 +26,13 @@
 #include <set>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "arrow/type.h"
 #include "fmt/format.h"
 #include "paimon/common/predicate/predicate_filter.h"
 #include "paimon/common/types/data_field.h"
+#include "paimon/core/bucket/bucket_select_converter.h"
 #include "paimon/core/core_options.h"
 #include "paimon/core/io/data_file_meta.h"
 #include "paimon/core/io/file_index_evaluator.h"
@@ -42,6 +44,7 @@
 #include "paimon/core/utils/field_mapping.h"
 #include "paimon/file_index/file_index_result.h"
 #include "paimon/predicate/predicate_utils.h"
+#include "paimon/scan_context.h"
 #include "paimon/status.h"
 
 namespace paimon {
@@ -66,6 +69,21 @@ Result<std::unique_ptr<AppendOnlyFileStoreScan>> AppendOnlyFileStoreScan::Create
                                     table_schema, arrow_schema, core_options, executor, pool));
     PAIMON_RETURN_NOT_OK(
         scan->SplitAndSetFilter(table_schema->PartitionKeys(), arrow_schema, scan_filters));
+    const auto& bucket_keys = table_schema->BucketKeys();
+    int32_t num_buckets = core_options.GetBucket();
+    if (scan->predicates_ && !scan_filters->GetBucketFilter().has_value() && num_buckets > 0 &&
+        !bucket_keys.empty()) {
+        std::vector<std::shared_ptr<arrow::DataType>> bucket_key_types;
+        bucket_key_types.reserve(bucket_keys.size());
+        for (const auto& key : bucket_keys) {
+            PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(key));
+            bucket_key_types.push_back(field.Type());
+        }
+        PAIMON_ASSIGN_OR_RAISE(scan->predicate_bucket_,
+                               BucketSelectConverter::Convert(
+                                   scan->predicates_, bucket_keys, bucket_key_types,
+                                   core_options.GetBucketFunctionType(), num_buckets, pool.get()));
+    }
     return scan;
 }
 
@@ -86,6 +104,12 @@ Result<bool> AppendOnlyFileStoreScan::FilterByStats(const ManifestEntry& entry)
     if (!predicates_) {
         return true;
     }
+    // A historical file may use a different schema or bucket count after a rescale.
+    // Keep the inferred bucket separate from the caller's explicit bucket filter.
+    if (predicate_bucket_ && entry.TotalBuckets() == core_options_.GetBucket() &&
+        entry.File()->schema_id == table_schema_->Id() && entry.Bucket() != *predicate_bucket_) {
+        return false;
+    }
     const auto& meta = entry.File();
     std::shared_ptr<TableSchema> data_schema = table_schema_;
     std::shared_ptr<Predicate> trimmed_predicates = predicates_;

src/paimon/core/operation/append_only_file_store_scan.h

@@ -18,7 +18,9 @@
 
 #pragma once
 
+#include <cstdint>
 #include <memory>
+#include <optional>
 #include <vector>
 
 #include "paimon/core/operation/file_store_scan.h"
@@ -79,5 +81,6 @@ class AppendOnlyFileStoreScan : public FileStoreScan {
 
  private:
     std::shared_ptr<SimpleStatsEvolutions> evolutions_;
+    std::optional<int32_t> predicate_bucket_;
 };
 }  // namespace paimon

src/paimon/core/operation/append_only_file_store_scan_test.cpp

@@ -20,15 +20,18 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <map>
 #include <optional>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "arrow/type.h"
 #include "gtest/gtest.h"
 #include "paimon/common/data/binary_row.h"
 #include "paimon/common/data/binary_row_writer.h"
 #include "paimon/common/io/cache/lru_cache.h"
+#include "paimon/core/bucket/default_bucket_function.h"
 #include "paimon/core/manifest/manifest_entry.h"
 #include "paimon/core/manifest/partition_entry.h"
 #include "paimon/core/schema/schema_manager.h"
@@ -37,6 +40,7 @@
 #include "paimon/core/table/source/abstract_table_scan.h"
 #include "paimon/core/table/source/snapshot/snapshot_reader.h"
 #include "paimon/defs.h"
+#include "paimon/executor.h"
 #include "paimon/fs/local/local_file_system.h"
 #include "paimon/memory/memory_pool.h"
 #include "paimon/metrics.h"
@@ -46,10 +50,105 @@
 #include "paimon/status.h"
 #include "paimon/table/source/scan_metrics.h"
 #include "paimon/table/source/table_scan.h"
+#include "paimon/testing/utils/binary_row_generator.h"
 #include "paimon/testing/utils/testharness.h"
 #include "paimon/testing/utils/timezone_guard.h"
 namespace paimon::test {
 
+class AppendBucketPruningTest : public testing::Test {
+ public:
+    Result<std::unique_ptr<AppendOnlyFileStoreScan>> CreateScan(
+        const std::shared_ptr<Predicate>& predicate,
+        const std::optional<int32_t>& bucket = std::nullopt) const {
+        auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(
+            {DataField(0, arrow::field("rowkey", arrow::utf8())),
+             DataField(1, arrow::field("value", arrow::int32()))});
+        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> schema,
+                               TableSchema::Create(schema_id_, arrow_schema, {}, {}, options_));
+        PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_));
+        auto filters = std::make_shared<ScanFilter>(
+            predicate, std::vector<std::map<std::string, std::string>>(), bucket);
+        return AppendOnlyFileStoreScan::Create(nullptr, schema_manager_, nullptr, nullptr, schema,
+                                               arrow_schema, filters, options,
+                                               GetGlobalDefaultExecutor(), pool_);
+    }
+
+    void CheckBuckets(const std::shared_ptr<Predicate>& predicate,
+                      const std::optional<int32_t>& expected_bucket,
+                      const std::optional<int32_t>& explicit_bucket = std::nullopt,
+                      int32_t total_buckets = kNumBuckets) const {
+        ASSERT_OK_AND_ASSIGN(auto scan, CreateScan(predicate, explicit_bucket));
+        // Every file's stats match the lookup, so only bucket pruning can discard a file.
+        SimpleStats stats = BinaryRowGenerator::GenerateStats(
+            {std::string("a"), 0}, {std::string("z"), 100}, {0, 0}, pool_.get());
+        ASSERT_OK_AND_ASSIGN(
+            auto file,
+            DataFileMeta::ForAppend("data.parquet", 100, 10, stats, 0, 9, 0, std::nullopt,
+                                    std::nullopt, std::nullopt, std::nullopt, std::nullopt));
+        for (int32_t bucket = 0; bucket < total_buckets; ++bucket) {
+            ManifestEntry entry(FileKind::Add(), BinaryRow::EmptyRow(), bucket, total_buckets,
+                                file);
+            ASSERT_OK_AND_ASSIGN(auto stats_scan, CreateScan(predicate, bucket));
+            ASSERT_OK_AND_ASSIGN(bool stats_match, stats_scan->FilterByStats(entry));
+            ASSERT_TRUE(stats_match);
+            ASSERT_OK_AND_ASSIGN(bool keep, scan->FilterManifestEntry(entry));
+            ASSERT_EQ(keep, !expected_bucket.has_value() || bucket == expected_bucket.value());
+        }
+    }
+
+    std::shared_ptr<Predicate> KeyEquals() const {
+        return PredicateBuilder::Equal(0, "rowkey", FieldType::STRING,
+                                       Literal(FieldType::STRING, "key", 3));
+    }
+
+    static constexpr int32_t kNumBuckets = 4;
+    int64_t schema_id_ = 0;
+    std::shared_ptr<SchemaManager> schema_manager_;
+    std::shared_ptr<MemoryPool> pool_ = GetDefaultPool();
+    std::map<std::string, std::string> options_ = {{Options::BUCKET, "4"},
+                                                   {Options::BUCKET_KEY, "rowkey"}};
+};
+
+TEST_F(AppendBucketPruningTest, PrunesStringKeyLookup) {
+    BinaryRow key = BinaryRowGenerator::GenerateRow({std::string("key")}, pool_.get());
+    int32_t expected_bucket = DefaultBucketFunction().Bucket(key, kNumBuckets);
+    CheckBuckets(KeyEquals(), expected_bucket);
+}
+
+TEST_F(AppendBucketPruningTest, PreservesExplicitBucketFilter) {
+    BinaryRow key = BinaryRowGenerator::GenerateRow({std::string("key")}, pool_.get());
+    int32_t other_bucket = (DefaultBucketFunction().Bucket(key, kNumBuckets) + 1) % kNumBuckets;
+    CheckBuckets(KeyEquals(), other_bucket, other_bucket);
+}
+
+TEST_F(AppendBucketPruningTest, DoesNotPruneWithoutCompleteEqualKeys) {
+    CheckBuckets(nullptr, std::nullopt);
+    CheckBuckets(PredicateBuilder::GreaterThan(0, "rowkey", FieldType::STRING,
+                                               Literal(FieldType::STRING, "key", 3)),
+                 std::nullopt);
+    options_[Options::BUCKET_KEY] = "rowkey,value";
+    CheckBuckets(KeyEquals(), std::nullopt);
+}
+
+TEST_F(AppendBucketPruningTest, DoesNotPruneBucketUnawareTable) {
+    options_[Options::BUCKET] = "-1";
+    CheckBuckets(KeyEquals(), std::nullopt);
+}
+
+TEST_F(AppendBucketPruningTest, DoesNotPruneDifferentBucketCount) {
+    CheckBuckets(KeyEquals(), std::nullopt, std::nullopt, 8);
+}
+
+TEST_F(AppendBucketPruningTest, DoesNotPruneDifferentSchema) {
+    auto test_dir = UniqueTestDirectory::Create("local");
+    schema_manager_ =
+        std::make_shared<SchemaManager>(std::make_shared<LocalFileSystem>(), test_dir->Str());
+    ASSERT_OK_AND_ASSIGN(auto scan, CreateScan(nullptr));
+    ASSERT_OK(schema_manager_->CreateTable(scan->schema_, {}, {}, options_));
+    schema_id_ = 1;
+    CheckBuckets(KeyEquals(), std::nullopt);
+}
+
 TEST(AppendOnlyFileStoreScanTest, TestReconstructPredicateWithNonCastedFields) {
     std::string table_root =
         paimon::test::GetDataDir() +

test/inte/scan_and_read_inte_test.cpp

@@ -29,8 +29,12 @@
 #include <vector>
 
 #include "arrow/api.h"
+#include "arrow/c/bridge.h"
 #include "arrow/ipc/json_simple.h"
+#include "fmt/format.h"
+#include "fmt/ranges.h"
 #include "gtest/gtest.h"
+#include "paimon/bucket/bucket_id_calculator.h"
 #include "paimon/common/factories/io_hook.h"
 #include "paimon/common/io/cache/lru_cache.h"
 #include "paimon/common/table/special_fields.h"
@@ -421,6 +425,93 @@ TEST_P(ScanAndReadInteTest, TestWithAppendSnapshot3) {
     ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString();
 }
 
+TEST_P(ScanAndReadInteTest, TestWithAppendBucketKeyPointLookup) {
+    for (const auto& key_type : {arrow::utf8(), arrow::binary()}) {
+        SCOPED_TRACE(key_type->ToString());
+        auto test_dir = UniqueTestDirectory::Create("local");
+        arrow::FieldVector fields = {arrow::field("rowkey", key_type),
+                                     arrow::field("value", arrow::int32())};
+        std::map<std::string, std::string> options = {{Options::FILE_FORMAT, FileFormat()},
+                                                      {Options::FILE_SYSTEM, "local"},
+                                                      {Options::BUCKET, "4"},
+                                                      {Options::BUCKET_KEY, "rowkey"}};
+        ASSERT_OK_AND_ASSIGN(
+            auto helper, TestHelper::Create(test_dir->Str(), arrow::schema(fields), {}, {}, options,
+                                            /*is_streaming_mode=*/false));
+        std::string table_path = test_dir->Str() + "/foo.db/bar";
+        // Route writes through the public bucket calculator, as a KV client does.
+        constexpr int32_t kNumKeys = 64;
+        constexpr int32_t kNumBuckets = 4;
+        std::vector<std::string> keys;
+        for (int32_t i = 0; i < kNumKeys; ++i) {
+            keys.push_back(fmt::format(R"(["key{:03}"])", i));
+        }
+        auto key_array = arrow::ipc::internal::json::ArrayFromJSON(
+                             arrow::struct_({fields[0]}), fmt::format("[{}]", fmt::join(keys, ",")))
+                             .ValueOrDie();
+        ::ArrowArray c_keys;
+        ::ArrowSchema c_schema;
+        ASSERT_TRUE(arrow::ExportArray(*key_array, &c_keys, &c_schema).ok());
+        ASSERT_OK_AND_ASSIGN(auto calculator,
+                             BucketIdCalculator::Create(false, kNumBuckets, GetDefaultPool()));
+        std::vector<int32_t> bucket_ids(kNumKeys);
+        ASSERT_OK(calculator->CalculateBucketIds(&c_keys, &c_schema, bucket_ids.data()));
+        std::vector<std::vector<std::string>> rows(kNumBuckets);
+        for (int32_t i = 0; i < kNumKeys; ++i) {
+            rows[bucket_ids[i]].push_back(fmt::format(R"(["key{:03}", {}])", i, i));
+        }
+        std::vector<std::unique_ptr<RecordBatch>> batches;
+        for (int32_t bucket = 0; bucket < kNumBuckets; ++bucket) {
+            ASSERT_FALSE(rows[bucket].empty());
+            ASSERT_OK_AND_ASSIGN(
+                auto batch, TestHelper::MakeRecordBatch(
+                                arrow::struct_(fields),
+                                fmt::format("[{}]", fmt::join(rows[bucket], ",")), {}, bucket, {}));
+            batches.push_back(std::move(batch));
+        }
+        ASSERT_OK(helper->WriteAndCommit(std::move(batches), 0, std::nullopt));
+
+        FieldType field_type =
+            key_type->id() == arrow::Type::STRING ? FieldType::STRING : FieldType::BINARY;
+        auto predicate =
+            PredicateBuilder::Equal(0, "rowkey", field_type, Literal(field_type, "key032", 6));
+        // All buckets have overlapping stats for this key. An explicit filter bypasses
+        // inference, proving that ordinary statistics cannot account for the pruning.
+        for (int32_t bucket = 0; bucket < kNumBuckets; ++bucket) {
+            ScanContextBuilder builder(table_path);
+            builder.SetPredicate(predicate).SetBucketFilter(bucket);
+            ASSERT_OK_AND_ASSIGN(auto context, FinishScanContext(builder));
+            ASSERT_OK_AND_ASSIGN(auto scan, TableScan::Create(std::move(context)));
+            ASSERT_OK_AND_ASSIGN(auto plan, scan->CreatePlan());
+            ASSERT_FALSE(plan->Splits().empty());
+        }
+        ScanContextBuilder scan_builder(table_path);
+        scan_builder.SetPredicate(predicate);
+        ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_builder));
+        ASSERT_OK_AND_ASSIGN(auto scan, TableScan::Create(std::move(scan_context)));
+        ASSERT_OK_AND_ASSIGN(auto plan, scan->CreatePlan());
+        ASSERT_FALSE(plan->Splits().empty());
+        for (const auto& split : plan->Splits()) {
+            auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
+            ASSERT_TRUE(data_split);
+            ASSERT_EQ(data_split->Bucket(), bucket_ids[32]);
+        }
+        ReadContextBuilder read_builder(table_path);
+        AddReadOptionsForPrefetch(&read_builder);
+        read_builder.SetPredicate(predicate).EnablePredicateFilter(true);
+        ASSERT_OK_AND_ASSIGN(auto read_context, read_builder.Finish());
+        ASSERT_OK_AND_ASSIGN(auto read, TableRead::Create(std::move(read_context)));
+        ASSERT_OK_AND_ASSIGN(auto reader, read->CreateReader(plan->Splits()));
+        ASSERT_OK_AND_ASSIGN(auto result, ReadResultCollector::CollectResult(std::move(reader)));
+        fields.insert(fields.begin(), arrow::field("_VALUE_KIND", arrow::int8()));
+        auto expected_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields),
+                                                                        R"([[0, "key032", 32]])")
+                                  .ValueOrDie();
+        auto expected = std::make_shared<arrow::ChunkedArray>(expected_array);
+        ASSERT_TRUE(expected->Equals(result)) << result->ToString();
+    }
+}
+
 TEST_P(ScanAndReadInteTest, TestWithAppendSnapshot5) {
     auto file_format = FileFormat();
     std::string table_path = GetDataDir() + "/" + file_format + "/append_09.db/append_09";