| /* |
| * 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 <cstdint> |
| #include <map> |
| #include <memory> |
| #include <string> |
| #include <utility> |
| #include <vector> |
| |
| #include "arrow/api.h" |
| #include "arrow/c/abi.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/ipc/json_simple.h" |
| #include "gtest/gtest.h" |
| #include "paimon/commit_context.h" |
| #include "paimon/common/types/data_field.h" |
| #include "paimon/common/utils/path_util.h" |
| #include "paimon/common/utils/string_utils.h" |
| #include "paimon/core/schema/schema_manager.h" |
| #include "paimon/core/schema/table_schema.h" |
| #include "paimon/defs.h" |
| #include "paimon/file_store_commit.h" |
| #include "paimon/file_store_write.h" |
| #include "paimon/fs/file_system_factory.h" |
| #include "paimon/predicate/literal.h" |
| #include "paimon/predicate/predicate_builder.h" |
| #include "paimon/read_context.h" |
| #include "paimon/reader/batch_reader.h" |
| #include "paimon/record_batch.h" |
| #include "paimon/result.h" |
| #include "paimon/scan_context.h" |
| #include "paimon/status.h" |
| #include "paimon/table/source/startup_mode.h" |
| #include "paimon/table/source/table_read.h" |
| #include "paimon/table/source/table_scan.h" |
| #include "paimon/testing/utils/dict_array_converter.h" |
| #include "paimon/testing/utils/read_result_collector.h" |
| #include "paimon/testing/utils/test_helper.h" |
| #include "paimon/testing/utils/testharness.h" |
| #include "paimon/write_context.h" |
| |
| namespace paimon { |
| class DataSplit; |
| class RecordBatch; |
| } // namespace paimon |
| |
| namespace paimon::test { |
| |
| class NestedColumnPruningInteTest : public ::testing::Test, |
| public ::testing::WithParamInterface<std::string> { |
| void SetUp() override { |
| file_format_ = GetParam(); |
| dir_ = UniqueTestDirectory::Create("local"); |
| test_dir_ = dir_->Str(); |
| table_path_ = PathUtil::JoinPath(test_dir_, "foo.db/bar"); |
| } |
| void TearDown() override { |
| dir_.reset(); |
| } |
| |
| void AssertChunkedArrayEquals(const std::shared_ptr<arrow::ChunkedArray>& expected, |
| const std::shared_ptr<arrow::ChunkedArray>& actual) const { |
| arrow::EqualOptions equal_options = arrow::EqualOptions::Defaults(); |
| bool is_equal = expected->Equals(actual, equal_options.diff_sink(&std::cout)); |
| if (!is_equal) { |
| std::cout << "[expected_type] " << expected->type()->ToString() << std::endl; |
| std::cout << "[actual_type] " << actual->type()->ToString() << std::endl; |
| std::cout << "[expected] " << expected->ToString() << std::endl; |
| std::cout << "[actual] " << actual->ToString() << std::endl; |
| } |
| ASSERT_TRUE(is_equal); |
| } |
| |
| void ScanReadAndCheck(const std::string& table_path, |
| const std::shared_ptr<arrow::Schema>& expected_schema, |
| const std::string& expected_json, |
| const std::shared_ptr<Predicate>& predicate = nullptr) const { |
| ScanContextBuilder scan_context_builder(table_path); |
| scan_context_builder.AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()); |
| if (predicate) { |
| scan_context_builder.SetPredicate(predicate); |
| } |
| ASSERT_OK_AND_ASSIGN(auto scan_context, scan_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); |
| ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); |
| ASSERT_FALSE(result_plan->Splits().empty()); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*expected_schema, c_schema.get()).ok()); |
| ReadContextBuilder read_context_builder(table_path); |
| read_context_builder.SetReadSchema(std::move(c_schema)); |
| if (predicate) { |
| read_context_builder.SetPredicate(predicate); |
| } |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); |
| ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); |
| |
| arrow::FieldVector expected_fields = expected_schema->fields(); |
| expected_fields.insert(expected_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); |
| auto expected_type = arrow::struct_(expected_fields); |
| auto expected = std::make_shared<arrow::ChunkedArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(expected_type, expected_json).ValueOrDie()); |
| AssertChunkedArrayEquals(expected, actual); |
| } |
| |
| protected: |
| std::string file_format_; |
| std::string test_dir_; |
| std::string table_path_; |
| std::unique_ptr<UniqueTestDirectory> dir_; |
| }; |
| |
| // Test: Table has struct field with 3 sub-fields, read only 1 sub-field via SetReadSchema. |
| TEST_P(NestedColumnPruningInteTest, PruneStructSubFields) { |
| // Table schema: f0 (int32), f1 (struct{a: int32, b: utf8, c: float64}) |
| auto struct_type = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("b", arrow::utf8()), |
| arrow::field("c", arrow::float64()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| // Write data |
| std::string data = R"([ |
| [1, [10, "hello", 1.1]], |
| [2, [20, "world", 2.2]], |
| [3, [30, "foo", 3.3]], |
| [4, [40, "bar", 4.4]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Build projected schema: only read f0 (full) and f1.a (sub-field of struct) |
| auto pruned_struct_type = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", pruned_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [10]], |
| [0, 2, [20]], |
| [0, 3, [30]], |
| [0, 4, [40]] |
| ])"); |
| } |
| |
| // Test: Projecting a STRUCT column as empty struct should return this column |
| // as all null values. |
| TEST_P(NestedColumnPruningInteTest, ProjectStructColumnAsEmptyStructReturnsNullColumn) { |
| auto struct_type = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("b", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "x"]], |
| [2, [22, "y"]], |
| [3, [33, "z"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Project f1 as empty struct. |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({})), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, null], |
| [0, 2, null], |
| [0, 3, null] |
| ])"); |
| } |
| |
| // Test: Two top-level struct columns have the same nested field name; projection should |
| // distinguish by parent column. |
| TEST_P(NestedColumnPruningInteTest, PruneSameNestedFieldNameFromDifferentStructColumns) { |
| // Table schema: f0 (int32), s0 (struct{f1: int32, a: utf8}), s1 (struct{f1: int32, b: utf8}) |
| auto s0_type = arrow::struct_({ |
| arrow::field("f1", arrow::int32()), |
| arrow::field("a", arrow::utf8()), |
| }); |
| auto s1_type = arrow::struct_({ |
| arrow::field("f1", arrow::int32()), |
| arrow::field("b", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("s0", s0_type), |
| arrow::field("s1", s1_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "left-1"], [101, "right-1"]], |
| [2, [22, "left-2"], [202, "right-2"]], |
| [3, [33, "left-3"], [303, "right-3"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Project only s0.f1 and s1.f1; both nested field names are identical. |
| auto projected_s0 = arrow::struct_({arrow::field("f1", arrow::int32())}); |
| auto projected_s1 = arrow::struct_({arrow::field("f1", arrow::int32())}); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("s0", projected_s0), |
| arrow::field("s1", projected_s1), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [11], [101]], |
| [0, 2, [22], [202]], |
| [0, 3, [33], [303]] |
| ])"); |
| } |
| |
| // Test: Querying only non-existent struct sub-fields should fail fast. |
| TEST_P(NestedColumnPruningInteTest, QueryStructSubFieldsAllNonExistent) { |
| // Table schema: f0 (int32), f1 (struct{f1: int32, f2: utf8, f3: float64}) |
| auto struct_type = arrow::struct_({ |
| arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::utf8()), |
| arrow::field("f3", arrow::float64()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "a", 1.1]], |
| [2, [22, "b", 2.2]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| // Query struct sub-fields that do not exist in table schema. |
| auto projected_struct_type = arrow::struct_({ |
| arrow::field("f4", arrow::int32()), |
| arrow::field("f5", arrow::int64()), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", projected_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), |
| "does not support schema evolution inside struct"); |
| } |
| |
| // Test: Querying a mix of existent and non-existent struct sub-fields should fail fast. |
| TEST_P(NestedColumnPruningInteTest, QueryStructSubFieldsWithNonExistentField) { |
| // Table schema: f0 (int32), f1 (struct{f1: int32, f2: utf8, f3: float64}) |
| auto struct_type = arrow::struct_({ |
| arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::utf8()), |
| arrow::field("f3", arrow::float64()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "a", 1.1]], |
| [2, [22, "b", 2.2]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| // Query struct sub-fields f2,f3,f4 where f4 does not exist in table schema. |
| auto projected_struct_type = arrow::struct_({ |
| arrow::field("f2", arrow::utf8()), |
| arrow::field("f3", arrow::float64()), |
| arrow::field("f4", arrow::int32()), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", projected_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), |
| "does not support schema evolution inside struct"); |
| } |
| |
| // Test: Nested schema divergence (simulated evolution mismatch) must fail fast |
| // instead of silently skipping nested fields. |
| TEST_P(NestedColumnPruningInteTest, QueryStructSubFieldsWithTypeMismatchShouldFail) { |
| // File schema (old): f1.a is INT32. |
| auto struct_type = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("b", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "x"]], |
| [2, [22, "y"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Persist schema-2 on disk so table latest schema diverges from old data files. |
| std::string file_system_identifier = "local"; |
| auto fs_iter = options.find(Options::FILE_SYSTEM); |
| if (fs_iter != options.end()) { |
| file_system_identifier = StringUtils::ToLowerCase(fs_iter->second); |
| } |
| ASSERT_OK_AND_ASSIGN(auto file_system, |
| FileSystemFactory::Get(file_system_identifier, table_path_, options)); |
| std::shared_ptr<FileSystem> schema_fs(std::move(file_system)); |
| |
| SchemaManager schema_manager(schema_fs, table_path_); |
| ASSERT_OK_AND_ASSIGN(auto latest_schema_opt, schema_manager.Latest()); |
| ASSERT_TRUE(latest_schema_opt.has_value()); |
| auto latest_schema = latest_schema_opt.value(); |
| |
| auto schema_v2_arrow = arrow::schema({ |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({ |
| arrow::field("a", arrow::utf8()), |
| arrow::field("b", arrow::utf8()), |
| })), |
| }); |
| ASSERT_OK_AND_ASSIGN( |
| auto schema_v2, |
| TableSchema::Create(/*schema_id=*/latest_schema->Id() + 1, schema_v2_arrow, |
| latest_schema->PartitionKeys(), latest_schema->PrimaryKeys(), |
| latest_schema->Options())); |
| ASSERT_OK_AND_ASSIGN(auto schema_v2_json, schema_v2->ToJsonString()); |
| auto schema_v2_path = PathUtil::JoinPath(schema_manager.SchemaDirectory(), |
| "schema-" + std::to_string(schema_v2->Id())); |
| ASSERT_OK(schema_fs->AtomicStore(schema_v2_path, schema_v2_json)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| |
| ASSERT_NOK_WITH_MSG(table_read->CreateReader(data_splits), |
| "PruneDataType nested field type mismatch for 'a': read string vs " |
| "data int32"); |
| } |
| |
| // Test: With SetReadSchema using the new schema, context build should pass, |
| // and mismatch against old file type should be rejected in reader creation. |
| TEST_P(NestedColumnPruningInteTest, |
| QueryStructSubFieldsWithTypeMismatchAndSetReadSchemaFailAtContext) { |
| // File schema (old): f1.a is INT32. |
| auto struct_type = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("b", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [11, "x"]], |
| [2, [22, "y"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Persist schema-2 on disk so table latest schema diverges from old data files. |
| std::string file_system_identifier = "local"; |
| auto fs_iter = options.find(Options::FILE_SYSTEM); |
| if (fs_iter != options.end()) { |
| file_system_identifier = StringUtils::ToLowerCase(fs_iter->second); |
| } |
| ASSERT_OK_AND_ASSIGN(auto file_system, |
| FileSystemFactory::Get(file_system_identifier, table_path_, options)); |
| std::shared_ptr<FileSystem> schema_fs(std::move(file_system)); |
| |
| SchemaManager schema_manager(schema_fs, table_path_); |
| ASSERT_OK_AND_ASSIGN(auto latest_schema_opt, schema_manager.Latest()); |
| ASSERT_TRUE(latest_schema_opt.has_value()); |
| auto latest_schema = latest_schema_opt.value(); |
| |
| auto schema_v2_arrow = arrow::schema({ |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({ |
| arrow::field("a", arrow::utf8()), |
| arrow::field("b", arrow::utf8()), |
| })), |
| }); |
| ASSERT_OK_AND_ASSIGN( |
| auto schema_v2, |
| TableSchema::Create(/*schema_id=*/latest_schema->Id() + 1, schema_v2_arrow, |
| latest_schema->PartitionKeys(), latest_schema->PrimaryKeys(), |
| latest_schema->Options())); |
| ASSERT_OK_AND_ASSIGN(auto schema_v2_json, schema_v2->ToJsonString()); |
| auto schema_v2_path = PathUtil::JoinPath(schema_manager.SchemaDirectory(), |
| "schema-" + std::to_string(schema_v2->Id())); |
| ASSERT_OK(schema_fs->AtomicStore(schema_v2_path, schema_v2_json)); |
| |
| // User-provided read schema uses latest nested type (a:string). |
| auto projected_schema = arrow::schema({ |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({ |
| arrow::field("a", arrow::utf8()), |
| arrow::field("b", arrow::utf8()), |
| })), |
| }); |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| ASSERT_OK_AND_ASSIGN(auto data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| ASSERT_NOK_WITH_MSG(table_read->CreateReader(data_splits), |
| "PruneDataType nested field type mismatch for 'a': read string vs " |
| "data int32"); |
| } |
| |
| // Test: Read only top-level fields, skip struct entirely. |
| TEST_P(NestedColumnPruningInteTest, PruneEntireStructField) { |
| auto struct_type = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| arrow::field("y", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| arrow::field("f2", arrow::float64()), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [100, [1, "aa"], 0.1], |
| [200, [2, "bb"], 0.2], |
| [300, [3, "cc"], 0.3] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Only read f0 and f2, skip f1 entirely. |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f2", arrow::float64()), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 100, 0.1], |
| [0, 200, 0.2], |
| [0, 300, 0.3] |
| ])"); |
| } |
| |
| // Test: Nested struct — prune sub-fields of a struct inside another struct. |
| TEST_P(NestedColumnPruningInteTest, PruneDeepNestedStruct) { |
| // Table schema: f0 (int32), f1 (struct{a: int32, inner: struct{x: int64, y: utf8}}) |
| auto inner_struct = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| arrow::field("y", arrow::utf8()), |
| }); |
| auto outer_struct = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("inner", inner_struct), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", outer_struct), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [10, [100, "aaa"]]], |
| [2, [20, [200, "bbb"]]], |
| [3, [30, [300, "ccc"]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Projected: f0, f1{inner{x}} — skip f1.a and f1.inner.y |
| auto pruned_inner = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| }); |
| auto pruned_outer = arrow::struct_({ |
| arrow::field("inner", pruned_inner), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", pruned_outer), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [[100]]], |
| [0, 2, [[200]]], |
| [0, 3, [[300]]] |
| ])"); |
| } |
| |
| // Test: Nested projected schema with special fields under row tracking. |
| TEST_P(NestedColumnPruningInteTest, PruneNestedStructWithSpecialFields) { |
| // Table schema: f0 (int32), f1 (struct{a: int32, inner: struct{x: int64, y: utf8}}) |
| auto inner_struct = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| arrow::field("y", arrow::utf8()), |
| }); |
| auto outer_struct = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("inner", inner_struct), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", outer_struct), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| {Options::ROW_TRACKING_ENABLED, "true"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [10, [100, "aaa"]]], |
| [2, [20, [200, "bbb"]]], |
| [3, [30, [300, "ccc"]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, |
| helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); |
| |
| // Projected: f0, f1{inner{x}}, _SEQUENCE_NUMBER, _ROW_ID |
| auto pruned_inner = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| }); |
| auto pruned_outer = arrow::struct_({ |
| arrow::field("inner", pruned_inner), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", pruned_outer), |
| arrow::field("_SEQUENCE_NUMBER", arrow::int64()), |
| arrow::field("_ROW_ID", arrow::int64()), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); |
| ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); |
| |
| ASSERT_EQ(read_result->num_chunks(), 1); |
| auto result_array = std::dynamic_pointer_cast<arrow::StructArray>(read_result->chunk(0)); |
| ASSERT_TRUE(result_array); |
| |
| ASSERT_TRUE(result_array->GetFieldByName("_SEQUENCE_NUMBER")); |
| ASSERT_TRUE(result_array->GetFieldByName("_ROW_ID")); |
| auto nested_col = result_array->GetFieldByName("f1"); |
| ASSERT_TRUE(nested_col); |
| |
| auto expected_nested_type = arrow::struct_({ |
| arrow::field("inner", arrow::struct_({arrow::field("x", arrow::int64())})), |
| }); |
| ASSERT_TRUE(nested_col->type()->Equals(expected_nested_type)); |
| |
| auto expected_nested_array = |
| arrow::ipc::internal::json::ArrayFromJSON(expected_nested_type, R"([ |
| [[100]], |
| [[200]], |
| [[300]] |
| ])") |
| .ValueOrDie(); |
| ASSERT_TRUE(nested_col->Equals(expected_nested_array)); |
| } |
| |
| // Test: Table has MAP<STRING, INT32> field, read with selected keys filter. |
| TEST_P(NestedColumnPruningInteTest, MapSelectedKeys) { |
| // Table schema: f0 (int32), f1 (map<string, int32>) |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| // Write data: each row has a map with keys "a", "b", "c" |
| std::string data = R"([ |
| [1, [["a", 10], ["b", 20], ["c", 30]]], |
| [2, [["a", 100], ["c", 300]]], |
| [3, [["b", 200], ["c", 400], ["d", 500]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Build projected schema: read f0 and f1 with selected keys "a,c" |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c"}); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type)->WithMetadata(selected_keys_metadata), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| // Expected: only keys "a" and "c" remain in each map |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [["a", 10], ["c", 30]]], |
| [0, 2, [["a", 100], ["c", 300]]], |
| [0, 3, [["c", 400]]] |
| ])"); |
| } |
| |
| // Test: Selected-keys metadata on MAP nested inside STRUCT should be applied. |
| TEST_P(NestedColumnPruningInteTest, NestedMapSelectedKeysInStruct) { |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| auto struct_type = arrow::struct_({ |
| arrow::field("m", map_type), |
| arrow::field("tag", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [[ ["a", 10], ["b", 20], ["c", 30] ], "r1"]], |
| [2, [[ ["a", 100], ["c", 300] ], "r2"]], |
| [3, [[ ["b", 200], ["c", 400], ["d", 500] ], "r3"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c"}); |
| auto projected_struct_type = arrow::struct_({ |
| arrow::field("m", map_type)->WithMetadata(selected_keys_metadata), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", projected_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [[ ["a", 10], ["c", 30] ]]], |
| [0, 2, [[ ["a", 100], ["c", 300] ]]], |
| [0, 3, [[ ["c", 400] ]]] |
| ])"); |
| } |
| |
| // Test: Partial STRUCT sub-field recall where one recalled child is MAP with selected keys. |
| TEST_P(NestedColumnPruningInteTest, PruneStructSubFieldsWithNestedMapSelectedKeys) { |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| auto struct_type = arrow::struct_({ |
| arrow::field("m", map_type), |
| arrow::field("keep", arrow::int64()), |
| arrow::field("drop", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [[ ["a", 10], ["b", 20], ["c", 30] ], 1001, "x1"]], |
| [2, [[ ["a", 100], ["c", 300] ], 1002, "x2"]], |
| [3, [[ ["b", 200], ["c", 400], ["d", 500] ], 1003, "x3"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c"}); |
| auto projected_struct_type = arrow::struct_({ |
| arrow::field("m", map_type)->WithMetadata(selected_keys_metadata), |
| arrow::field("keep", arrow::int64()), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", projected_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [[ ["a", 10], ["c", 30] ], 1001]], |
| [0, 2, [[ ["a", 100], ["c", 300] ], 1002]], |
| [0, 3, [[ ["c", 400] ], 1003]] |
| ])"); |
| } |
| |
| // Test: Null semantics should be preserved when pruning STRUCT sub-fields and |
| // applying selected-keys filtering on nested MAP. |
| TEST_P(NestedColumnPruningInteTest, PruneStructSubFieldsWithNestedMapSelectedKeysAndNulls) { |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| auto struct_type = arrow::struct_({ |
| arrow::field("m", map_type), |
| arrow::field("keep", arrow::int64()), |
| arrow::field("drop", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [[ ["a", 10], ["b", 20], ["c", 30] ], 1001, "x1"]], |
| [2, null], |
| [3, [null, 1003, "x3"]], |
| [4, [[ ["b", 200], ["c", 400], ["d", 500] ], null, "x4"]], |
| [5, [[ ["a", 500], ["c", null] ], 1005, "x5"]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c"}); |
| auto projected_struct_type = arrow::struct_({ |
| arrow::field("m", map_type)->WithMetadata(selected_keys_metadata), |
| arrow::field("keep", arrow::int64()), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", projected_struct_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [[ ["a", 10], ["c", 30] ], 1001]], |
| [0, 2, null], |
| [0, 3, [null, 1003]], |
| [0, 4, [[ ["c", 400] ], null]], |
| [0, 5, [[ ["a", 500], ["c", null] ], 1005]] |
| ])"); |
| } |
| |
| // Test: MAP_SELECTED_KEYS metadata value is empty string, select empty-string map key. |
| TEST_P(NestedColumnPruningInteTest, MapSelectedKeysEmptyStringKey) { |
| // Table schema: f0 (int32), f1 (map<string, int32>) |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| // Write data: each row has a map that may contain empty-string key. |
| std::string data = R"([ |
| [1, [["", 9], ["a", 10], ["c", 30]]], |
| [2, [["a", 100], ["", 99], ["c", 300]]], |
| [3, [["b", 200], ["c", 400], ["d", 500]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Build projected schema: read f0 and f1 with selected keys metadata set to empty string. |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {""}); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type)->WithMetadata(selected_keys_metadata), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| // Expected: only empty-string key remains. |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [["", 9]]], |
| [0, 2, [["", 99]]], |
| [0, 3, []] |
| ])"); |
| } |
| |
| // Test: MAP_SELECTED_KEYS output map entry order should follow selected key order. |
| TEST_P(NestedColumnPruningInteTest, MapSelectedKeysPreserveOrder) { |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| // Write data with map key order different from selected key order. |
| std::string data = R"([ |
| [1, [["a", 10], ["b", 20], ["c", 30]]], |
| [2, [["a", 100], ["c", 300]]], |
| [3, [["c", 400], ["a", 500], ["d", 600]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Query key order is c,a and output should follow this order. |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"c,a"}); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type)->WithMetadata(selected_keys_metadata), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [["c", 30], ["a", 10]]], |
| [0, 2, [["c", 300], ["a", 100]]], |
| [0, 3, [["c", 400], ["a", 500]]] |
| ])"); |
| } |
| |
| TEST_P(NestedColumnPruningInteTest, NestedStructMapSelectedKeysWithPredicate) { |
| if (file_format_ == "avro") { |
| return; |
| } |
| |
| auto map_type = arrow::map(arrow::utf8(), arrow::int32()); |
| auto info_type = arrow::struct_({ |
| arrow::field("score", arrow::int64()), |
| arrow::field("label", arrow::utf8()), |
| arrow::field("drop", arrow::utf8()), |
| }); |
| auto payload_type = arrow::struct_({ |
| arrow::field("attrs", map_type), |
| arrow::field("info", info_type), |
| arrow::field("note", arrow::utf8()), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("id", arrow::int32()), |
| arrow::field("payload", payload_type), |
| arrow::field("category", arrow::utf8()), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1048576"}, |
| {Options::BUCKET, "-1"}, |
| {Options::WRITE_BATCH_SIZE, "1"}, |
| {"parquet.page.size", "1"}, |
| {"parquet.enable-dictionary", "false"}, |
| {"parquet.write.enable-page-index", "true"}, |
| {"parquet.write.max-row-group-length", "1"}, |
| {"parquet.read.enable-page-index-filter", "true"}, |
| {"orc.stripe.size", "1"}, |
| {"orc.row.index.stride", "1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| WriteContextBuilder write_context_builder(table_path_, "commit_user_1"); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<WriteContext> write_context, |
| write_context_builder.SetOptions(options).Finish()); |
| ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); |
| |
| auto write_one_row = [&](const std::string& data) -> Status { |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| return file_store_write->Write(std::move(batch)); |
| }; |
| |
| ASSERT_OK(write_one_row( |
| R"([[1, [[["a", 10], ["b", 20], ["c", 30]], [1001, "low", "x"], "n1"], "hot"]])")); |
| ASSERT_OK(write_one_row( |
| R"([[12, [[["a", 100], ["c", 300], ["d", 400]], [1002, "mid", "y"], "n2"], "warm"]])")); |
| ASSERT_OK(write_one_row( |
| R"([[21, [[["b", 200], ["c", 500], ["a", 600]], [1003, "high", "z"], "n3"], "cold"]])")); |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| file_store_write->PrepareCommit(/*wait_compaction=*/false, |
| /*commit_identifier=*/0)); |
| ASSERT_OK(file_store_write->Close()); |
| |
| CommitContextBuilder commit_context_builder(table_path_, "commit_user_1"); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommitContext> commit_context, |
| commit_context_builder.SetOptions(options).Finish()); |
| ASSERT_OK_AND_ASSIGN(auto commit, FileStoreCommit::Create(std::move(commit_context))); |
| ASSERT_OK(commit->Commit(commit_msgs, /*commit_identifier=*/0)); |
| |
| // Read selected MAP keys together with nested STRUCT sub-fields. |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"c,a"}); |
| auto selected_payload_type = arrow::struct_({ |
| arrow::field("attrs", map_type)->WithMetadata(selected_keys_metadata), |
| arrow::field("info", arrow::struct_({arrow::field("score", arrow::int64())})), |
| }); |
| auto selected_schema = arrow::schema({ |
| arrow::field("id", arrow::int32()), |
| arrow::field("payload", selected_payload_type), |
| }); |
| |
| ScanReadAndCheck(table_path_, selected_schema, R"([ |
| [0, 1, [[["c", 30], ["a", 10]], [1001]]], |
| [0, 12, [[["c", 300], ["a", 100]], [1002]]], |
| [0, 21, [[["c", 500], ["a", 600]], [1003]]] |
| ])"); |
| |
| // Read only part of top-level columns and part of nested STRUCT fields. |
| auto partial_payload_type = arrow::struct_({ |
| arrow::field("info", arrow::struct_({arrow::field("label", arrow::utf8())})), |
| }); |
| auto partial_schema = arrow::schema({ |
| arrow::field("payload", partial_payload_type), |
| arrow::field("category", arrow::utf8()), |
| }); |
| |
| ScanReadAndCheck(table_path_, partial_schema, R"([ |
| [0, [["low"]], "hot"], |
| [0, [["mid"]], "warm"], |
| [0, [["high"]], "cold"] |
| ])"); |
| |
| // Read selected nested fields with predicate pushdown. |
| auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", |
| FieldType::INT, Literal(10)); |
| |
| ScanReadAndCheck(table_path_, selected_schema, R"([ |
| [0, 12, [[["c", 300], ["a", 100]], [1002]]], |
| [0, 21, [[["c", 500], ["a", 600]], [1003]]] |
| ])", |
| predicate); |
| } |
| |
| // Test: ORC dictionary-encoded map key/value should work with MAP_SELECTED_KEYS. |
| TEST_P(NestedColumnPruningInteTest, MapSelectedKeysWithOrcDictionaryEncodedMap) { |
| if (file_format_ != "orc") { |
| GTEST_SKIP() << "ORC-only dictionary encoding case"; |
| } |
| |
| auto map_type = arrow::map(arrow::utf8(), arrow::utf8()); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| {"orc.read.enable-lazy-decoding", "true"}, |
| {"orc.dictionary-key-size-threshold", "1.0"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| // Low-cardinality map keys/values increase dictionary-encoding probability for ORC. |
| std::string data = R"([ |
| [1, [["a", "v1"], ["b", "v2"], ["c", "v3"]]], |
| [2, [["a", "v1"], ["c", "v3"], ["d", "v4"]]], |
| [3, [["a", "v1"], ["b", "v2"], ["e", "v5"]]], |
| [4, [["a", "v1"], ["c", "v3"], ["e", "v5"]]], |
| [5, [["a", "v1"], ["b", "v2"], ["c", "v3"]]], |
| [6, [["a", "v1"], ["c", "v3"], ["d", "v4"]]], |
| [7, [["a", "v1"], ["b", "v2"], ["e", "v5"]]], |
| [8, [["a", "v1"], ["c", "v3"], ["e", "v5"]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, helper->NewScan(StartupMode::LatestFull(), |
| /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| auto selected_keys_metadata = |
| arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c"}); |
| auto projected_schema = arrow::schema({ |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", map_type)->WithMetadata(selected_keys_metadata), |
| }); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); |
| ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto decoded_result, |
| DictArrayConverter::ConvertDictArray(read_result->chunk(0), arrow::default_memory_pool())); |
| auto actual_chunked = std::make_shared<arrow::ChunkedArray>(decoded_result); |
| |
| auto expected_type = arrow::struct_({ |
| arrow::field("_VALUE_KIND", arrow::int8()), |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::map(arrow::utf8(), arrow::utf8())), |
| }); |
| auto expected_array = arrow::ipc::internal::json::ArrayFromJSON(expected_type, R"([ |
| [0, 1, [["a", "v1"], ["c", "v3"]]], |
| [0, 2, [["a", "v1"], ["c", "v3"]]], |
| [0, 3, [["a", "v1"]]], |
| [0, 4, [["a", "v1"], ["c", "v3"]]], |
| [0, 5, [["a", "v1"], ["c", "v3"]]], |
| [0, 6, [["a", "v1"], ["c", "v3"]]], |
| [0, 7, [["a", "v1"]]], |
| [0, 8, [["a", "v1"], ["c", "v3"]]] |
| ])") |
| .ValueOrDie(); |
| auto expected_chunked = std::make_shared<arrow::ChunkedArray>(expected_array); |
| |
| AssertChunkedArrayEquals(expected_chunked, actual_chunked); |
| } |
| |
| // Test: Deeper nested struct — prune sub-fields of a struct inside a struct inside another |
| // struct. |
| TEST_P(NestedColumnPruningInteTest, PruneDeeperNestedStruct) { |
| // Table schema: f0 (int32), f1 (struct{a: int32, inner1: struct{x: int64, inner2: struct{p: |
| // utf8, q: float64}}}) |
| auto inner2_struct = arrow::struct_({ |
| arrow::field("p", arrow::utf8()), |
| arrow::field("q", arrow::float64()), |
| }); |
| auto inner1_struct = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| arrow::field("inner2", inner2_struct), |
| }); |
| auto outer_struct = arrow::struct_({ |
| arrow::field("a", arrow::int32()), |
| arrow::field("inner1", inner1_struct), |
| }); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", outer_struct), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [10, [100, ["ppp", 1.1]]]], |
| [2, [20, [200, ["qqq", 2.2]]]], |
| [3, [30, [300, ["rrr", 3.3]]]] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| // Projected: f0, f1{inner1{inner2{p}}} |
| auto pruned_inner2 = arrow::struct_({ |
| arrow::field("p", arrow::utf8()), |
| }); |
| auto pruned_inner1 = arrow::struct_({ |
| arrow::field("inner2", pruned_inner2), |
| }); |
| auto pruned_outer = arrow::struct_({ |
| arrow::field("inner1", pruned_inner1), |
| }); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", pruned_outer), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| ScanReadAndCheck(table_path_, projected_schema, R"([ |
| [0, 1, [[[ "ppp" ]]]], |
| [0, 2, [[[ "qqq" ]]]], |
| [0, 3, [[[ "rrr" ]]]] |
| ])"); |
| } |
| |
| // Test: Nested pruning for LIST<STRUCT<...>> in integration path. |
| TEST_P(NestedColumnPruningInteTest, PruneListStructSubFields) { |
| auto list_elem_struct = arrow::struct_({ |
| arrow::field("x", arrow::int64()), |
| arrow::field("y", arrow::utf8()), |
| arrow::field("z", arrow::float64()), |
| }); |
| auto list_struct_type = arrow::list(arrow::field("item", list_elem_struct)); |
| arrow::FieldVector table_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", list_struct_type), |
| }; |
| auto table_schema = arrow::schema(table_fields); |
| |
| std::map<std::string, std::string> options = { |
| {Options::MANIFEST_FORMAT, "AVRO"}, |
| {Options::FILE_FORMAT, StringUtils::ToUpperCase(file_format_)}, |
| {Options::TARGET_FILE_SIZE, "1024"}, |
| {Options::BUCKET, "-1"}, |
| }; |
| |
| ASSERT_OK_AND_ASSIGN( |
| auto helper, TestHelper::Create(test_dir_, table_schema, /*partition_keys=*/{}, |
| /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); |
| |
| std::string data = R"([ |
| [1, [[100, "a", 1.1], [200, "b", 2.2]]], |
| [2, [[300, "c", 3.3]]], |
| [3, []] |
| ])"; |
| ASSERT_OK_AND_ASSIGN(auto batch, |
| TestHelper::MakeRecordBatch(arrow::struct_(table_fields), data, |
| /*partition_map=*/{}, /*bucket=*/0, {})); |
| int64_t commit_identifier = 0; |
| ASSERT_OK_AND_ASSIGN(auto commit_msgs, |
| helper->WriteAndCommit(std::move(batch), commit_identifier++, |
| /*expected_commit_messages=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto data_splits, helper->NewScan(StartupMode::LatestFull(), |
| /*snapshot_id=*/std::nullopt)); |
| ASSERT_FALSE(data_splits.empty()); |
| |
| auto pruned_list_elem_struct = arrow::struct_({arrow::field("x", arrow::int64())}); |
| auto pruned_list_type = arrow::list(arrow::field("item", pruned_list_elem_struct)); |
| arrow::FieldVector projected_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", pruned_list_type), |
| }; |
| auto projected_schema = arrow::schema(projected_fields); |
| |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); |
| |
| ReadContextBuilder read_context_builder(table_path_); |
| read_context_builder.SetOptions(options).SetReadSchema(std::move(c_schema)); |
| ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); |
| ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); |
| auto create_reader_result = table_read->CreateReader(data_splits); |
| ASSERT_NOK_WITH_MSG(create_reader_result, "partial projection inside list"); |
| } |
| |
| std::vector<std::string> GetTestValuesForNestedColumnPruningInteTest() { |
| std::vector<std::string> values; |
| values.emplace_back("parquet"); |
| #ifdef PAIMON_ENABLE_ORC |
| values.emplace_back("orc"); |
| #endif |
| return values; |
| } |
| |
| INSTANTIATE_TEST_SUITE_P(FileFormats, NestedColumnPruningInteTest, |
| ::testing::ValuesIn(GetTestValuesForNestedColumnPruningInteTest())); |
| |
| } // namespace paimon::test |
