| /* |
| * 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 <map> |
| #include <memory> |
| #include <string> |
| #include <vector> |
| |
| #include "arrow/api.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/io/file.h" |
| #include "gtest/gtest.h" |
| #include "paimon/common/data/shredding/shredding_file_reader.h" |
| #include "paimon/common/data/variant/generic_variant.h" |
| #include "paimon/common/data/variant/variant_defs.h" |
| #include "paimon/common/data/variant/variant_schema.h" |
| #include "paimon/common/data/variant/variant_shredding_batch_converter.h" |
| #include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" |
| #include "paimon/common/data/variant/variant_shredding_utils.h" |
| #include "paimon/common/data/variant/variant_shredding_write_plan.h" |
| #include "paimon/common/data/variant/variant_shredding_write_plan_factory.h" |
| #include "paimon/common/data/variant/variant_type_utils.h" |
| #include "paimon/common/types/data_field.h" |
| #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" |
| #include "paimon/common/utils/arrow/mem_utils.h" |
| #include "paimon/common/utils/checked_cast.h" |
| #include "paimon/common/utils/path_util.h" |
| #include "paimon/core/core_options.h" |
| #include "paimon/data/variant.h" |
| #include "paimon/format/parquet/parquet_field_id_converter.h" |
| #include "paimon/format/parquet/parquet_file_batch_reader.h" |
| #include "paimon/format/parquet/parquet_format_defs.h" |
| #include "paimon/format/parquet/parquet_format_writer.h" |
| #include "paimon/fs/local/local_file_system.h" |
| #include "paimon/memory/memory_pool.h" |
| #include "paimon/testing/utils/read_result_collector.h" |
| #include "paimon/testing/utils/testharness.h" |
| #include "paimon/testing/utils/variant_test_data.h" |
| #include "parquet/arrow/reader.h" |
| #include "parquet/file_reader.h" |
| #include "parquet/metadata.h" |
| #include "parquet/properties.h" |
| #include "parquet/schema.h" |
| |
| namespace paimon::parquet::test { |
| |
| class VariantParquetTest : public ::testing::Test { |
| public: |
| void SetUp() override { |
| dir_ = paimon::test::UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir_); |
| fs_ = std::make_shared<LocalFileSystem>(); |
| pool_ = GetDefaultPool(); |
| arrow_pool_ = GetArrowPool(pool_); |
| file_path_ = PathUtil::JoinPath(dir_->Str(), "variant.parquet"); |
| |
| std::vector<DataField> fields = {DataField(1, arrow::field("id", arrow::int32())), |
| DataField(2, VariantTypeUtils::ToArrowField("v"))}; |
| paimon_schema_ = DataField::ConvertDataFieldsToArrowSchema(fields); |
| |
| // `[id, s: struct<nv: VARIANT, t: STRING>]`: a variant nested inside a ROW column, next to |
| // a plain sibling that must survive the physical substitution and the reassembly. |
| nested_sibling_field_ = arrow::field("t", arrow::utf8()); |
| std::vector<DataField> nested_fields = { |
| DataField(1, arrow::field("id", arrow::int32())), |
| DataField(2, arrow::field("s", arrow::struct_({VariantTypeUtils::ToArrowField("nv"), |
| nested_sibling_field_})))}; |
| nested_schema_ = DataField::ConvertDataFieldsToArrowSchema(nested_fields); |
| nested_variant_field_ = nested_schema_->field(1)->type()->field(0); |
| |
| // `[id, arr: ARRAY<VARIANT>]` and `[id, m: MAP<STRING, VARIANT>]`: variants inside |
| // repeated groups, which are never shredded but can still be read as a projection. |
| list_element_field_ = VariantTypeUtils::ToArrowField("element"); |
| std::vector<DataField> list_fields = { |
| DataField(1, arrow::field("id", arrow::int32())), |
| DataField(2, arrow::field("arr", arrow::list(list_element_field_)))}; |
| list_schema_ = DataField::ConvertDataFieldsToArrowSchema(list_fields); |
| list_element_field_ = list_schema_->field(1)->type()->field(0); |
| |
| map_item_field_ = VariantTypeUtils::ToArrowField("value"); |
| std::vector<DataField> map_fields = { |
| DataField(1, arrow::field("id", arrow::int32())), |
| DataField(2, arrow::field("m", arrow::map(arrow::utf8(), map_item_field_)))}; |
| map_schema_ = DataField::ConvertDataFieldsToArrowSchema(map_fields); |
| map_item_field_ = map_schema_->field(1)->type()->field(0)->type()->field(1); |
| |
| // `[id, arr2: ARRAY<ROW<v: VARIANT, t: STRING>>]`: a variant one struct level below a |
| // repeated group, where the read and file children must line up field by field. |
| std::vector<DataField> list_struct_fields = { |
| DataField(1, arrow::field("id", arrow::int32())), |
| DataField( |
| 2, arrow::field("arr2", |
| arrow::list(arrow::field( |
| "element", arrow::struct_({VariantTypeUtils::ToArrowField("v"), |
| nested_sibling_field_})))))}; |
| list_struct_schema_ = DataField::ConvertDataFieldsToArrowSchema(list_struct_fields); |
| list_struct_variant_field_ = |
| list_struct_schema_->field(1)->type()->field(0)->type()->field(0); |
| } |
| |
| std::shared_ptr<arrow::StructArray> BuildArray(const std::vector<const char*>& jsons) { |
| EXPECT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> batch, |
| paimon::test::VariantTestData::BuildVariantBatch( |
| paimon_schema_->field(0), paimon_schema_->field(1), jsons, pool_)); |
| return batch; |
| } |
| |
| // Writes one batch with the given logical schema through the production parquet write path |
| // (mapping paimon field ids to parquet field ids). |
| void WriteFile(const std::shared_ptr<arrow::Schema>& schema, ArrowArray* c_array) { |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Schema> write_schema, |
| ParquetFieldIdConverter::AddParquetIdsFromPaimonIds(schema)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out, |
| fs_->Create(file_path_, /*overwrite=*/true)); |
| ::parquet::WriterProperties::Builder builder; |
| auto writer_properties = builder.build(); |
| ASSERT_OK_AND_ASSIGN( |
| auto format_writer, |
| ParquetFormatWriter::Create(out, write_schema, writer_properties, |
| DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); |
| ASSERT_OK(format_writer->AddBatch(c_array)); |
| ASSERT_OK(format_writer->Flush()); |
| ASSERT_OK(format_writer->Finish()); |
| ASSERT_OK(out->Flush()); |
| ASSERT_OK(out->Close()); |
| } |
| |
| void WriteFile(const std::shared_ptr<arrow::StructArray>& array) { |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok()); |
| WriteFile(paimon_schema_, arrow_array.get()); |
| } |
| |
| void WriteShreddedFile(const std::vector<const char*>& jsons, |
| const std::shared_ptr<VariantShreddingWritePlan>& plan) { |
| ASSERT_NE(plan, nullptr); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<VariantShreddingBatchConverter> converter, |
| VariantShreddingBatchConverter::Create(plan, pool_)); |
| auto logical = BuildArray(jsons); |
| auto c_logical = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*logical, c_logical.get()).ok()); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<ArrowArray> c_physical, |
| converter->Convert(c_logical.get())); |
| WriteFile(converter->GetPhysicalSchema(), c_physical.get()); |
| } |
| |
| // Writes `jsons` shredded according to the configured ROW-type shredding schema JSON. |
| void WriteShreddedFile(const std::vector<const char*>& jsons, |
| const char* shredding_schema_json) { |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<VariantShreddingWritePlan> plan, |
| VariantShreddingWritePlan::FromConfiguredSchema(paimon_schema_, shredding_schema_json)); |
| WriteShreddedFile(jsons, plan); |
| } |
| |
| // Writes `jsons` using the given inferred shredding type for the top-level Variant column. |
| void WriteShreddedFile(const std::vector<const char*>& jsons, |
| const std::shared_ptr<arrow::DataType>& shredding_type) { |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<VariantShreddingWritePlan> plan, |
| VariantShreddingWritePlan::Create(paimon_schema_, {{"v", shredding_type}})); |
| WriteShreddedFile(jsons, plan); |
| } |
| |
| static std::string NestedSiblingValue(size_t row) { |
| return "t" + std::to_string(row); |
| } |
| |
| std::shared_ptr<arrow::ArrayData> BuildIdColumn(size_t rows) { |
| arrow::Int32Builder id_builder; |
| for (size_t i = 0; i < rows; ++i) { |
| EXPECT_TRUE(id_builder.Append(static_cast<int32_t>(i)).ok()); |
| } |
| std::shared_ptr<arrow::Array> ids; |
| EXPECT_TRUE(id_builder.Finish(&ids).ok()); |
| return ids->data(); |
| } |
| |
| // Builds the offsets buffer of a repeated column, flattening its elements into `flat`. |
| std::shared_ptr<arrow::Buffer> BuildOffsetsBuffer( |
| const std::vector<std::vector<const char*>>& rows, std::vector<const char*>* flat) { |
| arrow::Int32Builder offset_builder; |
| for (const auto& row : rows) { |
| EXPECT_TRUE(offset_builder.Append(static_cast<int32_t>(flat->size())).ok()); |
| flat->insert(flat->end(), row.begin(), row.end()); |
| } |
| EXPECT_TRUE(offset_builder.Append(static_cast<int32_t>(flat->size())).ok()); |
| std::shared_ptr<arrow::Array> offsets; |
| EXPECT_TRUE(offset_builder.Finish(&offsets).ok()); |
| return offsets->data()->buffers[1]; |
| } |
| |
| // Builds a `[id, arr: list<VARIANT>]` batch: row `i` holds the variants of `rows[i]`. |
| // |
| // The repeated columns here are assembled at the `ArrayData` level on purpose: Arrow's |
| // `FromArrays` helpers `checked_cast` their arguments, which is a `dynamic_cast` in a debug |
| // build and fails across the test binary / libpaimon boundary for templated array classes. |
| std::shared_ptr<arrow::StructArray> BuildListArray( |
| const std::vector<std::vector<const char*>>& rows) { |
| std::vector<const char*> flat; |
| auto offsets_buffer = BuildOffsetsBuffer(rows, &flat); |
| // The element variants are built as one flat batch that the offsets slice into rows. |
| auto elements = paimon::test::VariantTestData::BuildVariantBatch( |
| list_schema_->field(0), list_element_field_, flat, pool_); |
| EXPECT_TRUE(elements.ok()) << elements.status().ToString(); |
| auto list_data = arrow::ArrayData::Make( |
| list_schema_->field(1)->type(), static_cast<int64_t>(rows.size()), |
| {nullptr, offsets_buffer}, {elements.value()->field(1)->data()}, /*null_count=*/0); |
| auto batch_data = arrow::ArrayData::Make( |
| arrow::struct_(list_schema_->fields()), static_cast<int64_t>(rows.size()), {nullptr}, |
| {BuildIdColumn(rows.size()), list_data}, /*null_count=*/0); |
| return std::make_shared<arrow::StructArray>(batch_data); |
| } |
| |
| // Writes `rows` into the unshredded `arr: list<VARIANT>` column. |
| void WriteListFile(const std::vector<std::vector<const char*>>& rows) { |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*BuildListArray(rows), arrow_array.get()).ok()); |
| WriteFile(list_schema_, arrow_array.get()); |
| } |
| |
| // Builds a `[id, m: map<STRING, VARIANT>]` batch: row `i` maps `keys[i][k]` to the variant of |
| // `rows[i][k]`. |
| std::shared_ptr<arrow::StructArray> BuildMapArray( |
| const std::vector<std::vector<std::string>>& keys, |
| const std::vector<std::vector<const char*>>& rows) { |
| std::vector<const char*> flat; |
| arrow::StringBuilder key_builder; |
| EXPECT_EQ(keys.size(), rows.size()); |
| for (size_t i = 0; i < keys.size() && i < rows.size(); ++i) { |
| EXPECT_EQ(keys[i].size(), rows[i].size()) << "row " << i; |
| for (const auto& key : keys[i]) { |
| EXPECT_TRUE(key_builder.Append(key).ok()); |
| } |
| } |
| auto offsets_buffer = BuildOffsetsBuffer(rows, &flat); |
| std::shared_ptr<arrow::Array> map_keys; |
| EXPECT_TRUE(key_builder.Finish(&map_keys).ok()); |
| auto elements = paimon::test::VariantTestData::BuildVariantBatch( |
| map_schema_->field(0), map_item_field_, flat, pool_); |
| EXPECT_TRUE(elements.ok()) << elements.status().ToString(); |
| // A map array is a list of `struct<key, value>` entries. |
| auto entries_data = arrow::ArrayData::Make( |
| map_schema_->field(1)->type()->field(0)->type(), static_cast<int64_t>(flat.size()), |
| {nullptr}, {map_keys->data(), elements.value()->field(1)->data()}, /*null_count=*/0); |
| auto map_data = |
| arrow::ArrayData::Make(map_schema_->field(1)->type(), static_cast<int64_t>(rows.size()), |
| {nullptr, offsets_buffer}, {entries_data}, /*null_count=*/0); |
| auto batch_data = arrow::ArrayData::Make( |
| arrow::struct_(map_schema_->fields()), static_cast<int64_t>(rows.size()), {nullptr}, |
| {BuildIdColumn(rows.size()), map_data}, /*null_count=*/0); |
| return std::make_shared<arrow::StructArray>(batch_data); |
| } |
| |
| // Writes `rows` into the unshredded `m: map<STRING, VARIANT>` column. |
| void WriteMapFile(const std::vector<std::vector<std::string>>& keys, |
| const std::vector<std::vector<const char*>>& rows) { |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*BuildMapArray(keys, rows), arrow_array.get()).ok()); |
| WriteFile(map_schema_, arrow_array.get()); |
| } |
| |
| // Builds a `[id, arr2: list<struct<v: VARIANT, t: STRING>>]` batch. The sibling of the |
| // `k`-th element over the whole column is `NestedSiblingValue(k)`. |
| std::shared_ptr<arrow::StructArray> BuildListStructArray( |
| const std::vector<std::vector<const char*>>& rows) { |
| std::vector<const char*> flat; |
| auto offsets_buffer = BuildOffsetsBuffer(rows, &flat); |
| auto elements = paimon::test::VariantTestData::BuildVariantBatch( |
| list_struct_schema_->field(0), list_struct_variant_field_, flat, pool_); |
| EXPECT_TRUE(elements.ok()) << elements.status().ToString(); |
| arrow::StringBuilder sibling_builder; |
| for (size_t i = 0; i < flat.size(); ++i) { |
| EXPECT_TRUE(sibling_builder.Append(NestedSiblingValue(i)).ok()); |
| } |
| std::shared_ptr<arrow::Array> sibling; |
| EXPECT_TRUE(sibling_builder.Finish(&sibling).ok()); |
| auto element_data = arrow::ArrayData::Make( |
| list_struct_schema_->field(1)->type()->field(0)->type(), |
| static_cast<int64_t>(flat.size()), {nullptr}, |
| {elements.value()->field(1)->data(), sibling->data()}, /*null_count=*/0); |
| auto list_data = arrow::ArrayData::Make( |
| list_struct_schema_->field(1)->type(), static_cast<int64_t>(rows.size()), |
| {nullptr, offsets_buffer}, {element_data}, /*null_count=*/0); |
| auto batch_data = arrow::ArrayData::Make( |
| arrow::struct_(list_struct_schema_->fields()), static_cast<int64_t>(rows.size()), |
| {nullptr}, {BuildIdColumn(rows.size()), list_data}, /*null_count=*/0); |
| return std::make_shared<arrow::StructArray>(batch_data); |
| } |
| |
| // Writes `rows` into the unshredded `arr2: list<struct<v: VARIANT, t: STRING>>` column. |
| void WriteListStructFile(const std::vector<std::vector<const char*>>& rows) { |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*BuildListStructArray(rows), arrow_array.get()).ok()); |
| WriteFile(list_struct_schema_, arrow_array.get()); |
| } |
| |
| // Builds a `[id, s: struct<nv, t>]` batch holding the variant encodings of `jsons`. |
| std::shared_ptr<arrow::StructArray> BuildNestedArray(const std::vector<const char*>& jsons) { |
| auto batch = paimon::test::VariantTestData::BuildVariantBatch( |
| nested_schema_->field(0), nested_variant_field_, jsons, pool_); |
| EXPECT_TRUE(batch.ok()) << batch.status().ToString(); |
| arrow::StringBuilder sibling_builder; |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| EXPECT_TRUE(sibling_builder.Append(NestedSiblingValue(i)).ok()); |
| } |
| std::shared_ptr<arrow::Array> sibling; |
| EXPECT_TRUE(sibling_builder.Finish(&sibling).ok()); |
| auto struct_column = arrow::StructArray::Make( |
| {batch.value()->field(1), sibling}, {nested_variant_field_, nested_sibling_field_}); |
| EXPECT_TRUE(struct_column.ok()) << struct_column.status().ToString(); |
| auto nested = |
| arrow::StructArray::Make({batch.value()->field(0), struct_column.ValueOrDie()}, |
| {nested_schema_->field(0), nested_schema_->field(1)}); |
| EXPECT_TRUE(nested.ok()) << nested.status().ToString(); |
| return nested.ValueOrDie(); |
| } |
| |
| // Writes `jsons` into `s.nv` unshredded. |
| void WriteNestedFile(const std::vector<const char*>& jsons) { |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*BuildNestedArray(jsons), arrow_array.get()).ok()); |
| WriteFile(nested_schema_, arrow_array.get()); |
| } |
| |
| // Writes `jsons` into `s.nv` shredded by `shredding_type` (the nested variant is addressed by |
| // its field-index path `{1, 0}`). |
| void WriteShreddedNestedFile(const std::vector<const char*>& jsons, |
| const std::shared_ptr<arrow::DataType>& shredding_type) { |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<VariantShreddingWritePlan> plan, |
| VariantShreddingWritePlan::CreateFromPaths( |
| nested_schema_, {{std::vector<int32_t>{1, 0}, shredding_type}})); |
| ASSERT_NE(plan, nullptr); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<VariantShreddingBatchConverter> converter, |
| VariantShreddingBatchConverter::Create(plan, pool_)); |
| auto c_logical = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*BuildNestedArray(jsons), c_logical.get()).ok()); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<ArrowArray> c_physical, |
| converter->Convert(c_logical.get())); |
| WriteFile(converter->GetPhysicalSchema(), c_physical.get()); |
| } |
| |
| // Builds the read schema projecting `s.nv` as the given variant-access projection, keeping |
| // the plain sibling `s.t`. |
| std::shared_ptr<arrow::Schema> BuildNestedAccessReadSchema( |
| const std::vector<std::pair<std::shared_ptr<arrow::DataType>, std::string>>& accesses) { |
| auto access_field = BuildAccessField(accesses, "nv"); |
| auto read_struct = nested_schema_->field(1)->WithType( |
| arrow::struct_({access_field, nested_sibling_field_})); |
| return arrow::schema({nested_schema_->field(0), read_struct}); |
| } |
| |
| // Asserts that the plain sibling column of `s` round-tripped unchanged. |
| void ExpectNestedSibling(const std::shared_ptr<arrow::StructArray>& s_column) { |
| const auto& sibling = static_cast<const arrow::StringArray&>(*s_column->field(1)); |
| for (int64_t i = 0; i < s_column->length(); ++i) { |
| EXPECT_EQ(sibling.GetString(i), NestedSiblingValue(static_cast<size_t>(i))); |
| } |
| } |
| |
| // Opens the written file and returns the reader plus its imported file schema. |
| void OpenFile(std::unique_ptr<FileBatchReader>* file_reader, |
| std::shared_ptr<arrow::Schema>* file_schema) { |
| ASSERT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_path_)); |
| auto length = fs_->GetFileStatus(file_path_).value().GetLen(); |
| auto in_stream = |
| std::make_unique<ArrowInputStreamAdapter>(std::move(input_stream), length, arrow_pool_); |
| std::map<std::string, std::string> options = {}; |
| ASSERT_OK_AND_ASSIGN(auto parquet_reader, ParquetFileBatchReader::Create( |
| std::move(in_stream), options, |
| /*batch_size=*/1024, |
| /*file_metadata=*/nullptr, |
| /*storage_read_bytes=*/nullptr, arrow_pool_)); |
| *file_reader = std::move(parquet_reader); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> c_file_schema, |
| (*file_reader)->GetFileSchema()); |
| auto imported = arrow::ImportSchema(c_file_schema.get()); |
| ASSERT_TRUE(imported.ok()) << imported.status().ToString(); |
| *file_schema = imported.ValueOrDie(); |
| } |
| |
| // Builds a variant-access projection field for a variant column via the public builder. |
| std::shared_ptr<arrow::Field> BuildAccessField( |
| const std::vector<std::pair<std::shared_ptr<arrow::DataType>, std::string>>& accesses, |
| const std::string& field_name = "v") { |
| VariantAccessBuilder builder; |
| for (const auto& [type, path] : accesses) { |
| auto c_target = std::make_unique<ArrowSchema>(); |
| EXPECT_TRUE(arrow::ExportField(arrow::Field("t", type), c_target.get()).ok()); |
| EXPECT_OK(builder.AddField(c_target.get(), path, /*fail_on_error=*/false)); |
| } |
| auto c_field = builder.Build(field_name); |
| EXPECT_TRUE(c_field.ok()) << c_field.status().ToString(); |
| auto imported = arrow::ImportField(c_field.value().get()); |
| EXPECT_TRUE(imported.ok()) << imported.status().ToString(); |
| return imported.ValueOrDie(); |
| } |
| |
| // Reads the whole file through the shredding reader with the given read schema and returns |
| // the second (variant) column. |
| void ReadColumn(const std::shared_ptr<arrow::Schema>& read_schema, |
| std::shared_ptr<arrow::Array>* column) { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| ASSERT_EQ(plans.size(), 1); |
| auto shredding_reader = |
| std::make_unique<ShreddingFileReader>(std::move(file_reader), std::move(plans), pool_); |
| auto c_read_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_read_schema.get()).ok()); |
| ASSERT_OK(shredding_reader->SetReadSchema(c_read_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, shredding_reader->NextBatchWithBitmap()); |
| ASSERT_FALSE(BatchReader::IsEofBatch(batch_with_bitmap)); |
| auto& [read_batch, bitmap] = batch_with_bitmap; |
| auto imported = arrow::ImportArray(read_batch.first.get(), read_batch.second.get()); |
| ASSERT_TRUE(imported.ok()) << imported.status().ToString(); |
| auto result_struct = checked_pointer_cast<arrow::StructArray>(imported.ValueOrDie()); |
| *column = result_struct->field(1); |
| shredding_reader->Close(); |
| // The assembled arrays borrow the reader's memory pool; keep the reader alive until the |
| // fixture is torn down (fixture members outlive test-body locals). |
| live_readers_.push_back(std::move(shredding_reader)); |
| } |
| |
| // `ReadColumn` for the cases whose second column is a struct. |
| void ReadVariantColumn(const std::shared_ptr<arrow::Schema>& read_schema, |
| std::shared_ptr<arrow::StructArray>* v_column) { |
| std::shared_ptr<arrow::Array> column; |
| ReadColumn(read_schema, &column); |
| ASSERT_EQ(column->type_id(), arrow::Type::STRUCT); |
| *v_column = checked_pointer_cast<arrow::StructArray>(column); |
| } |
| |
| protected: |
| std::unique_ptr<paimon::test::UniqueTestDirectory> dir_; |
| std::shared_ptr<FileSystem> fs_; |
| std::shared_ptr<MemoryPool> pool_; |
| std::shared_ptr<arrow::MemoryPool> arrow_pool_; |
| std::string file_path_; |
| std::shared_ptr<arrow::Schema> paimon_schema_; |
| std::shared_ptr<arrow::Schema> nested_schema_; |
| std::shared_ptr<arrow::Field> nested_variant_field_; |
| std::shared_ptr<arrow::Field> nested_sibling_field_; |
| std::shared_ptr<arrow::Schema> list_schema_; |
| std::shared_ptr<arrow::Field> list_element_field_; |
| std::shared_ptr<arrow::Schema> map_schema_; |
| std::shared_ptr<arrow::Field> map_item_field_; |
| std::shared_ptr<arrow::Schema> list_struct_schema_; |
| std::shared_ptr<arrow::Field> list_struct_variant_field_; |
| std::vector<std::unique_ptr<FileBatchReader>> live_readers_; |
| }; |
| |
| namespace { |
| |
| constexpr const char* kAgeCityShreddingSchema = R"({ |
| "type": "ROW", |
| "fields": [ { |
| "id": 0, |
| "name": "v", |
| "type": { |
| "type": "ROW", |
| "fields": [ |
| {"id": 1, "name": "age", "type": "INT"}, |
| {"id": 2, "name": "city", "type": "STRING"} |
| ] |
| } |
| } ] |
| })"; |
| |
| } // namespace |
| |
| TEST_F(VariantParquetTest, PhysicalLayoutMatchesJava) { |
| auto array = BuildArray({R"({"a": 1, "b": "hello"})", nullptr, "[1,2,3]"}); |
| WriteFile(array); |
| |
| // The on-disk layout must match the Java ParquetSchemaConverter: an (unannotated) group |
| // with two REQUIRED BINARY fields `value` (id 0) and `metadata` (id 1). The raw parquet |
| // reader is required because these parquet-level properties (repetition, physical types, |
| // field ids, the absence of a logical-type annotation) are not visible in the Arrow schema |
| // surfaced by the paimon reader. |
| auto file = arrow::io::ReadableFile::Open(file_path_, arrow_pool_.get()); |
| ASSERT_TRUE(file.ok()); |
| std::unique_ptr<::parquet::arrow::FileReader> reader; |
| auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); |
| ASSERT_TRUE(status.ok()) << status.ToString(); |
| const ::parquet::SchemaDescriptor* schema = reader->parquet_reader()->metadata()->schema(); |
| ASSERT_EQ(schema->num_columns(), 3); |
| const auto* root = schema->group_node(); |
| ASSERT_EQ(root->field_count(), 2); |
| const auto& variant_group_node = root->field(1); |
| ASSERT_TRUE(variant_group_node->is_group()); |
| ASSERT_EQ(variant_group_node->name(), "v"); |
| ASSERT_EQ(variant_group_node->field_id(), 2); |
| ASSERT_EQ(variant_group_node->logical_type()->type(), ::parquet::LogicalType::Type::NONE); |
| const auto* variant_group = |
| checked_cast<const ::parquet::schema::GroupNode*>(variant_group_node.get()); |
| ASSERT_EQ(variant_group->field_count(), 2); |
| const auto& value_node = variant_group->field(0); |
| ASSERT_EQ(value_node->name(), "value"); |
| ASSERT_TRUE(value_node->is_primitive()); |
| ASSERT_TRUE(value_node->is_required()); |
| ASSERT_EQ(value_node->field_id(), 0); |
| ASSERT_EQ( |
| checked_cast<const ::parquet::schema::PrimitiveNode*>(value_node.get())->physical_type(), |
| ::parquet::Type::BYTE_ARRAY); |
| const auto& metadata_node = variant_group->field(1); |
| ASSERT_EQ(metadata_node->name(), "metadata"); |
| ASSERT_TRUE(metadata_node->is_primitive()); |
| ASSERT_TRUE(metadata_node->is_required()); |
| ASSERT_EQ(metadata_node->field_id(), 1); |
| } |
| |
| TEST_F(VariantParquetTest, WriteAndReadRoundTrip) { |
| std::vector<const char*> jsons = { |
| R"({"a": 1, "b": "hello"})", |
| nullptr, |
| "[1,2,3]", |
| "{\"nested\": {\"x\": [true, null, 1.5]}, \"s\": \"ä¸æ–‡\"}", |
| "12345678901234", |
| "100.99", |
| }; |
| auto array = BuildArray(jsons); |
| WriteFile(array); |
| |
| { |
| // Sanity-check the raw file through the plain parquet-arrow reader: the struct child |
| // arrays must align with the logical rows. |
| auto file = arrow::io::ReadableFile::Open(file_path_, arrow_pool_.get()); |
| ASSERT_TRUE(file.ok()); |
| std::unique_ptr<::parquet::arrow::FileReader> raw_reader; |
| ASSERT_TRUE( |
| ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &raw_reader).ok()); |
| std::shared_ptr<arrow::Table> table; |
| ASSERT_TRUE(raw_reader->ReadTable(&table).ok()); |
| auto raw_variant = checked_pointer_cast<arrow::StructArray>(table->column(1)->chunk(0)); |
| auto raw_value = checked_pointer_cast<arrow::BinaryArray>(raw_variant->field(0)); |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| SCOPED_TRACE("raw row " + std::to_string(i)); |
| if (jsons[i] != nullptr) { |
| ASSERT_FALSE(raw_variant->IsNull(i)); |
| ASSERT_GT(raw_value->GetView(i).size(), 0); |
| } else { |
| ASSERT_TRUE(raw_variant->IsNull(i)); |
| } |
| } |
| } |
| |
| ASSERT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_path_)); |
| auto length = fs_->GetFileStatus(file_path_).value().GetLen(); |
| auto in_stream = |
| std::make_unique<ArrowInputStreamAdapter>(std::move(input_stream), length, arrow_pool_); |
| std::map<std::string, std::string> options = {}; |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, ParquetFileBatchReader::Create( |
| std::move(in_stream), options, |
| /*batch_size=*/1024, |
| /*file_metadata=*/nullptr, |
| /*storage_read_bytes=*/nullptr, arrow_pool_)); |
| auto c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*paimon_schema_, c_schema.get()).ok()); |
| ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(auto result_chunked, |
| paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); |
| batch_reader->Close(); |
| ASSERT_EQ(result_chunked->length(), static_cast<int64_t>(jsons.size())); |
| ASSERT_EQ(result_chunked->num_chunks(), 1); |
| auto result_struct = checked_pointer_cast<arrow::StructArray>(result_chunked->chunk(0)); |
| |
| auto variant_column = checked_pointer_cast<arrow::StructArray>(result_struct->field(1)); |
| ASSERT_EQ(variant_column->length(), static_cast<int64_t>(jsons.size())); |
| ASSERT_EQ(variant_column->field(0)->length(), variant_column->length()); |
| auto value_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(0)); |
| auto metadata_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(1)); |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| SCOPED_TRACE("row " + std::to_string(i)); |
| if (jsons[i] == nullptr) { |
| ASSERT_TRUE(variant_column->IsNull(i)); |
| continue; |
| } |
| ASSERT_FALSE(variant_column->IsNull(i)); |
| auto value_view = value_column->GetView(i); |
| auto metadata_view = metadata_column->GetView(i); |
| SCOPED_TRACE("value size " + std::to_string(value_view.size()) + ", metadata size " + |
| std::to_string(metadata_view.size())); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> variant, |
| GenericVariant::Create(value_view, metadata_view, pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> expected, |
| GenericVariant::FromJson(jsons[i], pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); |
| ASSERT_EQ(actual_json, expected_json); |
| } |
| } |
| |
| TEST_F(VariantParquetTest, ShreddedWriteAndReadRoundTrip) { |
| std::vector<const char*> jsons = { |
| R"({"age": 35, "city": "Hangzhou"})", |
| nullptr, |
| R"({"age": "not a number", "extra": [1, 2]})", |
| "[\"top level array\"]", |
| }; |
| WriteShreddedFile(jsons, kAgeCityShreddingSchema); |
| |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| auto file_variant_field = file_schema->GetFieldByName("v"); |
| ASSERT_NE(file_variant_field, nullptr); |
| ASSERT_TRUE(VariantShreddingUtils::IsShreddedFileType(file_variant_field->type())) |
| << file_variant_field->type()->ToString(); |
| file_reader->Close(); |
| } |
| |
| // Reading the column as a plain VARIANT reassembles every physical shape back to the |
| // original logical value. |
| std::shared_ptr<arrow::StructArray> variant_column; |
| ReadVariantColumn(paimon_schema_, &variant_column); |
| ASSERT_EQ(variant_column->length(), static_cast<int64_t>(jsons.size())); |
| auto value_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(0)); |
| auto metadata_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(1)); |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| SCOPED_TRACE("row " + std::to_string(i)); |
| if (jsons[i] == nullptr) { |
| ASSERT_TRUE(variant_column->IsNull(i)); |
| continue; |
| } |
| ASSERT_FALSE(variant_column->IsNull(i)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> variant, |
| GenericVariant::Create(value_column->GetView(i), metadata_column->GetView(i), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> expected, |
| GenericVariant::FromJson(jsons[i], pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); |
| ASSERT_EQ(actual_json, expected_json); |
| } |
| } |
| |
| TEST_F(VariantParquetTest, UntypedPhysicalVariantWriteAndReadRoundTrip) { |
| std::vector<const char*> jsons = { |
| R"({"a": 1, "b": "hello"})", |
| nullptr, |
| "[1,2,3]", |
| }; |
| WriteShreddedFile(jsons, arrow::null()); |
| |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| auto file_variant_field = file_schema->GetFieldByName("v"); |
| ASSERT_NE(file_variant_field, nullptr); |
| const auto& physical_type = |
| static_cast<const arrow::StructType&>(*file_variant_field->type()); |
| ASSERT_EQ(physical_type.num_fields(), 2); |
| ASSERT_EQ(physical_type.field(0)->name(), VariantDefs::kMetadataFieldName); |
| ASSERT_EQ(physical_type.field(1)->name(), VariantDefs::kValueFieldName); |
| ASSERT_FALSE(VariantShreddingUtils::IsShreddedFileType(file_variant_field->type())); |
| ASSERT_TRUE( |
| VariantShreddingUtils::IsUntypedPhysicalVariantType(file_variant_field->type())); |
| file_reader->Close(); |
| } |
| |
| std::shared_ptr<arrow::StructArray> variant_column; |
| ReadVariantColumn(paimon_schema_, &variant_column); |
| ASSERT_EQ(variant_column->length(), static_cast<int64_t>(jsons.size())); |
| auto value_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(0)); |
| auto metadata_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(1)); |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| SCOPED_TRACE("row " + std::to_string(i)); |
| if (jsons[i] == nullptr) { |
| ASSERT_TRUE(variant_column->IsNull(i)); |
| continue; |
| } |
| ASSERT_FALSE(variant_column->IsNull(i)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> variant, |
| GenericVariant::Create(value_column->GetView(i), metadata_column->GetView(i), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> expected, |
| GenericVariant::FromJson(jsons[i], pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); |
| ASSERT_EQ(actual_json, expected_json); |
| } |
| } |
| |
| TEST_F(VariantParquetTest, AdaptiveInferenceUntypedPhysicalWriteAndReadRoundTrip) { |
| std::vector<const char*> jsons = { |
| R"({"a": 1})", |
| "[1,2,3]", |
| nullptr, |
| }; |
| auto logical = BuildArray(jsons); |
| ASSERT_OK_AND_ASSIGN(CoreOptions options, |
| CoreOptions::FromMap({ |
| {Options::MANIFEST_FORMAT, "parquet"}, |
| {Options::VARIANT_INFER_SHREDDING_SCHEMA, "true"}, |
| {Options::VARIANT_SHREDDING_INFERENCE_MODE, "adaptive"}, |
| })); |
| auto factory = VariantShreddingWritePlanFactory::Create(options, paimon_schema_, pool_); |
| std::vector<std::shared_ptr<arrow::Array>> samples = {logical}; |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<ShreddingBatchConverter> converter, |
| factory->CreateConverter("parquet", samples)); |
| auto physical_variant = converter->GetPhysicalSchema()->GetFieldByName("v"); |
| ASSERT_NE(physical_variant, nullptr); |
| ASSERT_FALSE(VariantShreddingUtils::IsShreddedFileType(physical_variant->type())); |
| ASSERT_TRUE(VariantShreddingUtils::IsUntypedPhysicalVariantType(physical_variant->type())); |
| |
| auto c_logical = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*logical, c_logical.get()).ok()); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<ArrowArray> c_physical, |
| converter->Convert(c_logical.get())); |
| WriteFile(converter->GetPhysicalSchema(), c_physical.get()); |
| ASSERT_OK(factory->OnFileCompleted(converter)); |
| |
| std::shared_ptr<arrow::StructArray> variant_column; |
| ReadVariantColumn(paimon_schema_, &variant_column); |
| ASSERT_EQ(variant_column->length(), static_cast<int64_t>(jsons.size())); |
| auto value_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(0)); |
| auto metadata_column = checked_pointer_cast<arrow::BinaryArray>(variant_column->field(1)); |
| for (size_t i = 0; i < jsons.size(); ++i) { |
| SCOPED_TRACE("row " + std::to_string(i)); |
| if (jsons[i] == nullptr) { |
| ASSERT_TRUE(variant_column->IsNull(i)); |
| continue; |
| } |
| ASSERT_FALSE(variant_column->IsNull(i)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> variant, |
| GenericVariant::Create(value_column->GetView(i), metadata_column->GetView(i), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> expected, |
| GenericVariant::FromJson(jsons[i], pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); |
| ASSERT_EQ(actual_json, expected_json); |
| } |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadMixedTypedAndBinary) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", |
| R"({"age": 25, "other": "Hello"})", nullptr}; |
| WriteShreddedFile(jsons, kAgeCityShreddingSchema); |
| |
| auto access_field = BuildAccessField( |
| {{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}, {arrow::utf8(), "$.missing"}}); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); |
| |
| // The plan prunes `typed_value` to the requested keys and keeps `value` because `$.other` |
| // and `$.missing` are not shredded. |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| ASSERT_EQ(plans.size(), 1); |
| const auto& physical_type = |
| static_cast<const arrow::StructType&>(*plans.at("v")->PhysicalField()->type()); |
| ASSERT_NE(physical_type.GetFieldByName(VariantDefs::kMetadataFieldName), nullptr); |
| ASSERT_NE(physical_type.GetFieldByName(VariantDefs::kValueFieldName), nullptr); |
| auto typed_value = physical_type.GetFieldByName(VariantDefs::kTypedValueFieldName); |
| ASSERT_NE(typed_value, nullptr); |
| const auto& typed_struct = static_cast<const arrow::StructType&>(*typed_value->type()); |
| ASSERT_EQ(typed_struct.num_fields(), 1); |
| ASSERT_NE(typed_struct.GetFieldByName("age"), nullptr); |
| file_reader->Close(); |
| } |
| |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| ASSERT_EQ(v_column->length(), 3); |
| const auto& age = static_cast<const arrow::Int64Array&>(*v_column->field(0)); |
| const auto& other = static_cast<const arrow::StringArray&>(*v_column->field(1)); |
| const auto& missing = static_cast<const arrow::StringArray&>(*v_column->field(2)); |
| ASSERT_EQ(age.Value(0), 35); |
| ASSERT_EQ(age.Value(1), 25); |
| ASSERT_TRUE(v_column->IsNull(2)); |
| ASSERT_TRUE(other.IsNull(0)); |
| ASSERT_EQ(other.GetString(1), "Hello"); |
| ASSERT_TRUE(missing.IsNull(0)); |
| ASSERT_TRUE(missing.IsNull(1)); |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadTypedOnlyPrunesValue) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", |
| R"({"age": 25, "other": "Hello"})"}; |
| WriteShreddedFile(jsons, kAgeCityShreddingSchema); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.city"}}); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); |
| |
| // All requested keys are shredded: neither `value` nor the unrequested typed keys are read. |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| const auto& physical_type = |
| static_cast<const arrow::StructType&>(*plans.at("v")->PhysicalField()->type()); |
| ASSERT_EQ(physical_type.GetFieldByName(VariantDefs::kValueFieldName), nullptr); |
| auto typed_value = physical_type.GetFieldByName(VariantDefs::kTypedValueFieldName); |
| ASSERT_NE(typed_value, nullptr); |
| ASSERT_EQ(typed_value->type()->num_fields(), 2); |
| file_reader->Close(); |
| } |
| |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| const auto& age = static_cast<const arrow::Int64Array&>(*v_column->field(0)); |
| const auto& city = static_cast<const arrow::StringArray&>(*v_column->field(1)); |
| ASSERT_EQ(age.Value(0), 35); |
| ASSERT_EQ(age.Value(1), 25); |
| ASSERT_EQ(city.GetString(0), "Chicago"); |
| // Row 1 has no "city" key: the shredded field is missing, which reads as null. |
| ASSERT_TRUE(city.IsNull(1)); |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadUnshreddedFile) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", |
| R"({"age": 25, "other": "Hello"})", nullptr}; |
| WriteFile(BuildArray(jsons)); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}}); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); |
| |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| ASSERT_EQ(v_column->length(), 3); |
| const auto& age = static_cast<const arrow::Int64Array&>(*v_column->field(0)); |
| const auto& other = static_cast<const arrow::StringArray&>(*v_column->field(1)); |
| ASSERT_EQ(age.Value(0), 35); |
| ASSERT_EQ(age.Value(1), 25); |
| ASSERT_TRUE(v_column->IsNull(2)); |
| ASSERT_TRUE(other.IsNull(0)); |
| ASSERT_EQ(other.GetString(1), "Hello"); |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadSemicolonKey) { |
| // Object keys may contain the description delimiter; the description parser anchors on the |
| // trailing failOnError/timeZoneId tokens instead of splitting on every delimiter. |
| std::vector<const char*> jsons = {R"({"a;b": 7})"}; |
| WriteFile(BuildArray(jsons)); |
| auto access_field = BuildAccessField({{arrow::int64(), "$['a;b']"}}); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| ASSERT_EQ(static_cast<const arrow::Int64Array&>(*v_column->field(0)).Value(0), 7); |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadVariantTarget) { |
| std::vector<const char*> jsons = {R"({"user": {"name": "Paimon", "age": 1}})", |
| R"({"user": "flat"})"}; |
| WriteFile(BuildArray(jsons)); |
| |
| // A variant-marked target re-encodes the extracted sub-variant instead of casting it to a |
| // plain struct; the marker on the target field must survive AddField. |
| VariantAccessBuilder builder; |
| ASSERT_OK_AND_ASSIGN(auto variant_target, Variant::ArrowField("t")); |
| ASSERT_OK(builder.AddField(variant_target.get(), "$.user")); |
| ASSERT_OK_AND_ASSIGN(auto c_access_field, builder.Build("v")); |
| auto imported = arrow::ImportField(c_access_field.get()); |
| ASSERT_TRUE(imported.ok()) << imported.status().ToString(); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), imported.ValueOrDie()}); |
| |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| const auto& user = static_cast<const arrow::StructArray&>(*v_column->field(0)); |
| const auto& value_column = static_cast<const arrow::BinaryArray&>(*user.field(0)); |
| const auto& metadata_column = static_cast<const arrow::BinaryArray&>(*user.field(1)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> row0, |
| GenericVariant::Create(value_column.GetView(0), metadata_column.GetView(0), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string row0_json, row0->ToJson()); |
| ASSERT_EQ(row0_json, R"({"age":1,"name":"Paimon"})"); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> row1, |
| GenericVariant::Create(value_column.GetView(1), metadata_column.GetView(1), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string row1_json, row1->ToJson()); |
| ASSERT_EQ(row1_json, R"("flat")"); |
| } |
| |
| TEST_F(VariantParquetTest, VariantAccessReadNestedPath) { |
| const char* nested_shredding_schema = R"({ |
| "type": "ROW", |
| "fields": [ { |
| "id": 0, |
| "name": "v", |
| "type": { |
| "type": "ROW", |
| "fields": [ { |
| "id": 1, |
| "name": "address", |
| "type": { |
| "type": "ROW", |
| "fields": [ {"id": 2, "name": "city", "type": "STRING"} ] |
| } |
| } ] |
| } |
| } ] |
| })"; |
| std::vector<const char*> jsons = {R"({"address": {"city": "Hangzhou"}})", |
| R"({"address": "oops"})", R"({"address": {"zip": 12345}})"}; |
| WriteShreddedFile(jsons, nested_shredding_schema); |
| |
| auto access_field = BuildAccessField({{arrow::utf8(), "$.address.city"}}); |
| auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); |
| |
| std::shared_ptr<arrow::StructArray> v_column; |
| ReadVariantColumn(read_schema, &v_column); |
| const auto& city = static_cast<const arrow::StringArray&>(*v_column->field(0)); |
| ASSERT_EQ(city.GetString(0), "Hangzhou"); |
| // Row 1's address is not an object; row 2's address has no "city" key. |
| ASSERT_TRUE(city.IsNull(1)); |
| ASSERT_TRUE(city.IsNull(2)); |
| } |
| |
| TEST_F(VariantParquetTest, NestedVariantPlainReadReassembles) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", nullptr}; |
| WriteShreddedNestedFile(jsons, arrow::struct_({arrow::field("age", arrow::int32()), |
| arrow::field("city", arrow::utf8())})); |
| |
| // Read as a plain nested VARIANT: the shredded sub-columns are reassembled back into |
| // `struct<value, metadata>`. |
| std::shared_ptr<arrow::StructArray> s_column; |
| ReadVariantColumn(nested_schema_, &s_column); |
| const auto& nv = static_cast<const arrow::StructArray&>(*s_column->field(0)); |
| ASSERT_TRUE(nv.type()->Equals(*nested_variant_field_->type())); |
| const auto& value_column = static_cast<const arrow::BinaryArray&>(*nv.field(0)); |
| const auto& metadata_column = static_cast<const arrow::BinaryArray&>(*nv.field(1)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<GenericVariant> row0, |
| GenericVariant::Create(value_column.GetView(0), metadata_column.GetView(0), pool_)); |
| ASSERT_OK_AND_ASSIGN(std::string row0_json, row0->ToJson()); |
| ASSERT_EQ(row0_json, R"({"age":35,"city":"Chicago"})"); |
| ASSERT_TRUE(nv.IsNull(1)); |
| ExpectNestedSibling(s_column); |
| } |
| |
| TEST_F(VariantParquetTest, NestedVariantAccessReadShreddedFile) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", |
| R"({"age": 25, "other": "Hello"})", nullptr}; |
| WriteShreddedNestedFile(jsons, arrow::struct_({arrow::field("age", arrow::int32()), |
| arrow::field("city", arrow::utf8())})); |
| |
| auto read_schema = |
| BuildNestedAccessReadSchema({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}}); |
| |
| // A nested projection prunes the scan the same way a top-level one does: `typed_value` is |
| // narrowed to `age`, and `value` is kept because `$.other` is not shredded. |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| ASSERT_EQ(plans.size(), 1); |
| const auto& physical_struct = |
| static_cast<const arrow::StructType&>(*plans.at("s")->PhysicalField()->type()); |
| auto physical_nv = physical_struct.GetFieldByName("nv"); |
| ASSERT_NE(physical_nv, nullptr); |
| const auto& physical_nv_type = static_cast<const arrow::StructType&>(*physical_nv->type()); |
| ASSERT_NE(physical_nv_type.GetFieldByName(VariantDefs::kMetadataFieldName), nullptr); |
| ASSERT_NE(physical_nv_type.GetFieldByName(VariantDefs::kValueFieldName), nullptr); |
| auto typed_value = physical_nv_type.GetFieldByName(VariantDefs::kTypedValueFieldName); |
| ASSERT_NE(typed_value, nullptr); |
| const auto& typed_struct = static_cast<const arrow::StructType&>(*typed_value->type()); |
| ASSERT_EQ(typed_struct.num_fields(), 1); |
| ASSERT_NE(typed_struct.GetFieldByName("age"), nullptr); |
| file_reader->Close(); |
| } |
| |
| std::shared_ptr<arrow::StructArray> s_column; |
| ReadVariantColumn(read_schema, &s_column); |
| ASSERT_EQ(s_column->length(), 3); |
| const auto& nv = static_cast<const arrow::StructArray&>(*s_column->field(0)); |
| const auto& age = static_cast<const arrow::Int64Array&>(*nv.field(0)); |
| const auto& other = static_cast<const arrow::StringArray&>(*nv.field(1)); |
| // The extracted values must match the requested access struct, not `struct<value, metadata>`. |
| ASSERT_EQ(nv.num_fields(), 2); |
| ASSERT_EQ(age.Value(0), 35); |
| ASSERT_EQ(age.Value(1), 25); |
| ASSERT_TRUE(other.IsNull(0)); |
| ASSERT_EQ(other.GetString(1), "Hello"); |
| // Row 2's nested variant is null, so the whole projection is null. |
| ASSERT_TRUE(nv.IsNull(2)); |
| ExpectNestedSibling(s_column); |
| } |
| |
| TEST_F(VariantParquetTest, NestedVariantAccessReadUnshreddedFile) { |
| std::vector<const char*> jsons = {R"({"age": 35, "city": "Chicago"})", |
| R"({"age": 25, "other": "Hello"})", nullptr}; |
| WriteNestedFile(jsons); |
| |
| // The nested column is stored unshredded: the paths are extracted from the `value` binary, |
| // which still requires a read plan (the projection type must never reach the format reader). |
| auto read_schema = |
| BuildNestedAccessReadSchema({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}}); |
| std::shared_ptr<arrow::StructArray> s_column; |
| ReadVariantColumn(read_schema, &s_column); |
| ASSERT_EQ(s_column->length(), 3); |
| const auto& nv = static_cast<const arrow::StructArray&>(*s_column->field(0)); |
| const auto& age = static_cast<const arrow::Int64Array&>(*nv.field(0)); |
| const auto& other = static_cast<const arrow::StringArray&>(*nv.field(1)); |
| ASSERT_EQ(age.Value(0), 35); |
| ASSERT_EQ(age.Value(1), 25); |
| ASSERT_TRUE(other.IsNull(0)); |
| ASSERT_EQ(other.GetString(1), "Hello"); |
| ASSERT_TRUE(nv.IsNull(2)); |
| ExpectNestedSibling(s_column); |
| } |
| |
| TEST_F(VariantParquetTest, ListVariantPlainReadNeedsNoPlan) { |
| // A variant inside a repeated group is never shredded, so a plain read of it still needs no |
| // plan at all: the logical type is exactly what the file stores. |
| WriteListFile({{R"({"x": 1})", R"({"x": 2})"}, {R"({"x": 3})"}}); |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| list_schema_, file_schema, pool_)); |
| ASSERT_TRUE(plans.empty()); |
| file_reader->Close(); |
| } |
| |
| TEST_F(VariantParquetTest, ListVariantAccessRead) { |
| // A variant-access projection inside an ARRAY extracts the paths per element (as in Java's |
| // testReadNestedVariantInArray). The empty row and the null element cover the offsets and |
| // the element validity being carried over by the reassembly. |
| WriteListFile( |
| {{R"({"x": 1, "y": 2})", R"({"x": 3, "y": 4})"}, {}, {R"({"x": 5, "y": 6})", nullptr}}); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.x"}}, "element"); |
| auto read_schema = arrow::schema( |
| {list_schema_->field(0), list_schema_->field(1)->WithType(arrow::list(access_field))}); |
| |
| // The parquet reader rejects partial projection inside a repeated group, so the whole file |
| // subtree is pushed down and only reassembled back. |
| { |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| ASSERT_EQ(plans.size(), 1); |
| ASSERT_TRUE(plans.at("arr")->PhysicalField()->type()->Equals( |
| *file_schema->GetFieldByName("arr")->type())); |
| file_reader->Close(); |
| } |
| |
| std::shared_ptr<arrow::Array> arr_column; |
| ReadColumn(read_schema, &arr_column); |
| ASSERT_NE(arr_column, nullptr); |
| const auto& list = static_cast<const arrow::ListArray&>(*arr_column); |
| ASSERT_EQ(list.length(), 3); |
| ASSERT_EQ(list.value_length(0), 2); |
| ASSERT_EQ(list.value_length(1), 0); |
| ASSERT_EQ(list.value_length(2), 2); |
| const auto& elements = static_cast<const arrow::StructArray&>(*list.values()); |
| const auto& x = static_cast<const arrow::Int64Array&>(*elements.field(0)); |
| ASSERT_EQ(x.Value(list.value_offset(0)), 1); |
| ASSERT_EQ(x.Value(list.value_offset(0) + 1), 3); |
| ASSERT_EQ(x.Value(list.value_offset(2)), 5); |
| // The null element's whole projection is null. |
| ASSERT_TRUE(elements.IsNull(list.value_offset(2) + 1)); |
| } |
| |
| TEST_F(VariantParquetTest, ListOfStructVariantAccessRead) { |
| // A variant one struct level below the ARRAY: the plan descends list -> struct -> variant |
| // and must leave the struct's plain sibling untouched. |
| WriteListStructFile({{R"({"x": 1, "y": 2})", R"({"x": 3})"}, {R"({"x": 5})"}}); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.x"}}, "v"); |
| auto element_field = |
| arrow::field("element", arrow::struct_({access_field, nested_sibling_field_})); |
| auto read_schema = |
| arrow::schema({list_struct_schema_->field(0), |
| list_struct_schema_->field(1)->WithType(arrow::list(element_field))}); |
| |
| std::shared_ptr<arrow::Array> arr_column; |
| ReadColumn(read_schema, &arr_column); |
| ASSERT_NE(arr_column, nullptr); |
| const auto& list = static_cast<const arrow::ListArray&>(*arr_column); |
| ASSERT_EQ(list.length(), 2); |
| const auto& elements = static_cast<const arrow::StructArray&>(*list.values()); |
| // Each element is `struct<v: <access projection>, t: STRING>`, so the extracted path sits one |
| // struct level below the element. |
| const auto& access = static_cast<const arrow::StructArray&>(*elements.field(0)); |
| const auto& x = static_cast<const arrow::Int64Array&>(*access.field(0)); |
| const auto& sibling = static_cast<const arrow::StringArray&>(*elements.field(1)); |
| ASSERT_EQ(x.Value(0), 1); |
| ASSERT_EQ(x.Value(1), 3); |
| ASSERT_EQ(x.Value(2), 5); |
| for (int64_t i = 0; i < 3; ++i) { |
| EXPECT_EQ(sibling.GetString(i), NestedSiblingValue(static_cast<size_t>(i))); |
| } |
| } |
| |
| TEST_F(VariantParquetTest, ListPartialProjectionNeedsNoPlan) { |
| // Projecting a subset of a struct inside a repeated group is unsupported, so no plan is |
| // built: the read fails in the reader instead of assembling a mistyped column. |
| WriteListStructFile({{R"({"x": 1})"}}); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.x"}}, "v"); |
| auto element_field = arrow::field("element", arrow::struct_({access_field})); |
| auto read_schema = |
| arrow::schema({list_struct_schema_->field(0), |
| list_struct_schema_->field(1)->WithType(arrow::list(element_field))}); |
| |
| std::unique_ptr<FileBatchReader> file_reader; |
| std::shared_ptr<arrow::Schema> file_schema; |
| OpenFile(&file_reader, &file_schema); |
| ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( |
| read_schema, file_schema, pool_)); |
| ASSERT_TRUE(plans.empty()); |
| file_reader->Close(); |
| } |
| |
| TEST_F(VariantParquetTest, MapVariantAccessRead) { |
| // The same for a variant value inside a MAP: the plan descends through the map entries |
| // struct and rewrites only the item child. |
| WriteMapFile({{"a", "b"}, {"c", "d"}}, |
| {{R"({"x": 1, "y": 2})", R"({"x": 3})"}, {R"({"x": 5, "y": 6})", nullptr}}); |
| |
| auto access_field = BuildAccessField({{arrow::int64(), "$.x"}}, "value"); |
| auto read_schema = |
| arrow::schema({map_schema_->field(0), |
| map_schema_->field(1)->WithType(arrow::map(arrow::utf8(), access_field))}); |
| |
| std::shared_ptr<arrow::Array> m_column; |
| ReadColumn(read_schema, &m_column); |
| ASSERT_NE(m_column, nullptr); |
| const auto& map = static_cast<const arrow::MapArray&>(*m_column); |
| ASSERT_EQ(map.length(), 2); |
| ASSERT_EQ(map.value_length(0), 2); |
| ASSERT_EQ(map.value_length(1), 2); |
| // The keys must survive untouched next to the rewritten values. |
| const auto& keys = static_cast<const arrow::StringArray&>(*map.keys()); |
| ASSERT_EQ(keys.GetString(map.value_offset(0)), "a"); |
| ASSERT_EQ(keys.GetString(map.value_offset(0) + 1), "b"); |
| ASSERT_EQ(keys.GetString(map.value_offset(1)), "c"); |
| ASSERT_EQ(keys.GetString(map.value_offset(1) + 1), "d"); |
| const auto& items = static_cast<const arrow::StructArray&>(*map.items()); |
| const auto& x = static_cast<const arrow::Int64Array&>(*items.field(0)); |
| ASSERT_EQ(x.Value(map.value_offset(0)), 1); |
| ASSERT_EQ(x.Value(map.value_offset(0) + 1), 3); |
| ASSERT_EQ(x.Value(map.value_offset(1)), 5); |
| // A null map value projects to a null row. |
| ASSERT_TRUE(items.IsNull(map.value_offset(1) + 1)); |
| } |
| |
| } // namespace paimon::parquet::test |
