| /* |
| * 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/avro/avro_file_batch_reader.h" |
| |
| #include <memory> |
| #include <string> |
| #include <utility> |
| |
| #include "arrow/api.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/ipc/api.h" |
| #include "gtest/gtest.h" |
| #include "paimon/common/utils/path_util.h" |
| #include "paimon/core/manifest/manifest_file.h" |
| #include "paimon/core/manifest/manifest_list.h" |
| #include "paimon/format/file_format.h" |
| #include "paimon/format/file_format_factory.h" |
| #include "paimon/fs/local/local_file_system.h" |
| #include "paimon/testing/utils/read_result_collector.h" |
| #include "paimon/testing/utils/testharness.h" |
| #include "paimon/testing/utils/timezone_guard.h" |
| |
| namespace paimon::avro::test { |
| |
| class AvroFileBatchReaderTest : public ::testing::Test, public ::testing::WithParamInterface<bool> { |
| public: |
| void SetUp() override { |
| ASSERT_OK_AND_ASSIGN(file_format_, |
| FileFormatFactory::Get("avro", {{Options::FILE_FORMAT, "avro"}})); |
| fs_ = std::make_shared<LocalFileSystem>(); |
| dir_ = ::paimon::test::UniqueTestDirectory::Create(); |
| ASSERT_TRUE(dir_); |
| pool_ = GetDefaultPool(); |
| } |
| void TearDown() override {} |
| |
| void WriteData(const std::shared_ptr<arrow::Array>& src_array, const std::string& file_path, |
| const std::string& compression) { |
| arrow::Schema src_schema(src_array->type()->fields()); |
| ::ArrowSchema c_schema; |
| ASSERT_TRUE(arrow::ExportSchema(src_schema, &c_schema).ok()); |
| ASSERT_OK_AND_ASSIGN(auto writer_builder, |
| file_format_->CreateWriterBuilder(&c_schema, /*batch_size=*/-1)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out, |
| fs_->Create(file_path, /*overwrite=*/false)); |
| ASSERT_OK_AND_ASSIGN(auto writer, writer_builder->Build(out, compression)); |
| |
| ::ArrowArray arrow_array; |
| ASSERT_TRUE(arrow::ExportArray(*src_array, &arrow_array).ok()); |
| ASSERT_OK(writer->AddBatch(&arrow_array)); |
| ASSERT_OK(writer->Flush()); |
| ASSERT_OK(writer->Finish()); |
| ASSERT_OK(out->Flush()); |
| ASSERT_OK(out->Close()); |
| } |
| |
| std::pair<std::unique_ptr<FileBatchReader>, std::shared_ptr<arrow::ChunkedArray>> ReadData( |
| const std::string& file_path, int32_t read_batch_size) { |
| EXPECT_OK_AND_ASSIGN(auto reader_builder, |
| file_format_->CreateReaderBuilder(read_batch_size)); |
| EXPECT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(file_path)); |
| EXPECT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); |
| EXPECT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( |
| batch_reader.get())); |
| return std::make_pair(std::move(batch_reader), result_array); |
| } |
| |
| private: |
| std::shared_ptr<MemoryPool> pool_; |
| std::shared_ptr<FileFormat> file_format_; |
| std::shared_ptr<FileSystem> fs_; |
| std::unique_ptr<paimon::test::UniqueTestDirectory> dir_; |
| }; |
| |
| TEST_F(AvroFileBatchReaderTest, TestReadDataWithNull) { |
| std::string path = paimon::test::GetDataDir() + "/avro/data/avro_with_null"; |
| auto [reader_holder, result_array] = ReadData(path, /*read_batch_size=*/1024); |
| |
| arrow::FieldVector fields = { |
| arrow::field("_KEY_f0", arrow::utf8(), /*nullable=*/true), |
| arrow::field("_SEQUENCE_NUMBER", arrow::int64(), /*nullable=*/true), |
| arrow::field("_VALUE_KIND", arrow::int32(), /*nullable=*/true), |
| arrow::field("f0", arrow::utf8(), /*nullable=*/true), |
| arrow::field("f1", arrow::utf8(), /*nullable=*/true), |
| arrow::field("f2", arrow::int32(), /*nullable=*/true), |
| arrow::field("f3", arrow::float64(), /*nullable=*/true)}; |
| |
| auto arrow_data_type = arrow::struct_(fields); |
| |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ |
| ["Alex", 2, 3, "Alex", "20250326", 18, 10.1], |
| ["Bob", 3, 3, "Bob", "20250326", 19, 11.1], |
| ["Evan", 1, 0, "Evan", "20250326", null, 14.1] |
| ])"}, |
| &expected_array); |
| ASSERT_TRUE(array_status.ok()) << array_status.ToString(); |
| ASSERT_TRUE(result_array->Equals(expected_array)); |
| ASSERT_TRUE(expected_array->Equals(result_array)); |
| auto read_metrics = reader_holder->GetReaderMetrics(); |
| ASSERT_TRUE(read_metrics); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestReadWithDifferentBatchSize) { |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), "file.avro"); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::boolean()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int64()), arrow::field("f3", arrow::float32()), |
| arrow::field("f4", arrow::float64()), arrow::field("f5", arrow::utf8()), |
| arrow::field("f6", arrow::binary())}; |
| auto arrow_data_type = arrow::struct_(fields); |
| |
| size_t length = 600; |
| std::string data_str = "["; |
| for (size_t i = 0; i < length; i++) { |
| if (i % 3 == 0) { |
| data_str.append(fmt::format(R"([{}, {}, {}, {}, {}, "str_{}", "bin_{}"])", "true", i, |
| i * 100000000000L, i * 0.12, i * 123.45678901, i, i)); |
| } else if (i % 3 == 1) { |
| data_str.append(fmt::format(R"([{}, -{}, -{}, -{}, -{}, "string_{}", "binary_{}"])", |
| "false", i, i * 100000000000L, i * 0.12, i * 123.45678901, |
| i, i)); |
| } else { |
| data_str.append("[null, null, null, null, null, null, null]"); |
| } |
| if (i != length - 1) { |
| data_str.append(","); |
| } |
| } |
| data_str.append("]"); |
| |
| std::shared_ptr<arrow::Array> src_array = |
| arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, data_str).ValueOrDie(); |
| ASSERT_TRUE(src_array); |
| WriteData(src_array, file_path, /*compression=*/"zstd"); |
| |
| for (int32_t batch_size : {1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1}) { |
| auto [reader_holder, result_array] = ReadData(file_path, batch_size); |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON( |
| arrow_data_type, {data_str}, &expected_array); |
| ASSERT_TRUE(array_status.ok()) << array_status.ToString(); |
| ASSERT_TRUE(result_array->Equals(expected_array)); |
| ASSERT_TRUE(expected_array->Equals(result_array)); |
| } |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestReadAllTypes) { |
| std::string path = paimon::test::GetDataDir() + "/avro/data/avro_all_types"; |
| auto [reader_holder, result_array] = ReadData(path, /*read_batch_size=*/1024); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::boolean()), |
| arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int32()), |
| arrow::field("f3", arrow::int32()), |
| arrow::field("f4", arrow::int64()), |
| arrow::field("f5", arrow::float32()), |
| arrow::field("f6", arrow::float64()), |
| arrow::field("f7", arrow::utf8()), |
| arrow::field("f8", arrow::binary()), |
| arrow::field("f10", arrow::list(arrow::float32())), |
| arrow::field("f11", arrow::struct_({arrow::field("f0", arrow::boolean()), |
| arrow::field("f1", arrow::int64())})), |
| arrow::field("f12", arrow::timestamp(arrow::TimeUnit::MICRO)), |
| arrow::field("f13", arrow::date32()), |
| arrow::field("f14", arrow::decimal128(2, 2)), |
| arrow::field("f15", arrow::decimal128(10, 10)), |
| arrow::field("f16", arrow::decimal128(19, 19))}; |
| |
| auto arrow_data_type = arrow::struct_(fields); |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ |
| [true, 127, 32767, 2147483647, 9999999999999, 1234.56, 1234567890.0987654321, "aa", "qq", [0.1, 0.2], [true, null], "1970-01-01 00:02:03.123123", 2456, "0.22", "0.1234567890", "0.1234567890987654321"], |
| [false, -128, -32768, -2147483648, -9999999999999, -1234.56, -1234567890.0987654321, null, "ww", [-0.1, -0.2, null, 0.3, 0.4], [null, 2], "1970-01-01 00:16:39.999999", null, "-0.22", "-0.1234567890", null], |
| [null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null] |
| ])"}, |
| &expected_array); |
| ASSERT_TRUE(array_status.ok()) << array_status.ToString(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| ASSERT_TRUE(expected_array->Equals(result_array)) << result_array->ToString(); |
| } |
| |
| TEST_P(AvroFileBatchReaderTest, TestReadTimestampTypes) { |
| auto enable_tz = GetParam(); |
| std::string timezone_str = enable_tz ? "Asia/Tokyo" : "Asia/Shanghai"; |
| paimon::test::TimezoneGuard tz_guard(timezone_str); |
| |
| std::string path = paimon::test::GetDataDir() + |
| "/avro/append_with_multiple_ts_precision_and_timezone.db/" |
| "append_with_multiple_ts_precision_and_timezone/bucket-0/" |
| "data-441e233b-529d-4a8f-a0a4-25c2c84fb965-0.avro"; |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, |
| file_format_->CreateReaderBuilder(/*batch_size=*/1024)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(path)); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); |
| |
| auto timezone = DateTimeUtils::GetLocalTimezoneName(); |
| arrow::FieldVector read_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_tz_sec", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)), |
| arrow::field("ts_tz_milli", arrow::timestamp(arrow::TimeUnit::MILLI, timezone)), |
| arrow::field("ts_tz_micro", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)), |
| }; |
| auto read_schema = arrow::schema(read_fields); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| EXPECT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| |
| // check array |
| ASSERT_OK_AND_ASSIGN(auto result_array, |
| ::paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| auto array_status = |
| arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(read_fields), {R"([ |
| ["1970-01-01T00:00:01","1970-01-01T00:00:00.001","1970-01-01T00:00:00.000001","1970-01-01T00:00:02","1970-01-01T00:00:00.002","1970-01-01T00:00:00.000002"], |
| [null,"1970-01-01T00:00:00.003",null,null,"1970-01-01T00:00:00.004",null], |
| ["1970-01-01T00:00:05",null,"1970-01-01T00:00:00.000005","1970-01-01T00:00:06",null,"1970-01-01T00:00:00.000006"] |
| ])"}, |
| &expected_array); |
| ASSERT_TRUE(array_status.ok()) << array_status.ToString(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString(); |
| ASSERT_TRUE(expected_array->Equals(result_array)); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestReadMapTypes) { |
| std::string path = paimon::test::GetDataDir() + |
| "/avro/append_with_multiple_map.db/" |
| "append_with_multiple_map/bucket-0/" |
| "data-72442742-e49e-48a4-a736-a2475aac2d2c-0.avro"; |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, |
| file_format_->CreateReaderBuilder(/*batch_size=*/1024)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(path)); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); |
| |
| arrow::FieldVector read_fields = { |
| arrow::field("f0", arrow::map(arrow::int32(), arrow::int32())), |
| arrow::field("f1", arrow::map(arrow::float64(), arrow::float64())), |
| arrow::field("f2", arrow::map(arrow::utf8(), arrow::utf8())), |
| arrow::field("f3", arrow::map(arrow::utf8(), arrow::binary())), |
| arrow::field("f4", arrow::map(arrow::timestamp(arrow::TimeUnit::MICRO), |
| arrow::timestamp(arrow::TimeUnit::MICRO))), |
| arrow::field("f5", arrow::map(arrow::utf8(), arrow::list(arrow::float64()))), |
| arrow::field("f6", arrow::map(arrow::utf8(), arrow::map(arrow::float64(), arrow::utf8()))), |
| arrow::field("f7", arrow::map(arrow::int64(), |
| arrow::struct_({field("f0", arrow::int32()), |
| field("f1", arrow::utf8()), |
| field("f2", arrow::decimal128(5, 2))})))}; |
| auto read_schema = arrow::schema(read_fields); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| EXPECT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| |
| // check array |
| ASSERT_OK_AND_ASSIGN(auto result_array, |
| ::paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); |
| std::shared_ptr<arrow::ChunkedArray> expected_array; |
| auto array_status = |
| arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(read_fields), {R"([ |
| [ |
| [[1,10],[2,20]], |
| [[1.1,10.1],[2.2,20.2]], |
| [["key1","val1"],["key2","val2"]], |
| [["123456","abcdef"]], |
| [["2023-01-01 12:00:00.123000","2023-01-01 12:00:00.123000"],["2023-01-02 13:30:00.456000","2023-01-02 13:30:00.456000"]], |
| [["arr_key",[1.5, 2.5, 3.5]]], |
| [["outer_key",[[99.9,"nested_val"]]]], |
| [[1000, [42, "row_str", "123.45"]]] |
| ] |
| ])"}, |
| &expected_array); |
| ASSERT_TRUE(array_status.ok()) << array_status.ToString(); |
| ASSERT_TRUE(result_array->Equals(expected_array)) << result_array->ToString() << std::endl; |
| ASSERT_TRUE(expected_array->Equals(result_array)); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestSetReadSchemaRejectNestedSubFieldProjection) { |
| std::string path = PathUtil::JoinPath(dir_->Str(), "nested_projection_unsupported.avro"); |
| |
| arrow::FieldVector write_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({arrow::field("a", arrow::int32()), |
| arrow::field("b", arrow::utf8())}))}; |
| auto write_type = arrow::struct_(write_fields); |
| auto write_array = arrow::ipc::internal::json::ArrayFromJSON(write_type, R"([ |
| [1, [10, "x"]], |
| [2, [20, "y"]] |
| ])") |
| .ValueOrDie(); |
| WriteData(write_array, path, /*compression=*/"null"); |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, |
| file_format_->CreateReaderBuilder(/*batch_size=*/1024)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(path)); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); |
| |
| arrow::FieldVector read_fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::struct_({arrow::field("a", arrow::int32())}))}; |
| auto read_schema = arrow::schema(read_fields); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| |
| ASSERT_NOK_WITH_MSG(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt), |
| "does not support nested sub-field projection"); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestGetPreviousBatchFileRowId) { |
| std::string path = paimon::test::GetDataDir() + |
| "/avro/append_simple.db/" |
| "append_simple/bucket-0/" |
| "data-d7d1c416-6e34-4834-af87-341d09418f0c-0.avro"; |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, file_format_->CreateReaderBuilder(/*batch_size=*/1)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(path)); |
| ASSERT_OK_AND_ASSIGN(auto reader, reader_builder->Build(in)); |
| |
| arrow::FieldVector read_fields = { |
| arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::float64()), |
| arrow::field("f2", arrow::utf8()), |
| arrow::field("f3", |
| arrow::struct_({arrow::field("f0", arrow::map(arrow::utf8(), arrow::int32())), |
| arrow::field("f1", arrow::list(arrow::int32()))}))}; |
| |
| auto read_schema = arrow::schema(read_fields); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| EXPECT_OK(reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto num_rows, reader->GetNumberOfRows()); |
| ASSERT_EQ(4, num_rows); |
| ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_OK_AND_ASSIGN(auto batch1, reader->NextBatch()); |
| ArrowArrayRelease(batch1.first.get()); |
| ArrowSchemaRelease(batch1.second.get()); |
| ASSERT_EQ(0, reader->GetPreviousBatchFileRowId(0).value()); |
| ASSERT_OK_AND_ASSIGN(auto batch2, reader->NextBatch()); |
| ASSERT_EQ(1, reader->GetPreviousBatchFileRowId(0).value()); |
| ArrowArrayRelease(batch2.first.get()); |
| ArrowSchemaRelease(batch2.second.get()); |
| ASSERT_OK_AND_ASSIGN(auto batch3, reader->NextBatch()); |
| ASSERT_EQ(2, reader->GetPreviousBatchFileRowId(0).value()); |
| ArrowArrayRelease(batch3.first.get()); |
| ArrowSchemaRelease(batch3.second.get()); |
| ASSERT_OK_AND_ASSIGN(auto batch4, reader->NextBatch()); |
| ASSERT_EQ(3, reader->GetPreviousBatchFileRowId(0).value()); |
| ArrowArrayRelease(batch4.first.get()); |
| ArrowSchemaRelease(batch4.second.get()); |
| ASSERT_OK_AND_ASSIGN(auto batch5, reader->NextBatch()); |
| ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); |
| ASSERT_TRUE(BatchReader::IsEofBatch(batch5)); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestSetReadSchemaResetsReaderToFirstRow) { |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), "file.avro"); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::int32()), |
| arrow::field("f1", arrow::int32()), |
| }; |
| auto file_data_type = arrow::struct_(fields); |
| auto src_array = arrow::ipc::internal::json::ArrayFromJSON(file_data_type, R"([ |
| [1, 10], |
| [2, 20], |
| [3, 30], |
| [4, 40] |
| ])") |
| .ValueOrDie(); |
| WriteData(src_array, file_path, /*compression=*/"null"); |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, file_format_->CreateReaderBuilder(/*batch_size=*/2)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(file_path)); |
| ASSERT_OK_AND_ASSIGN(auto reader, reader_builder->Build(in)); |
| |
| ASSERT_OK_AND_ASSIGN(auto first_batch, reader->NextBatch()); |
| ASSERT_EQ(0, reader->GetPreviousBatchFileRowId(0).value()); |
| auto first_array = |
| arrow::ImportArray(first_batch.first.get(), first_batch.second.get()).ValueOrDie(); |
| ASSERT_TRUE(first_array->Equals(src_array->Slice(0, 2))) << first_array->ToString(); |
| |
| auto read_schema = arrow::schema({arrow::field("f1", arrow::int32())}); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| ASSERT_OK(reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); |
| |
| ASSERT_OK_AND_ASSIGN(auto projected_batch, reader->NextBatch()); |
| ASSERT_EQ(0, reader->GetPreviousBatchFileRowId(0).value()); |
| auto projected_array = |
| arrow::ImportArray(projected_batch.first.get(), projected_batch.second.get()).ValueOrDie(); |
| auto expected_projected_array = arrow::ipc::internal::json::ArrayFromJSON( |
| arrow::struct_({arrow::field("f1", arrow::int32())}), |
| R"([ |
| [10], |
| [20] |
| ])") |
| .ValueOrDie(); |
| ASSERT_TRUE(projected_array->Equals(expected_projected_array)) << projected_array->ToString(); |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, TestGetNumberOfRows) { |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), "file.avro"); |
| |
| arrow::FieldVector fields = { |
| arrow::field("f0", arrow::boolean()), arrow::field("f1", arrow::int32()), |
| arrow::field("f2", arrow::int64()), arrow::field("f3", arrow::float32()), |
| arrow::field("f4", arrow::float64()), arrow::field("f5", arrow::utf8()), |
| arrow::field("f6", arrow::binary())}; |
| auto arrow_data_type = arrow::struct_(fields); |
| |
| size_t length = 102400; |
| std::string data_str = "["; |
| for (size_t i = 0; i < length; i++) { |
| data_str.append(fmt::format(R"([{}, {}, {}, {}, {}, "str_{}", "bin_{}"])", "true", i, |
| i * 100000000000L, i * 0.12, i * 123.45678901, i, i)); |
| if (i != length - 1) { |
| data_str.append(","); |
| } |
| } |
| data_str.append("]"); |
| |
| std::shared_ptr<arrow::Array> src_array = |
| arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, data_str).ValueOrDie(); |
| ASSERT_TRUE(src_array); |
| WriteData(src_array, file_path, /*compression=*/"null"); |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, file_format_->CreateReaderBuilder(25600)); |
| |
| // check GetNumberOfRows can be called at any position, and continue read |
| int32_t expected_batches = 4; |
| for (int32_t pos = 0; pos < expected_batches; pos++) { |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(file_path)); |
| ASSERT_OK_AND_ASSIGN(auto reader, reader_builder->Build(in)); |
| |
| arrow::ArrayVector result_array_vector; |
| for (int32_t i = 0; i < pos; i++) { |
| ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatch()); |
| auto result_array = |
| arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); |
| result_array_vector.push_back(result_array); |
| } |
| // check number of rows, and continue read |
| ASSERT_OK_AND_ASSIGN(auto num_rows, reader->GetNumberOfRows()); |
| ASSERT_EQ(length, num_rows); |
| for (int32_t i = pos; i < expected_batches; i++) { |
| ASSERT_OK_AND_ASSIGN(auto batch, reader->NextBatch()); |
| auto result_array = |
| arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); |
| result_array_vector.push_back(result_array); |
| } |
| ASSERT_OK_AND_ASSIGN(auto eof_batch, reader->NextBatch()); |
| ASSERT_TRUE(BatchReader::IsEofBatch(eof_batch)); |
| ASSERT_OK_AND_ASSIGN(num_rows, reader->GetNumberOfRows()); |
| ASSERT_EQ(length, num_rows); |
| |
| auto result_array = arrow::ChunkedArray(result_array_vector); |
| ASSERT_TRUE(result_array.Equals(arrow::ChunkedArray(src_array))); |
| } |
| } |
| |
| TEST_F(AvroFileBatchReaderTest, 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_data_type = arrow::struct_({write_field}); |
| auto write_array = |
| arrow::ipc::internal::json::ArrayFromJSON(write_data_type, data_json).ValueOrDie(); |
| |
| std::string file_path = PathUtil::JoinPath(dir_->Str(), file_name); |
| WriteData(write_array, file_path, /*compression=*/"null"); |
| |
| // Read back with binary schema |
| auto read_field = arrow::field("f0", arrow::binary()); |
| auto read_data_type = arrow::struct_({read_field}); |
| |
| ASSERT_OK_AND_ASSIGN(auto reader_builder, |
| file_format_->CreateReaderBuilder(/*batch_size=*/1024)); |
| ASSERT_OK_AND_ASSIGN(std::shared_ptr<InputStream> in, fs_->Open(file_path)); |
| ASSERT_OK_AND_ASSIGN(auto batch_reader, reader_builder->Build(in)); |
| |
| // Check GetFileSchema: regardless of write type, avro file schema is always binary |
| ASSERT_OK_AND_ASSIGN(auto c_file_schema, batch_reader->GetFileSchema()); |
| auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); |
| arrow::Schema expected_file_schema({read_field}); |
| ASSERT_TRUE(file_schema->Equals(expected_file_schema)); |
| |
| auto read_schema = arrow::schema({read_field}); |
| std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); |
| EXPECT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, |
| /*selection_bitmap=*/std::nullopt)); |
| |
| ASSERT_OK_AND_ASSIGN(auto result_array, ::paimon::test::ReadResultCollector::CollectResult( |
| batch_reader.get())); |
| auto expected_array = |
| arrow::ipc::internal::json::ArrayFromJSON(read_data_type, data_json).ValueOrDie(); |
| auto expected_chunked_array = std::make_shared<arrow::ChunkedArray>(expected_array); |
| ASSERT_TRUE(result_array->Equals(expected_chunked_array)); |
| }; |
| |
| check_binary_read_result(arrow::binary(), "binary.avro"); |
| check_binary_read_result(arrow::large_binary(), "large-binary.avro"); |
| } |
| |
| INSTANTIATE_TEST_SUITE_P(TestParam, AvroFileBatchReaderTest, ::testing::Values(false, true)); |
| |
| } // namespace paimon::avro::test |
