feat(scan): batch read primary-key tables with dynamic buckets

src/paimon/core/table/table.cpp

@@ -871,9 +871,12 @@ Result<std::unique_ptr<TableScan>> NewDataTableScan(
         ValidateRealtimeScan(*table_schema, core_options, *context, read_optimized));
     // validate options
     if (core_options.GetBucket() == -1) {
-        if (!table_schema->PrimaryKeys().empty()) {
-            return Status::NotImplemented(fmt::format(
-                "do not support pk table bucket={} in scan process", core_options.GetBucket()));
+        // Batch planning of dynamic bucket primary key tables only depends on the buckets
+        // recorded in the manifest, and deletion vectors are looked up per (partition, bucket).
+        if (!table_schema->PrimaryKeys().empty() && context->IsStreamingMode()) {
+            return Status::NotImplemented(
+                fmt::format("do not support pk table bucket={} in streaming scan process",
+                            core_options.GetBucket()));
         }
     } else if (core_options.GetBucket() < 1 &&
                !SchemaValidation::IsPostponeBucketTable(*table_schema, core_options.GetBucket())) {

test/inte/scan_and_read_inte_test.cpp

@@ -3124,6 +3124,47 @@ TEST_P(ScanAndReadInteTest, TestScanAndReadWithDisableIndex) {
     ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString();
 }
 
+TEST(ScanAndReadDynamicBucketInteTest, TestBatchReadPkDynamicBucketWithDeletionVectors) {
+    // Written by Java Paimon. Keys spread over several dynamic buckets, each with its own
+    // deletion vector index, so deletion vectors must be looked up per (partition, bucket).
+    const std::string table_path =
+        GetDataDir() + "/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/";
+    const std::shared_ptr<arrow::DataType> result_type =
+        arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()),
+                        arrow::field("id", arrow::int32()), arrow::field("v", arrow::utf8())});
+    const std::vector<std::pair<std::string, std::string>> cases = {
+        {"2", R"([[0, 1, "a"], [0, 2, "b"], [0, 3, "x"], [0, 4, "y"], [0, 5, "z"]])"},
+        {"4", R"([[0, 1, "c"], [0, 2, "b"], [0, 3, "x"], [0, 4, "w"], [0, 5, "z"]])"},
+        {"6", R"([[0, 1, "c"], [0, 3, "x"], [0, 4, "w"]])"},
+    };
+    for (const auto& [snapshot_id, expected_json] : cases) {
+        SCOPED_TRACE("snapshot " + snapshot_id);
+        ScanContextBuilder scan_context_builder(table_path);
+        scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID, snapshot_id);
+        ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context,
+                             scan_context_builder.Finish());
+        ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> table_scan,
+                             TableScan::Create(std::move(scan_context)));
+        ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, table_scan->CreatePlan());
+        ASSERT_GT(plan->Splits().size(), 1);
+
+        ReadContextBuilder read_context_builder(table_path);
+        ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context,
+                             read_context_builder.Finish());
+        ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> table_read,
+                             TableRead::Create(std::move(read_context)));
+        ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> batch_reader,
+                             table_read->CreateReader(plan->Splits()));
+        ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> actual,
+                             ReadResultCollector::CollectResult(std::move(batch_reader)));
+        const std::shared_ptr<arrow::Schema> sort_schema = arrow::schema(result_type->fields());
+        ASSERT_OK_AND_ASSIGN(actual, ReadResultCollector::SortArray(actual, sort_schema));
+        const std::shared_ptr<arrow::ChunkedArray> expected = std::make_shared<arrow::ChunkedArray>(
+            arrow::ipc::internal::json::ArrayFromJSON(result_type, expected_json).ValueOrDie());
+        ASSERT_TRUE(expected->Equals(actual)) << actual->ToString();
+    }
+}
+
 TEST(ScanAndReadFileIndexInteTest, TestPkDvEmbeddedBitmapFiltersScanAndRead) {
     std::unique_ptr<UniqueTestDirectory> dir = UniqueTestDirectory::Create("local");
     ASSERT_NE(nullptr, dir);

test/inte/scan_inte_test.cpp

@@ -1730,7 +1730,7 @@ TEST_P(ScanInteTest, TestScanPKWithInvalidOptions) {
     context_builder.AddOption(Options::BUCKET, "-1").WithStreamingMode(true);
     ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(context_builder));
     ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)),
-                        "do not support pk table bucket=-1 in scan process");
+                        "do not support pk table bucket=-1 in streaming scan process");
 }
 
 TEST_P(ScanInteTest, TestReadWithNoSnapshot) {

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/README

@@ -0,0 +1,21 @@
+Written by Spark 3.5.3 with Paimon 1.3.2.
+
+id:int not null v:string
+primary key: id
+bucket: -1 (dynamic bucket)
+dynamic-bucket.target-row-num = 2
+deletion-vectors.enabled = true
+file.format = parquet
+
+CREATE TABLE dv (id INT, v STRING)
+TBLPROPERTIES ('primary-key' = 'id', 'bucket' = '-1', 'deletion-vectors.enabled' = 'true',
+               'file.format' = 'parquet', 'dynamic-bucket.target-row-num' = '2');
+
+INSERT INTO dv VALUES (1, 'a'), (2, 'b'), (3, 'x'), (4, 'y'), (5, 'z');
+-- snapshot 2: (1, a) (2, b) (3, x) (4, y) (5, z)
+INSERT INTO dv VALUES (1, 'c'), (4, 'w');
+-- snapshot 4: (1, c) (2, b) (3, x) (4, w) (5, z)
+DELETE FROM dv WHERE id IN (2, 5);
+-- snapshot 6: (1, c) (3, x) (4, w)
+
+Keys are spread over several buckets, and each bucket keeps its own deletion vector index.

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/index/index-226a01b7-1ea1-410b-a0f4-c3d72dac3dfb-0

@@ -0,0 +1 @@
+ÒåG

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/index/index-48eb8ead-486d-4255-9698-f2fc4b10e3cb-0

@@ -0,0 +1 @@
+VGpÊ

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/index/index-6edcb848-8159-4ae7-826d-8995c1a6f454-0

@@ -0,0 +1 @@
+‚Hô

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/index/index-ed55898c-8f16-432f-8fa3-95727d920a0a-0

@@ -0,0 +1 @@
+Oä»ð

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/index/index-f787a693-854f-4b72-99e0-40217aaca953-0

@@ -0,0 +1 @@
+WYùž

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/schema/schema-0

@@ -0,0 +1,24 @@
+{
+  "version" : 3,
+  "id" : 0,
+  "fields" : [ {
+    "id" : 0,
+    "name" : "id",
+    "type" : "INT NOT NULL"
+  }, {
+    "id" : 1,
+    "name" : "v",
+    "type" : "STRING"
+  } ],
+  "highestFieldId" : 1,
+  "partitionKeys" : [ ],
+  "primaryKeys" : [ "id" ],
+  "options" : {
+    "bucket" : "-1",
+    "owner" : "root",
+    "file.format" : "parquet",
+    "dynamic-bucket.target-row-num" : "2",
+    "deletion-vectors.enabled" : "true"
+  },
+  "timeMillis" : 1792572933936
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/EARLIEST

@@ -0,0 +1 @@
+1

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/LATEST

@@ -0,0 +1 @@
+6

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-1

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 1,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-0014d5c7-8244-459c-ba6c-3e66042c6b22-0",
+  "baseManifestListSize" : 884,
+  "deltaManifestList" : "manifest-list-0014d5c7-8244-459c-ba6c-3e66042c6b22-1",
+  "deltaManifestListSize" : 990,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "86f3fe45-aa7d-41e3-a32b-f49670b910da",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "APPEND",
+  "timeMillis" : 1792572934818,
+  "logOffsets" : { },
+  "totalRecordCount" : 5,
+  "deltaRecordCount" : 5,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-2

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 2,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-0014d5c7-8244-459c-ba6c-3e66042c6b22-2",
+  "baseManifestListSize" : 990,
+  "deltaManifestList" : "manifest-list-0014d5c7-8244-459c-ba6c-3e66042c6b22-3",
+  "deltaManifestListSize" : 990,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "86f3fe45-aa7d-41e3-a32b-f49670b910da",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "COMPACT",
+  "timeMillis" : 1792572934825,
+  "logOffsets" : { },
+  "totalRecordCount" : 5,
+  "deltaRecordCount" : 0,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-3

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 3,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-19ad07c9-fcfc-4be6-b685-d6918922e81b-0",
+  "baseManifestListSize" : 1001,
+  "deltaManifestList" : "manifest-list-19ad07c9-fcfc-4be6-b685-d6918922e81b-1",
+  "deltaManifestListSize" : 991,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "b110fa6e-4595-43c0-b479-8756f91b126b",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "APPEND",
+  "timeMillis" : 1792572936021,
+  "logOffsets" : { },
+  "totalRecordCount" : 7,
+  "deltaRecordCount" : 2,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-4

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 4,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-19ad07c9-fcfc-4be6-b685-d6918922e81b-2",
+  "baseManifestListSize" : 1043,
+  "deltaManifestList" : "manifest-list-19ad07c9-fcfc-4be6-b685-d6918922e81b-3",
+  "deltaManifestListSize" : 991,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "b110fa6e-4595-43c0-b479-8756f91b126b",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "COMPACT",
+  "timeMillis" : 1792572936030,
+  "logOffsets" : { },
+  "totalRecordCount" : 5,
+  "deltaRecordCount" : -2,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-5

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 5,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-51b6b5dd-86de-4a97-9a90-86dc8382f85e-0",
+  "baseManifestListSize" : 1051,
+  "deltaManifestList" : "manifest-list-51b6b5dd-86de-4a97-9a90-86dc8382f85e-1",
+  "deltaManifestListSize" : 991,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "9b190bb7-7a17-4f78-9035-60a019950355",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "APPEND",
+  "timeMillis" : 1792572937421,
+  "logOffsets" : { },
+  "totalRecordCount" : 7,
+  "deltaRecordCount" : 2,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}

test/test_data/parquet/pk_dynamic_bucket_dv.db/pk_dynamic_bucket_dv/snapshot/snapshot-6

@@ -0,0 +1,20 @@
+{
+  "version" : 3,
+  "id" : 6,
+  "schemaId" : 0,
+  "baseManifestList" : "manifest-list-51b6b5dd-86de-4a97-9a90-86dc8382f85e-2",
+  "baseManifestListSize" : 1091,
+  "deltaManifestList" : "manifest-list-51b6b5dd-86de-4a97-9a90-86dc8382f85e-3",
+  "deltaManifestListSize" : 991,
+  "changelogManifestList" : null,
+  "indexManifest" : "index-manifest-243d52b5-0cfc-4039-bf11-7436f05ebd99-0",
+  "commitUser" : "9b190bb7-7a17-4f78-9035-60a019950355",
+  "commitIdentifier" : 9223372036854775807,
+  "commitKind" : "COMPACT",
+  "timeMillis" : 1792572937428,
+  "logOffsets" : { },
+  "totalRecordCount" : 3,
+  "deltaRecordCount" : -4,
+  "changelogRecordCount" : 0,
+  "nextRowId" : 0
+}