| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| #include "paimon/format/parquet/parquet_file_batch_reader.h" |
| |
| #include <atomic> |
| #include <functional> |
| #include <iostream> |
| #include <limits> |
| #include <memory> |
| #include <string> |
| |
| #include "arrow/api.h" |
| #include "arrow/array/array_base.h" |
| #include "arrow/array/array_nested.h" |
| #include "arrow/array/builder_nested.h" |
| #include "arrow/array/builder_primitive.h" |
| #include "arrow/c/abi.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/io/caching.h" |
| #include "arrow/io/interfaces.h" |
| #include "arrow/ipc/api.h" |
| #include "arrow/ipc/json_simple.h" |
| #include "gtest/gtest.h" |
| #include "paimon/common/types/data_field.h" |
| #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" |
| #include "paimon/common/utils/arrow/arrow_utils.h" |
| #include "paimon/common/utils/arrow/mem_utils.h" |
| #include "paimon/common/utils/checked_cast.h" |
| #include "paimon/common/utils/date_time_utils.h" |
| #include "paimon/common/utils/path_util.h" |
| #include "paimon/defs.h" |
| #include "paimon/format/parquet/parquet_field_id_converter.h" |
| #include "paimon/format/parquet/parquet_format_defs.h" |
| #include "paimon/format/parquet/parquet_format_writer.h" |
| #include "paimon/format/parquet/parquet_reader_builder.h" |
| #include "paimon/fs/file_system.h" |
| #include "paimon/fs/local/local_file_system.h" |
| #include "paimon/memory/memory_pool.h" |
| #include "paimon/predicate/literal.h" |
| #include "paimon/predicate/predicate_builder.h" |
| #include "paimon/reader/batch_reader.h" |
| #include "paimon/testing/utils/counting_cache_test_utils.h" |
| #include "paimon/testing/utils/read_result_collector.h" |
| #include "paimon/testing/utils/testharness.h" |
| #include "paimon/testing/utils/timezone_guard.h" |
| #include "paimon/utils/roaring_bitmap32.h" |
| #include "parquet/properties.h" |
| |
| namespace paimon { |
| class Predicate; |
| } // namespace paimon |
| |
| namespace paimon::parquet::test { |
| |
| std::string SerializeSchemaToString(const std::shared_ptr<arrow::Schema>& schema) { |
| std::shared_ptr<arrow::Buffer> serialized = arrow::ipc::SerializeSchema(*schema).ValueOrDie(); |
| return std::string(reinterpret_cast<const char*>(serialized->data()), |
| static_cast<size_t>(serialized->size())); |
| } |
| |
| class FailedUriInputStream : public InputStream { |
| public: |
| explicit FailedUriInputStream(const std::shared_ptr<InputStream>& input) : input_(input) {} |
| |
| Status Seek(int64_t offset, SeekOrigin origin) override { |
| return input_->Seek(offset, origin); |
| } |
| |
| Result<int64_t> GetPos() const override { |
| return input_->GetPos(); |
| } |
| |
| Result<int64_t> Read(char* buffer, int64_t size) override { |
| return input_->Read(buffer, size); |
| } |
| |
| Result<int64_t> Read(char* buffer, int64_t size, int64_t offset) override { |
| return input_->Read(buffer, size, offset); |
| } |
| |
| void ReadAsync(char* buffer, int64_t size, int64_t offset, |
| std::function<void(Status)>&& callback) override { |
| return input_->ReadAsync(buffer, size, offset, std::move(callback)); |
| } |
| |
| Result<std::string> GetUri() const override { |
| return Status::Invalid("failed to get uri"); |
| } |
| |
| Result<int64_t> Length() const override { |
| return input_->Length(); |
| } |
| |
| Status Close() override { |
| return input_->Close(); |
| } |
| |
| private: |
| std::shared_ptr<InputStream> input_; |
| }; |
| |
| class ParquetFileBatchReaderTest : public ::testing::Test, |
| public ::testing::WithParamInterface<bool> { |
| public: |
| static std::shared_ptr<arrow::Field> WithMapSelectedKeys( |
| const std::shared_ptr<arrow::Field>& field, const std::string& selected_keys) { |
| auto metadata = |
| field->metadata() ? field->metadata()->Copy() : arrow::key_value_metadata({}); |
| auto set_status = metadata->Set(DataField::MAP_SELECTED_KEYS, selected_keys); |
| EXPECT_TRUE(set_status.ok()) << set_status.ToString(); |
| return field->WithMetadata(metadata); |
| } |
| |
| static std::shared_ptr<arrow::Schema> MakeReadSchema(const arrow::FieldVector& fields) { |
| return arrow::schema(fields); |
| } |
| |
| void SetUp() override { |
| dir_ = paimon::test::UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir_); |
| fs_ = std::make_shared<LocalFileSystem>(); |
| pool_ = GetArrowPool(GetDefaultPool()); |
| batch_size_ = 10; |
| file_path_ = PathUtil::JoinPath(dir_->Str(), "test.parquet"); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f1", arrow::boolean()), |
| arrow::field("f2", arrow::int8()), |
| arrow::field("f3", arrow::int16()), |
| arrow::field("f4", arrow::int32()), |
| arrow::field("f5", arrow::int64()), |
| arrow::field("f6", arrow::float32()), |
| arrow::field("f7", arrow::float64()), |
| arrow::field("f8", arrow::utf8()), |
| arrow::field("f9", arrow::binary()), |
| arrow::field("f10", arrow::map(arrow::list(arrow::float32()), |
| arrow::struct_({arrow::field("f0", arrow::boolean()), |
| arrow::field("f1", arrow::int64())}))), |
| arrow::field("f11", arrow::timestamp(arrow::TimeUnit::NANO)), |
| arrow::field("f12", arrow::date32()), |
| arrow::field("f13", arrow::decimal128(2, 2))}; |
| |
| schema_ = arrow::schema(fields); |
| struct_array_ = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ |
| [true, 31, null, 300001, 3000000001, 3.1, 300.000000001, "s31", "a31", [[[5.11, 5.21], [true, 61]]], "1970-01-01 00:00:00.000071", 81, "0.91"], |
| [false, 32, 302, 300002, 3000000002, 3.2, 300.000000002, "s32", "a32", [[[5.12, 5.22], [false, 62]]], "1970-01-01 00:00:00.000072", 82, "0.92"], |
| [true, 33, 303, 300003, 3000000003, 3.3, 300.000000003, "s33", "a33", null, "1970-01-01 00:00:00.000073", 83, "0.93"], |
| [false, 34, 304, 300004, 3000000004, 3.4, 300.000000004, "s34", "a34", [[[5.141, 5.241], [false, 641]], [[5.14, 5.24], [false, 64]]], "1970-01-01 00:00:00.000074", 84, "0.94"], |
| [true, 35, 305, 300005, 3000000005, 3.5, 300.000000005, "s35", "a35", [[[5.15, 5.25], [true, 65]]], "1970-01-01 00:00:00.000075", 85, "0.95"], |
| [false, 36, 306, 300006, 3000000006, 3.6, 300.000000006, "s36", "a36", [[[5.16, 5.26], null]], "1970-01-01 00:00:00.000076", 86, "0.96"] |
| ])") |
| .ValueOrDie()); |
| } |
| |
| void TearDown() override {} |
| |
| void WriteArray(const std::string& file_path, const std::shared_ptr<arrow::Array>& src_array, |
| const std::shared_ptr<arrow::Schema>& arrow_schema, int64_t write_batch_size, |
| bool enable_dictionary, int64_t max_row_group_length, |
| int64_t max_page_size = 1024 * 1024 * 1024) const { |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out, |
| fs_->Create(file_path, /*overwrite=*/true)); |
| ::parquet::WriterProperties::Builder builder; |
| builder.write_batch_size(write_batch_size); |
| builder.max_row_group_length(max_row_group_length); |
| builder.data_pagesize(max_page_size); |
| builder.enable_write_page_index(); |
| enable_dictionary ? builder.enable_dictionary() : builder.disable_dictionary(); |
| auto writer_properties = builder.build(); |
| ASSERT_OK_AND_ASSIGN(auto format_writer, ParquetFormatWriter::Create( |
| out, arrow_schema, writer_properties, |
| DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, pool_)); |
| |
| auto arrow_array = std::make_unique<ArrowArray>(); |
| ASSERT_TRUE(arrow::ExportArray(*src_array, arrow_array.get()).ok()); |
| ASSERT_OK(format_writer->AddBatch(arrow_array.get())); |
| ASSERT_OK(format_writer->Flush()); |
| ASSERT_OK(format_writer->Finish()); |
| ASSERT_OK(out->Flush()); |
| ASSERT_OK(out->Close()); |
| } |
| |
| std::unique_ptr<ParquetFileBatchReader> PrepareParquetFileBatchReader( |
| const std::string& file_name, const std::shared_ptr<arrow::Schema>& read_schema, |
| const std::shared_ptr<Predicate>& predicate, |
| const std::optional<RoaringBitmap32>& selection_bitmap, int32_t batch_size, |
| bool enable_page_level_filter = false) const { |
| EXPECT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_name)); |
| auto length = fs_->GetFileStatus(file_name).value().GetLen(); |
| auto in_stream = |
| std::make_unique<ArrowInputStreamAdapter>(std::move(input_stream), length, pool_); |
| auto storage_read_bytes = in_stream->StorageReadBytes(); |
| std::map<std::string, std::string> options; |
| options[PARQUET_READ_ENABLE_PAGE_INDEX_FILTER] = |
| enable_page_level_filter ? "true" : "false"; |
| return PrepareParquetFileBatchReader(std::move(in_stream), options, read_schema, predicate, |
| selection_bitmap, batch_size, |
| std::move(storage_read_bytes)); |
| } |
| |
| std::unique_ptr<paimon::parquet::ParquetFileBatchReader> PrepareParquetFileBatchReader( |
| std::unique_ptr<arrow::io::RandomAccessFile>&& in_stream, |
| const std::map<std::string, std::string>& options, |
| const std::shared_ptr<arrow::Schema>& read_schema, |
| const std::shared_ptr<Predicate>& predicate, |
| const std::optional<RoaringBitmap32>& selection_bitmap, int32_t batch_size, |
| std::shared_ptr<std::atomic<uint64_t>> storage_read_bytes = nullptr) const { |
| EXPECT_OK_AND_ASSIGN(auto parquet_batch_reader, |
| ParquetFileBatchReader::Create( |
| std::move(in_stream), options, batch_size, |
| /*file_metadata=*/nullptr, std::move(storage_read_bytes), pool_)); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); |
| EXPECT_TRUE(arrow_status.ok()); |
| EXPECT_OK(parquet_batch_reader->SetReadSchema(c_schema.get(), predicate, selection_bitmap)); |
| return parquet_batch_reader; |
| } |
| |
| protected: |
| std::string file_path_; |
| std::unique_ptr<paimon::test::UniqueTestDirectory> dir_; |
| std::shared_ptr<FileSystem> fs_; |
| std::shared_ptr<arrow::MemoryPool> pool_; |
| int32_t batch_size_; |
| std::shared_ptr<arrow::Schema> schema_; |
| std::shared_ptr<arrow::StructArray> struct_array_; |
| }; |
| |
| static std::shared_ptr<arrow::StructArray> MakeSequentialIntData(int32_t num_rows) { |
| arrow::Int32Builder val_builder; |
| EXPECT_TRUE(val_builder.Reserve(num_rows).ok()); |
| for (int32_t i = 0; i < num_rows; ++i) { |
| val_builder.UnsafeAppend(i); |
| } |
| auto val_array = val_builder.Finish().ValueOrDie(); |
| auto field = arrow::field("f0", arrow::int32()); |
| return arrow::StructArray::Make({val_array}, {field}).ValueOrDie(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestParquetMetadataCacheReusesSerializedFooter) { |
| WriteArray(file_path_, struct_array_, schema_, /*write_batch_size=*/struct_array_->length(), |
| /*enable_dictionary=*/false, |
| /*max_row_group_length=*/struct_array_->length()); |
| |
| auto cache = std::make_shared<paimon::test::CountingRoutingCache>(CacheKind::DATA_FILE_FOOTER, |
| 128 * 1024 * 1024); |
| auto open_reader = [&]() -> Result<std::unique_ptr<ParquetFileBatchReader>> { |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input_stream, fs_->Open(file_path_)); |
| std::map<std::string, std::string> options; |
| ParquetReaderBuilder builder(options, batch_size_); |
| builder.WithMemoryPool(GetDefaultPool())->WithCache(cache); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileBatchReader> reader, |
| builder.Build(input_stream)); |
| auto parquet_reader = dynamic_cast<ParquetFileBatchReader*>(reader.release()); |
| if (parquet_reader == nullptr) { |
| return Status::Invalid("failed to cast FileBatchReader to ParquetFileBatchReader"); |
| } |
| return std::unique_ptr<ParquetFileBatchReader>(parquet_reader); |
| }; |
| |
| ASSERT_OK_AND_ASSIGN(auto reader1, open_reader()); |
| ASSERT_OK_AND_ASSIGN(auto schema1, reader1->GetFileSchema()); |
| ASSERT_TRUE(schema1); |
| ASSERT_TRUE(schema1->release); |
| schema1->release(schema1.get()); |
| ASSERT_EQ(1, cache->GetCount()); |
| ASSERT_EQ(1, cache->SupplierCallCount()); |
| ASSERT_EQ(1, cache->Size()); |
| ASSERT_EQ(CacheKind::DATA_FILE_FOOTER, cache->LastKind()); |
| |
| ASSERT_OK_AND_ASSIGN(auto reader2, open_reader()); |
| ASSERT_OK_AND_ASSIGN(auto schema2, reader2->GetFileSchema()); |
| ASSERT_TRUE(schema2); |
| ASSERT_TRUE(schema2->release); |
| schema2->release(schema2.get()); |
| ASSERT_EQ(2, cache->GetCount()); |
| ASSERT_EQ(1, cache->SupplierCallCount()); |
| ASSERT_EQ(1, cache->Size()); |
| ASSERT_EQ(CacheKind::DATA_FILE_FOOTER, cache->LastKind()); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestParquetMetadataCacheBypassesWhenGetUriFails) { |
| WriteArray(file_path_, struct_array_, schema_, /*write_batch_size=*/struct_array_->length(), |
| /*enable_dictionary=*/false, |
| /*max_row_group_length=*/struct_array_->length()); |
| |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> input_stream, fs_->Open(file_path_)); |
| auto failed_uri_input_stream = std::make_shared<FailedUriInputStream>(input_stream); |
| auto cache = std::make_shared<paimon::test::CountingRoutingCache>(CacheKind::DATA_FILE_FOOTER, |
| 128 * 1024 * 1024); |
| |
| std::map<std::string, std::string> options; |
| ParquetReaderBuilder builder(options, batch_size_); |
| builder.WithMemoryPool(GetDefaultPool())->WithCache(cache); |
| ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileBatchReader> reader, |
| builder.Build(failed_uri_input_stream)); |
| auto parquet_reader = dynamic_cast<ParquetFileBatchReader*>(reader.get()); |
| ASSERT_TRUE(parquet_reader); |
| ASSERT_OK_AND_ASSIGN(auto file_schema, parquet_reader->GetFileSchema()); |
| ASSERT_TRUE(file_schema); |
| ASSERT_TRUE(file_schema->release); |
| file_schema->release(file_schema.get()); |
| |
| ASSERT_EQ(0, cache->GetCount()); |
| ASSERT_EQ(0, cache->SupplierCallCount()); |
| ASSERT_EQ(0, cache->Size()); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestReadBinaryWrittenFromBinaryAndLargeBinary) { |
| auto check_binary_read_result = [&](const std::shared_ptr<arrow::DataType>& write_type, |
| const std::string& file_name) { |
| std::string data_json = R"([ |
| ["descriptor-1"], |
| [""], |
| [null], |
| ["descriptor-2"] |
| ])"; |
| auto write_field = arrow::field("f0", write_type); |
| auto write_schema = arrow::schema({write_field}); |
| auto write_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({write_field}), data_json) |
| .ValueOrDie()); |
| |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), file_name); |
| WriteArray(file_path, write_array, write_schema, /*write_batch_size=*/write_array->length(), |
| /*enable_dictionary=*/false, /*max_row_group_length=*/write_array->length()); |
| |
| auto read_field = arrow::field("f0", arrow::binary()); |
| auto read_schema = arrow::schema({read_field}); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size_); |
| |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); |
| ASSERT_TRUE(file_schema->Equals(*read_schema)); |
| |
| auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({read_field}), data_json) |
| .ValueOrDie()); |
| auto expected_chunked_array = std::make_shared<arrow::ChunkedArray>(expected_array); |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| ASSERT_TRUE(result_array->Equals(expected_chunked_array)); |
| }; |
| |
| check_binary_read_result(arrow::binary(), "binary.parquet"); |
| check_binary_read_result(arrow::large_binary(), "large-binary.parquet"); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestSimple) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "/parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_name, schema_, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| std::shared_ptr<arrow::ChunkedArray> expected_array = |
| std::make_shared<arrow::ChunkedArray>(struct_array_); |
| ASSERT_TRUE(result_array->Equals(*expected_array, |
| arrow::EqualOptions::Defaults().diff_sink(&std::cout))); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestSetReadSchema) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> input_stream, fs_->Open(file_name)); |
| auto length = fs_->GetFileStatus(file_name).value().GetLen(); |
| auto in_stream = |
| std::make_unique<ArrowInputStreamAdapter>(std::move(input_stream), length, pool_); |
| std::map<std::string, std::string> options; |
| ASSERT_OK_AND_ASSIGN(auto parquet_batch_reader, |
| ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size_, |
| /*file_metadata=*/nullptr, |
| /*storage_read_bytes=*/nullptr, pool_)); |
| // test GetFileSchema() |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto arrow_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); |
| ASSERT_TRUE(arrow_file_schema->Equals(*schema_)); |
| |
| // NextBatch() without SetReadSchema(), will return data with |
| // file schema |
| ASSERT_OK_AND_ASSIGN( |
| auto result_with_file_schema, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| auto expected_file_chunk_array = std::make_shared<arrow::ChunkedArray>(struct_array_); |
| ASSERT_TRUE(result_with_file_schema->Equals(expected_file_chunk_array)); |
| // NextBatch() with SetReadSchema(), will return data with read schema |
| arrow::Schema read_schema({schema_->field(0), schema_->field(9), schema_->field(12)}); |
| std::unique_ptr<ArrowSchema> c_read_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(read_schema, c_read_schema.get()).ok()); |
| ASSERT_OK(parquet_batch_reader->SetReadSchema(c_read_schema.get(), |
| /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| auto expected_read_array = |
| arrow::StructArray::Make( |
| {struct_array_->field(0), struct_array_->field(9), struct_array_->field(12)}, |
| read_schema.fields()) |
| .ValueOrDie(); |
| auto expected_read_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_read_array); |
| ASSERT_OK_AND_ASSIGN( |
| auto result_with_read_schema, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_TRUE(result_with_read_schema->Equals(expected_read_chunk_array)); |
| |
| // NextBatch() with predicate |
| auto predicate = PredicateBuilder::IsNull( |
| /*field_index=*/0, /*field_name=*/"f1", FieldType::BOOLEAN); |
| c_read_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(read_schema, c_read_schema.get()).ok()); |
| ASSERT_OK(parquet_batch_reader->SetReadSchema(c_read_schema.get(), predicate, |
| /*selection_bitmap=*/std::nullopt)); |
| ASSERT_OK_AND_ASSIGN(result_with_read_schema, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| ASSERT_FALSE(result_with_read_schema); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestSetReadSchemaWithLegacyParquetMissingFieldIds) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "/parquet/append_09.db/append_09/f1=20/bucket-0/" |
| "data-b446f78a-2cfb-4b3b-add8-31295d24a277-0.parquet"; |
| |
| std::vector<DataField> read_fields = { |
| DataField(0, arrow::field("f0", arrow::utf8())), |
| DataField(2, arrow::field("f2", arrow::int32())), |
| DataField(3, arrow::field("f3", arrow::float64())), |
| }; |
| auto read_schema = DataField::ConvertDataFieldsToArrowSchema(read_fields); |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size_); |
| |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( |
| arrow::struct_(read_schema->fields()), {R"([ |
| ["Lucy", 1, 14.1] |
| ])"}, |
| &expected_array) |
| .ok()); |
| ASSERT_TRUE(result_array->Equals(expected_array)) |
| << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNextBatchSimple) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| for (auto batch_size : {1, 2, 3, 5, 8, 10}) { |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_name, schema_, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size); |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| parquet_batch_reader->Close(); |
| auto expected_array = std::make_shared<arrow::ChunkedArray>(struct_array_); |
| ASSERT_TRUE(result_array->Equals(expected_array)); |
| // test metrics |
| auto read_metrics = parquet_batch_reader->GetReaderMetrics(); |
| ASSERT_TRUE(read_metrics); |
| ASSERT_OK_AND_ASSIGN(uint64_t storage_read_bytes, |
| read_metrics->GetCounter(ParquetMetrics::READ_STORAGE_BYTES)); |
| ASSERT_GT(storage_read_bytes, 0u); |
| } |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithTargetSchema) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| auto read_schema = arrow::schema(arrow::FieldVector( |
| {schema_->field(4), schema_->field(9), schema_->field(10), schema_->field(12)})); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size_); |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| parquet_batch_reader->Close(); |
| auto expected_read_array = |
| arrow::StructArray::Make({struct_array_->field(4), struct_array_->field(9), |
| struct_array_->field(10), struct_array_->field(12)}, |
| read_schema->fields()) |
| .ValueOrDie(); |
| auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_read_array); |
| ASSERT_TRUE(result_array->Equals(expected_chunk_array)); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithOutofOrderTargetSchema) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| // read with f13, f11, f10, f5 |
| auto read_schema = arrow::schema(arrow::FieldVector( |
| {schema_->field(12), schema_->field(10), schema_->field(9), schema_->field(4)})); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size_); |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| parquet_batch_reader->Close(); |
| auto expected_read_array = |
| arrow::StructArray::Make({struct_array_->field(12), struct_array_->field(10), |
| struct_array_->field(9), struct_array_->field(4)}, |
| read_schema->fields()) |
| .ValueOrDie(); |
| auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_read_array); |
| ASSERT_TRUE(result_array->Equals(expected_chunk_array)); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNextBatchWithDictionary) { |
| auto f0 = arrow::field("f0", arrow::list(arrow::utf8())); |
| auto f1 = arrow::field("f1", arrow::map(arrow::utf8(), arrow::binary())); |
| auto f2 = arrow::field( |
| "f2", arrow::struct_({field("sub1", arrow::int64()), field("sub2", arrow::binary()), |
| field("sub3", arrow::utf8())})); |
| |
| arrow::FieldVector fields = {f0, f1, f2}; |
| auto src_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ |
| [["a", "a", "b"], [["a", "q"], ["b", "w"]], [10, "q", "a"]], |
| [["a", "c"], [["a", "e"], ["b", "r"], ["c", "e"]], [20, "w", "a"]], |
| [["a", "d"], [["d", "r"], ["e", "t"]], [null, "e", "b"]], |
| [["a"], [["a", "q"]], [null, "w", "c"]], |
| [null, [["a", "w"], ["f", "y"]], [50, "r", null]], |
| [["a"], null, [60, "r", "b"]], |
| [["a", "b", "e"], [["a", null], ["b", "w"]], null] |
| ])") |
| .ValueOrDie()); |
| auto src_schema = arrow::schema(fields); |
| auto arrow_schema = arrow::schema(fields); |
| auto expected_array = arrow::ChunkedArray::Make({src_array}).ValueOrDie(); |
| |
| auto check_result = [&](bool enable_dictionary) { |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1, enable_dictionary, |
| /*max_row_group_length=*/3); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, arrow_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| }; |
| check_result(true); |
| check_result(false); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNestedStructChildProjectionRecall) { |
| auto f0 = arrow::field("f0", arrow::int32()); |
| auto f1 = arrow::field( |
| "f1", arrow::struct_({arrow::field("c0", arrow::int64()), arrow::field("c1", arrow::utf8()), |
| arrow::field("c2", arrow::float64())})); |
| auto f2 = arrow::field("f2", arrow::utf8()); |
| |
| auto write_schema = arrow::schema({f0, f1, f2}); |
| auto write_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(write_schema->fields()), R"([ |
| [1, [100, "a", 1.1], "x"], |
| [2, [200, "b", 2.2], "y"], |
| [3, [300, null, 3.3], "z"] |
| ])") |
| .ValueOrDie()); |
| |
| WriteArray(file_path_, write_array, write_schema, |
| /*write_batch_size=*/write_array->length(), |
| /*enable_dictionary=*/false, /*max_row_group_length=*/write_array->length()); |
| |
| auto read_schema = MakeReadSchema({ |
| f0, |
| arrow::field("f1", arrow::struct_({arrow::field("c1", arrow::utf8())})), |
| }); |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( |
| arrow::struct_(read_schema->fields()), {R"([ |
| [1, ["a"]], |
| [2, ["b"]], |
| [3, [null]] |
| ])"}, |
| &expected_array) |
| .ok()); |
| |
| ASSERT_TRUE(result_array->Equals(expected_array)) |
| << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestReadSchemaWithMapSelectedKeysMetadata) { |
| auto id_field = arrow::field("id", arrow::int32()); |
| auto map_field = arrow::field("m", arrow::map(arrow::utf8(), arrow::int32())); |
| |
| auto write_schema = arrow::schema({id_field, map_field}); |
| auto write_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(write_schema->fields()), R"([ |
| [1, [["k1", 10], ["k2", 20], ["k3", 30]]], |
| [2, [["k2", 200]]], |
| [3, null] |
| ])") |
| .ValueOrDie()); |
| |
| WriteArray(file_path_, write_array, write_schema, |
| /*write_batch_size=*/write_array->length(), |
| /*enable_dictionary=*/false, /*max_row_group_length=*/write_array->length()); |
| |
| // selected-keys metadata is consumed by upper-level field mapping; format reader should |
| // still accept the schema and read data correctly. |
| auto read_schema = MakeReadSchema( |
| {id_field, WithMapSelectedKeys(map_field, "k1,k3")}); // NOLINT(whitespace/comma) |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| auto expected_array = arrow::ChunkedArray::Make({write_array}).ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) |
| << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNestedListTimestampTimezoneAndMapFieldName) { |
| const std::string timezone = "Asia/Shanghai"; |
| paimon::test::TimezoneGuard timezone_guard(timezone); |
| |
| auto write_attrs_type = |
| std::make_shared<arrow::MapType>(arrow::field("key", arrow::utf8(), /*nullable=*/false), |
| arrow::field("attrs", arrow::utf8())); |
| auto write_element_type = arrow::struct_({ |
| arrow::field("key", arrow::utf8()), |
| arrow::field("attrs", write_attrs_type), |
| arrow::field("updated_at", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)), |
| }); |
| auto write_schema = arrow::schema( |
| {arrow::field("annotations", arrow::list(arrow::field("element", write_element_type)))}); |
| |
| const std::string data_json = R"([ |
| [[ ["ann-1", [["source", "model"]], "2026-07-16 12:00:00.000001"] ]], |
| [[ ["ann-2", [], "2026-07-16 12:00:00.000002"], |
| ["ann-3", null, null] ]], |
| [null] |
| ])"; |
| auto write_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(write_schema->fields()), data_json) |
| .ValueOrDie()); |
| WriteArray(file_path_, write_array, write_schema, /*write_batch_size=*/write_array->length(), |
| /*enable_dictionary=*/false, /*max_row_group_length=*/write_array->length()); |
| |
| auto read_element_type = arrow::struct_({ |
| arrow::field("key", arrow::utf8()), |
| arrow::field("attrs", arrow::map(arrow::utf8(), arrow::utf8())), |
| arrow::field("updated_at", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)), |
| }); |
| auto read_schema = arrow::schema( |
| {arrow::field("annotations", arrow::list(arrow::field("element", read_element_type)))}); |
| auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_schema->fields()), data_json) |
| .ValueOrDie()); |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| auto expected_chunked_array = arrow::ChunkedArray::Make({expected_array}).ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_chunked_array)) |
| << "expected: " << expected_chunked_array->ToString() |
| << "\nactual: " << result_array->ToString(); |
| |
| auto projected_element_type = arrow::struct_({ |
| arrow::field("key", arrow::utf8()), |
| arrow::field("attrs", arrow::map(arrow::utf8(), arrow::utf8())), |
| }); |
| auto partial_read_schema = arrow::schema({arrow::field( |
| "annotations", arrow::list(arrow::field("element", projected_element_type)))}); |
| auto c_partial_read_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*partial_read_schema, c_partial_read_schema.get()).ok()); |
| ASSERT_NOK_WITH_MSG( |
| parquet_batch_reader->SetReadSchema(c_partial_read_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt), |
| "Parquet does not support partial projection inside list/map"); |
| |
| auto mismatched_element_type = arrow::struct_({ |
| arrow::field("key", arrow::utf8()), |
| arrow::field("attrs", arrow::map(arrow::utf8(), arrow::utf8())), |
| arrow::field("updated_at", arrow::utf8()), |
| }); |
| auto mismatched_read_schema = arrow::schema({arrow::field( |
| "annotations", arrow::list(arrow::field("element", mismatched_element_type)))}); |
| auto c_mismatched_read_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*mismatched_read_schema, c_mismatched_read_schema.get()).ok()); |
| ASSERT_NOK_WITH_MSG( |
| parquet_batch_reader->SetReadSchema(c_mismatched_read_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt), |
| "Parquet does not support partial projection inside list/map"); |
| |
| auto unsupported_timestamp_element_type = arrow::struct_({ |
| arrow::field("key", arrow::utf8()), |
| arrow::field("attrs", arrow::map(arrow::utf8(), arrow::utf8())), |
| arrow::field("updated_at", arrow::timestamp(arrow::TimeUnit::NANO, timezone)), |
| }); |
| auto unsupported_timestamp_schema = arrow::schema({arrow::field( |
| "annotations", arrow::list(arrow::field("element", unsupported_timestamp_element_type)))}); |
| auto c_unsupported_timestamp_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE( |
| arrow::ExportSchema(*unsupported_timestamp_schema, c_unsupported_timestamp_schema.get()) |
| .ok()); |
| ASSERT_NOK_WITH_MSG(parquet_batch_reader->SetReadSchema(c_unsupported_timestamp_schema.get(), |
| /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt), |
| "Parquet does not support partial projection inside list/map"); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestNestedTimestampSecondReadFromMilliFile) { |
| const std::string timezone = "Asia/Shanghai"; |
| paimon::test::TimezoneGuard timezone_guard(timezone); |
| |
| // Parquet has no second-precision timestamp, so the writer stores second timestamps as |
| // milliseconds. Reading them back with a second-precision schema must cast milli to second, |
| // including for timestamp leaves nested inside list/struct/map. |
| auto event_type = arrow::struct_({ |
| arrow::field("name", arrow::utf8()), |
| arrow::field("ts_sec", arrow::timestamp(arrow::TimeUnit::SECOND)), |
| arrow::field("ts_tz_sec", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)), |
| }); |
| auto schema = arrow::schema({ |
| arrow::field("events", arrow::list(arrow::field("element", event_type))), |
| arrow::field("marks", arrow::map(arrow::utf8(), arrow::timestamp(arrow::TimeUnit::SECOND))), |
| }); |
| |
| const std::string data_json = R"([ |
| [[ ["e-1", "2026-07-16 12:00:01", "2026-07-16 12:00:02"] ], |
| [["begin", "2026-07-16 12:00:03"]]], |
| [[ ["e-2", "2026-07-16 12:00:04", null], |
| ["e-3", null, "2026-07-16 12:00:05"] ], []], |
| [[null], null] |
| ])"; |
| auto write_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), data_json) |
| .ValueOrDie()); |
| WriteArray(file_path_, write_array, schema, /*write_batch_size=*/write_array->length(), |
| /*enable_dictionary=*/false, /*max_row_group_length=*/write_array->length()); |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| |
| // The nested second timestamps are physically stored as milliseconds in the file. |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOr(nullptr); |
| ASSERT_TRUE(file_schema); |
| auto file_event_type = |
| static_cast<const arrow::ListType&>(*file_schema->field(0)->type()).value_type(); |
| ASSERT_EQ(arrow::Type::STRUCT, file_event_type->id()); |
| ASSERT_EQ(arrow::TimeUnit::MILLI, |
| static_cast<const arrow::TimestampType&>(*file_event_type->field(1)->type()).unit()); |
| ASSERT_EQ(arrow::TimeUnit::MILLI, |
| static_cast<const arrow::TimestampType&>(*file_event_type->field(2)->type()).unit()); |
| auto file_mark_type = |
| static_cast<const arrow::MapType&>(*file_schema->field(1)->type()).item_type(); |
| ASSERT_EQ(arrow::TimeUnit::MILLI, |
| static_cast<const arrow::TimestampType&>(*file_mark_type).unit()); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| auto expected_array = arrow::ChunkedArray::Make({write_array}).ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) |
| << "expected: " << expected_array->ToString() << "\nactual: " << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestGetFileSchemaWithFieldId) { |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_name, schema_, /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto arrow_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOr(nullptr); |
| ASSERT_TRUE(arrow_file_schema); |
| ASSERT_OK_AND_ASSIGN(auto data_fields, |
| DataField::ConvertArrowSchemaToDataFields(arrow_file_schema)); |
| |
| auto list_type = arrow::list(DataField::ConvertDataFieldToArrowField( |
| DataField(536880129, arrow::field("item", arrow::float32())))); |
| std::vector<DataField> struct_fields = {DataField(10, arrow::field("f0", arrow::boolean())), |
| DataField(11, arrow::field("f1", arrow::int64()))}; |
| auto struct_type = DataField::ConvertDataFieldsToArrowStructType(struct_fields); |
| auto map_type = arrow::map(list_type, struct_type); |
| |
| std::vector<DataField> expected_data_fields = { |
| DataField(0, arrow::field("f1", arrow::boolean())), |
| DataField(1, arrow::field("f2", arrow::int8())), |
| DataField(2, arrow::field("f3", arrow::int16())), |
| DataField(3, arrow::field("f4", arrow::int32())), |
| DataField(4, arrow::field("f5", arrow::int64())), |
| DataField(5, arrow::field("f6", arrow::float32())), |
| DataField(6, arrow::field("f7", arrow::float64())), |
| DataField(7, arrow::field("f8", arrow::utf8())), |
| DataField(8, arrow::field("f9", arrow::binary())), |
| DataField(9, arrow::field("f10", map_type)), |
| DataField(12, arrow::field("f11", arrow::timestamp(arrow::TimeUnit::NANO))), |
| DataField(13, arrow::field("f12", arrow::date32())), |
| DataField(14, arrow::field("f13", arrow::decimal128(2, 2))), |
| }; |
| ASSERT_EQ(data_fields, expected_data_fields); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestCreateReaderProperties) { |
| { |
| // test default options |
| std::map<std::string, std::string> options; |
| ASSERT_OK_AND_ASSIGN(auto reader_properties, |
| ParquetFileBatchReader::CreateReaderProperties(pool_, options)); |
| ASSERT_EQ(reader_properties.is_buffered_stream_enabled(), true); |
| } |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { |
| { |
| // test default options |
| std::map<std::string, std::string> options; |
| int32_t batch_size = 1024; |
| ASSERT_OK_AND_ASSIGN( |
| auto arrow_reader_properties, |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); |
| ASSERT_EQ(arrow_reader_properties.pre_buffer(), true); |
| ASSERT_EQ(arrow_reader_properties.batch_size(), 1024); |
| ASSERT_EQ(arrow_reader_properties.use_threads(), true); |
| ASSERT_EQ(arrow::GetCpuThreadPoolCapacity(), 3); |
| auto expected_cache_options = arrow::io::CacheOptions::Defaults(); |
| expected_cache_options.hole_size_limit = DEFAULT_PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT; |
| ASSERT_EQ(arrow_reader_properties.cache_options(), expected_cache_options); |
| } |
| { |
| std::map<std::string, std::string> options = {{PARQUET_READ_EXECUTOR_THREAD_COUNT, "0"}}; |
| int32_t batch_size = 1024; |
| ASSERT_OK_AND_ASSIGN( |
| auto arrow_reader_properties, |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); |
| ASSERT_EQ(arrow_reader_properties.use_threads(), false); |
| } |
| { |
| std::map<std::string, std::string> options = {{PARQUET_READ_EXECUTOR_THREAD_COUNT, "6"}}; |
| int32_t batch_size = 1024; |
| ASSERT_OK_AND_ASSIGN( |
| auto arrow_reader_properties, |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); |
| ASSERT_EQ(arrow_reader_properties.use_threads(), true); |
| ASSERT_EQ(arrow::GetCpuThreadPoolCapacity(), 6); |
| } |
| { |
| std::map<std::string, std::string> options = { |
| {PARQUET_READ_CACHE_OPTION_LAZY, "true"}, |
| {PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT, "2"}, |
| {PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT, "1048576"}, |
| {PARQUET_READ_CACHE_OPTION_RANGE_SIZE_LIMIT, "8388608"}, |
| }; |
| ASSERT_OK_AND_ASSIGN( |
| auto arrow_reader_properties, |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024)); |
| const auto& cache_options = arrow_reader_properties.cache_options(); |
| ASSERT_TRUE(cache_options.lazy); |
| ASSERT_EQ(cache_options.prefetch_limit, 2); |
| ASSERT_EQ(cache_options.hole_size_limit, 1024 * 1024); |
| ASSERT_EQ(cache_options.range_size_limit, 8 * 1024 * 1024); |
| } |
| { |
| std::map<std::string, std::string> options = { |
| {PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT, "-1"}, |
| }; |
| ASSERT_NOK_WITH_MSG( |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024), |
| "parquet.read.cache-option.hole-size-limit must be non-negative"); |
| } |
| { |
| std::map<std::string, std::string> options = { |
| {PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT, "1048576"}, |
| {PARQUET_READ_CACHE_OPTION_RANGE_SIZE_LIMIT, "1048576"}, |
| }; |
| ASSERT_NOK_WITH_MSG( |
| ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024), |
| "parquet.read.cache-option.range-size-limit must be greater than " |
| "parquet.read.cache-option.hole-size-limit"); |
| } |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestBitmapRowGroupPushDownWithMultiRowGroups) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto arrow_type = arrow::struct_(fields); |
| auto src_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow_type, R"([ |
| [0], |
| [1], |
| [2], |
| [3], |
| [4], |
| [5], |
| [6], |
| [7], |
| [8], |
| [9], |
| [10], |
| [11] |
| ])") |
| .ValueOrDie()); |
| auto src_schema = arrow::schema(fields); |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({3, 5}); |
| // data in file rowGroup0:[0, 1, 2, 3, 4, 5] | rowGroup1:[6, 7, 8, 9, 10, 11] |
| |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/12, |
| /*enable_dictionary=*/true, |
| /*max_row_group_length=*/6); |
| |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, /*predicate=*/nullptr, bitmap, /*batch_size=*/12); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| |
| auto expected_array = arrow::ChunkedArray::Make({src_array->Slice(0, 6)}).ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| } |
| TEST_F(ParquetFileBatchReaderTest, TestBitmapPagePushDownWithMultiRowGroups) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto arrow_type = arrow::struct_(fields); |
| auto src_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow_type, R"([ |
| [0], |
| [1], |
| [2], |
| [3], |
| [4], |
| [5], |
| [6], |
| [7], |
| [8], |
| [9], |
| [10], |
| [11] |
| ])") |
| .ValueOrDie()); |
| auto src_schema = arrow::schema(fields); |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({3, 5}); |
| // data in file rowGroup0:[0, 1, 2, 3, 4, 5] | rowGroup1:[6, 7, 8, 9, 10, 11] |
| |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/12, |
| /*enable_dictionary=*/true, |
| /*max_row_group_length=*/6); |
| |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, arrow_schema, /*predicate=*/nullptr, bitmap, |
| /*batch_size=*/12, /*enable_page_level_filter=*/true); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| |
| auto expected_array = arrow::ChunkedArray(src_array->Slice(3, 3)); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapRowGroupPushDown) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto arrow_type = arrow::struct_(fields); |
| arrow::StructBuilder struct_builder(arrow_type, arrow::default_memory_pool(), |
| {std::make_shared<arrow::Int32Builder>()}); |
| auto int_builder = checked_cast<arrow::Int32Builder*>(struct_builder.field_builder(0)); |
| int32_t length = 1024; |
| for (int32_t i = 0; i < length; ++i) { |
| ASSERT_TRUE(struct_builder.Append().ok()); |
| ASSERT_TRUE(int_builder->Append(i).ok()); |
| } |
| // data file: |
| // rowGroup0: [0, 256) |
| // rowGroup1: [256, 512) |
| // rowGroup2: [512, 768) |
| // rowGroup3: [768, 1024) |
| std::shared_ptr<arrow::Array> src_array; |
| ASSERT_TRUE(struct_builder.Finish(&src_array).ok()); |
| auto src_schema = arrow::schema(fields); |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1024, |
| /*enable_dictionary=*/true, |
| /*max_row_group_length=*/256); |
| { |
| // simple case |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({100, 400, 600}); |
| ASSERT_OK_AND_ASSIGN( |
| auto predicate, |
| PredicateBuilder::Or( |
| {PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, |
| Literal(255)), |
| PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(600))})); |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, predicate, bitmap, /*batch_size=*/length); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| |
| auto expected_array = |
| arrow::ChunkedArray::Make({src_array->Slice(0, 256), src_array->Slice(512, 256)}) |
| .ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| } |
| { |
| // test all data has been filtered out with predicate and bitmap pushdown |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({100, 400, 600}); |
| auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(800)); |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, predicate, bitmap, /*batch_size=*/length); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_FALSE(result_array); |
| } |
| } |
| TEST_F(ParquetFileBatchReaderTest, TestPredicateAndBitmapPagePushDown) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto arrow_type = arrow::struct_(fields); |
| arrow::StructBuilder struct_builder(arrow_type, arrow::default_memory_pool(), |
| {std::make_shared<arrow::Int32Builder>()}); |
| auto int_builder = checked_cast<arrow::Int32Builder*>(struct_builder.field_builder(0)); |
| int32_t length = 1024; |
| for (int32_t i = 0; i < length; ++i) { |
| ASSERT_TRUE(struct_builder.Append().ok()); |
| ASSERT_TRUE(int_builder->Append(i).ok()); |
| } |
| // data file: |
| // rowGroup0: [0, 256) |
| // rowGroup1: [256, 512) |
| // rowGroup2: [512, 768) |
| // rowGroup3: [768, 1024) |
| std::shared_ptr<arrow::Array> src_array; |
| ASSERT_TRUE(struct_builder.Finish(&src_array).ok()); |
| auto src_schema = arrow::schema(fields); |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1024, |
| /*enable_dictionary=*/true, |
| /*max_row_group_length=*/256); |
| { |
| // simple case |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({100, 400, 600}); |
| ASSERT_OK_AND_ASSIGN( |
| auto predicate, |
| PredicateBuilder::Or( |
| {PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, |
| Literal(255)), |
| PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(600))})); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, arrow_schema, predicate, bitmap, |
| /*batch_size=*/length, /*enable_page_level_filter=*/true); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| |
| auto expected_array = |
| arrow::ChunkedArray::Make({src_array->Slice(100, 1), src_array->Slice(600, 1)}) |
| .ValueOrDie(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| } |
| { |
| // test all data has been filtered out with predicate and bitmap pushdown |
| std::optional<RoaringBitmap32> bitmap = RoaringBitmap32::From({100, 400, 600}); |
| auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(800)); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, arrow_schema, predicate, bitmap, |
| /*batch_size=*/length, /*enable_page_level_filter=*/true); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_FALSE(result_array); |
| } |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestReadNoField) { |
| // if only read partition fields, format reader will set empty read schema |
| std::string file_name = paimon::test::GetDataDir() + |
| "parquet/parquet_append_table.db/parquet_append_table/bucket-0/" |
| "data-9ea62f34-1dca-49c1-bf7a-d37303d8fb76-0.parquet"; |
| // read no field |
| auto read_schema = arrow::schema(arrow::FieldVector()); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_name, read_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| // read 2 rows |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN(auto batch1, parquet_batch_reader->NextBatch()); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 0); |
| // read 2 rows |
| ASSERT_OK_AND_ASSIGN(auto batch2, parquet_batch_reader->NextBatch()); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 2); |
| // read 2 rows |
| ASSERT_OK_AND_ASSIGN(auto batch3, parquet_batch_reader->NextBatch()); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 4); |
| // read rows with eof |
| ASSERT_OK_AND_ASSIGN(auto batch4, parquet_batch_reader->NextBatch()); |
| ASSERT_TRUE(BatchReader::IsEofBatch(batch4)); |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| parquet_batch_reader->Close(); |
| |
| arrow::FieldVector fields; |
| auto arrow_type = arrow::struct_(fields); |
| auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow_type, R"([ |
| [], |
| [] |
| ])") |
| .ValueOrDie()); |
| |
| auto check_batch = [](BatchReader::ReadBatch&& result_batch, |
| const std::shared_ptr<arrow::Array>& expected_array) { |
| auto& [c_array, c_schema] = result_batch; |
| auto result_array = arrow::ImportArray(c_array.get(), c_schema.get()).ValueOr(nullptr); |
| ASSERT_TRUE(result_array); |
| ASSERT_TRUE(result_array->Equals(expected_array)); |
| }; |
| |
| check_batch(std::move(batch1), expected_array); |
| check_batch(std::move(batch2), expected_array); |
| check_batch(std::move(batch3), expected_array); |
| } |
| |
| TEST_P(ParquetFileBatchReaderTest, TestTimestampType) { |
| auto enable_tz = GetParam(); |
| std::string timezone_str = enable_tz ? "Asia/Tokyo" : DateTimeUtils::GetLocalTimezoneName(); |
| paimon::test::TimezoneGuard tz_guard(timezone_str); |
| |
| arrow::FieldVector fields = { |
| arrow::field("ts_sec", arrow::timestamp(arrow::TimeUnit::SECOND)), |
| arrow::field("ts_milli", arrow::timestamp(arrow::TimeUnit::MILLI)), |
| arrow::field("ts_micro", arrow::timestamp(arrow::TimeUnit::MICRO)), |
| arrow::field("ts_nano", arrow::timestamp(arrow::TimeUnit::NANO)), |
| arrow::field("ts_tz_sec", arrow::timestamp(arrow::TimeUnit::SECOND, timezone_str)), |
| arrow::field("ts_tz_milli", arrow::timestamp(arrow::TimeUnit::MILLI, timezone_str)), |
| arrow::field("ts_tz_micro", arrow::timestamp(arrow::TimeUnit::MICRO, timezone_str)), |
| arrow::field("ts_tz_nano", arrow::timestamp(arrow::TimeUnit::NANO, timezone_str)), |
| }; |
| |
| auto array = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ |
| ["1970-01-01 00:00:01", "1970-01-01 00:00:00.001", "1970-01-01 00:00:00.000001", "1970-01-01 00:00:00.000000001", "1970-01-01 00:00:02", "1970-01-01 00:00:00.002", "1970-01-01 00:00:00.000002", "1970-01-01 00:00:00.000000002"], |
| ["1970-01-01 00:00:03", "1970-01-01 00:00:00.003", null, "1970-01-01 00:00:00.000000003", "1970-01-01 00:00:04", "1970-01-01 00:00:00.004", "1970-01-01 00:00:00.000004", "1970-01-01 00:00:00.000000004"], |
| ["1970-01-01 00:00:05", "1970-01-01 00:00:00.005", null, null, "1970-01-01 00:00:06", null, "1970-01-01 00:00:00.000006", null] |
| ])") |
| .ValueOrDie()); |
| std::shared_ptr<arrow::ChunkedArray> expected_array = |
| std::make_shared<arrow::ChunkedArray>(array); |
| |
| { |
| // read data generated by Java Paimon |
| std::string file_name = paimon::test::GetDataDir() + |
| "/parquet/append_with_multiple_ts_precision_and_timezone.db/" |
| "append_with_multiple_ts_precision_and_timezone/bucket-0/" |
| "data-9b8abdde-df4d-4655-bb4c-ffda164ef9d4-0.parquet"; |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_name, std::make_shared<arrow::Schema>(fields), |
| /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, batch_size_); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); |
| } |
| { |
| // read data generated by C++ Paimon |
| auto arrow_schema = std::make_shared<arrow::Schema>(fields); |
| WriteArray(file_path_, array, arrow_schema, |
| /*write_batch_size=*/1, /*enable_dictionary=*/true, |
| /*max_row_group_length=*/1); |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path_, arrow_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, /*batch_size=*/2); |
| |
| // check file schema |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto result_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOr(nullptr); |
| ASSERT_TRUE(result_file_schema); |
| |
| arrow::FieldVector expected_fields = { |
| arrow::field("ts_sec", arrow::timestamp(arrow::TimeUnit::MILLI)), |
| arrow::field("ts_milli", arrow::timestamp(arrow::TimeUnit::MILLI)), |
| arrow::field("ts_micro", arrow::timestamp(arrow::TimeUnit::MICRO)), |
| arrow::field("ts_nano", arrow::timestamp(arrow::TimeUnit::NANO)), |
| arrow::field("ts_tz_sec", arrow::timestamp(arrow::TimeUnit::MILLI, timezone_str)), |
| arrow::field("ts_tz_milli", arrow::timestamp(arrow::TimeUnit::MILLI, timezone_str)), |
| arrow::field("ts_tz_micro", arrow::timestamp(arrow::TimeUnit::MICRO, timezone_str)), |
| arrow::field("ts_tz_nano", arrow::timestamp(arrow::TimeUnit::NANO)), |
| }; |
| auto expected_file_schema = arrow::schema(expected_fields); |
| ASSERT_TRUE(result_file_schema->Equals(expected_file_schema)) |
| << result_file_schema->ToString(); |
| |
| // check array |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> result_array, |
| paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader.get())); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| } |
| } |
| |
| INSTANTIATE_TEST_SUITE_P(TestParam, ParquetFileBatchReaderTest, ::testing::Values(false, true)); |
| |
| TEST_F(ParquetFileBatchReaderTest, TestAddMetadataPerFieldMetadata) { |
| // Write a simple parquet file, call AddMetadata to inject per-field metadata |
| // before Finish, then read back and verify the file schema carries the metadata. |
| auto write_schema = arrow::schema({ |
| arrow::field("id", arrow::int32()), |
| arrow::field("name", arrow::utf8()), |
| arrow::field("score", arrow::float64()), |
| }); |
| |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), "update_schema_test.parquet"); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out, |
| fs_->Create(file_path, /*overwrite=*/true)); |
| |
| ::parquet::WriterProperties::Builder builder; |
| builder.write_batch_size(10); |
| 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, pool_)); |
| |
| // Write one batch of data. |
| auto data = arrow::ipc::internal::json::ArrayFromJSON( |
| arrow::struct_(write_schema->fields()), |
| R"([[1, "alice", 95.5], [2, "bob", 88.0], [3, "charlie", 72.3]])") |
| .ValueOrDie(); |
| ArrowArray c_array; |
| ASSERT_TRUE(arrow::ExportArray(*data, &c_array).ok()); |
| ASSERT_OK(format_writer->AddBatch(&c_array)); |
| ASSERT_OK(format_writer->Flush()); |
| |
| // Build an updated schema with per-field metadata on "name" and "score". |
| auto name_meta = std::make_shared<arrow::KeyValueMetadata>(); |
| name_meta->Append("shredding.field_mapping", "0:alice,1:bob,2:charlie"); |
| name_meta->Append("shredding.num_columns", "3"); |
| auto score_meta = std::make_shared<arrow::KeyValueMetadata>(); |
| score_meta->Append("custom.unit", "percent"); |
| |
| auto updated_schema = arrow::schema({ |
| write_schema->field(0), // id — no metadata |
| write_schema->field(1)->WithMetadata(name_meta), // name — shredding metadata |
| write_schema->field(2)->WithMetadata(score_meta), // score — custom metadata |
| }); |
| |
| // AddMetadata must be called before Finish. |
| ASSERT_OK(format_writer->AddMetadata( |
| {{ArrowUtils::kArrowSchemaMetadataKey, SerializeSchemaToString(updated_schema)}})); |
| ASSERT_OK(format_writer->Finish()); |
| ASSERT_OK(out->Flush()); |
| ASSERT_OK(out->Close()); |
| |
| // Read back: GetFileSchema should reflect the updated per-field metadata. |
| auto parquet_batch_reader = |
| PrepareParquetFileBatchReader(file_path, write_schema, /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt, batch_size_); |
| |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); |
| auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); |
| |
| // Field 0 "id": no metadata. |
| ASSERT_EQ("id", file_schema->field(0)->name()); |
| |
| // Field 1 "name": should have the shredding metadata we set. |
| ASSERT_EQ("name", file_schema->field(1)->name()); |
| auto read_name_meta = file_schema->field(1)->metadata(); |
| ASSERT_NE(nullptr, read_name_meta); |
| auto field_mapping_val = read_name_meta->Get("shredding.field_mapping").ValueOrDie(); |
| ASSERT_EQ("0:alice,1:bob,2:charlie", field_mapping_val); |
| auto num_columns_val = read_name_meta->Get("shredding.num_columns").ValueOrDie(); |
| ASSERT_EQ("3", num_columns_val); |
| |
| // Field 2 "score": should have the custom metadata. |
| ASSERT_EQ("score", file_schema->field(2)->name()); |
| auto read_score_meta = file_schema->field(2)->metadata(); |
| ASSERT_NE(nullptr, read_score_meta); |
| auto unit_val = read_score_meta->Get("custom.unit").ValueOrDie(); |
| ASSERT_EQ("percent", unit_val); |
| |
| // Also verify data integrity — read it back and compare content. |
| ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult( |
| parquet_batch_reader.get())); |
| ASSERT_EQ(result_array->num_chunks(), 1); |
| ASSERT_TRUE(data->Equals(*result_array->chunk(0))) << result_array->ToString(); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestRowMappingSimple) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto src_array = MakeSequentialIntData(12); |
| // data in file rowGroup0:[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] |
| // one row per page |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1, |
| /*enable_dictionary=*/true, /*max_row_group_length=*/12, /*max_page_size=*/1); |
| |
| // 1<=f0<=3 || 5<=f0<=6 |
| ASSERT_OK_AND_ASSIGN( |
| auto predicate, |
| PredicateBuilder::Or({PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(1), Literal(3)), |
| PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(5), Literal(6))})); |
| |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/2, |
| /*enable_page_level_filter=*/true); |
| |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch1, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| auto expected_batch1 = src_array->Slice(1, 2); |
| ASSERT_TRUE(batch1->chunk(0)->Equals(expected_batch1)) << batch1->ToString(); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 1); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(1).value(), 2); |
| // out of bound return invalid |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(2)); |
| |
| // Not adjacent pages |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch2, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| auto expected_batch2 = std::dynamic_pointer_cast<arrow::StructArray>( |
| arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ |
| [3], |
| [5] |
| ])") |
| .ValueOrDie()); |
| ASSERT_TRUE(batch2->chunk(0)->Equals(expected_batch2)) << batch2->ToString(); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 3); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(1).value(), 5); |
| |
| // Only one record read |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch3, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| auto expected_batch3 = src_array->Slice(6, 1); |
| ASSERT_TRUE(batch3->chunk(0)->Equals(expected_batch3)) << batch3->ToString(); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 6); |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(1)); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> eof_batch, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(nullptr, eof_batch); |
| // previous batch is eof, return invalid. |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestRowMappingFullyAndPartially) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto src_array = MakeSequentialIntData(12); |
| // data in file RowGroup0:[0, 1, 2] | RowGroup1:[3, 4, 5] | RowGroup2:[6, 7, 8] | RowGroup3:[9, |
| // 10, 11] one row per page |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1, |
| /*enable_dictionary=*/true, /*max_row_group_length=*/3, /*max_page_size=*/1); |
| |
| // 3<=f0<=5 || f0==6 || f0==8 |
| // RowGroup 1 is fully matched, RowGroup 2 is partially matched, RowGroup 0 and RowGroup 3 are |
| // not matched. |
| ASSERT_OK_AND_ASSIGN( |
| auto predicate, |
| PredicateBuilder::Or({PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(3), Literal(5)), |
| PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(6)), |
| PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(8))})); |
| |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3, |
| /*enable_page_level_filter=*/true); |
| |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch1, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 3); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(2).value(), 5); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch2, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 6); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(1).value(), 8); |
| } |
| |
| TEST_F(ParquetFileBatchReaderTest, TestRowMappingSetReadSchemaTwice) { |
| arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; |
| auto src_array = MakeSequentialIntData(12); |
| // data in file RowGroup0:[0, 1, 2] | RowGroup1:[3, 4, 5] | RowGroup2:[6, 7, 8] | RowGroup3:[9, |
| // 10, 11] one row per page |
| auto arrow_schema = arrow::schema(fields); |
| WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1, |
| /*enable_dictionary=*/true, /*max_row_group_length=*/3, /*max_page_size=*/1); |
| |
| // 1<=f0<=3 || 6<=f0<=7 |
| ASSERT_OK_AND_ASSIGN( |
| auto predicate, |
| PredicateBuilder::Or({PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(1), Literal(3)), |
| PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(6), Literal(7))})); |
| |
| auto parquet_batch_reader = PrepareParquetFileBatchReader( |
| file_path_, arrow_schema, /*predicate=*/predicate, std::nullopt, /*batch_size=*/3, |
| /*enable_page_level_filter=*/true); |
| |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch1, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 1); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(1).value(), 2); |
| |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch2, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 3); |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(1)); |
| |
| ASSERT_OK_AND_ASSIGN( |
| predicate, |
| PredicateBuilder::Or({PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"f0", |
| FieldType::INT, Literal(3), Literal(5))})); |
| |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*arrow_schema, c_schema.get()).ok()); |
| ASSERT_OK( |
| parquet_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/predicate, std::nullopt)); |
| ASSERT_NOK(parquet_batch_reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN( |
| std::shared_ptr<arrow::ChunkedArray> batch3, |
| paimon::test::ReadResultCollector::CollectResultOneBatch(parquet_batch_reader.get())); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(0).value(), 3); |
| ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(2).value(), 5); |
| } |
| |
| } // namespace paimon::parquet::test |
