/*
 * 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/common/reader/data_evolution_file_reader.h"

#include <map>

#include "arrow/api.h"
#include "arrow/array/array_base.h"
#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "arrow/ipc/api.h"
#include "arrow/util/range.h"
#include "gtest/gtest.h"
#include "paimon/common/utils/arrow/mem_utils.h"
#include "paimon/memory/memory_pool.h"
#include "paimon/testing/mock/mock_file_batch_reader.h"
#include "paimon/testing/utils/read_result_collector.h"
#include "paimon/testing/utils/testharness.h"

namespace paimon::test {
class DataEvolutionFileReaderTest : public ::testing::Test,
                                    public ::testing::WithParamInterface<bool> {
 public:
    void SetUp() override {
        pool_ = GetDefaultPool();
    }

    void TearDown() override {
        pool_.reset();
    }

    void CheckResult(const arrow::ArrayVector& src_array_vec,
                     const std::shared_ptr<arrow::Schema>& read_schema,
                     const std::vector<int32_t>& reader_offsets,
                     const std::vector<int32_t>& field_offsets,
                     const std::shared_ptr<arrow::Array>& expected_array,
                     const std::optional<RoaringBitmap32>& selection_bitmap = std::nullopt) const {
        for (auto batch_size : arrow::internal::Iota(1, 10)) {
            int32_t total_row_count = 0;
            std::vector<std::unique_ptr<BatchReader>> readers;
            for (const auto& array : src_array_vec) {
                if (array == nullptr) {
                    // simulate no fields read from current reader
                    readers.push_back(nullptr);
                    continue;
                }
                total_row_count += array->length();
                std::unique_ptr<MockFileBatchReader> file_batch_reader;
                if (selection_bitmap) {
                    file_batch_reader = std::make_unique<MockFileBatchReader>(
                        array, array->type(), selection_bitmap.value(), batch_size);
                } else {
                    file_batch_reader =
                        std::make_unique<MockFileBatchReader>(array, array->type(), batch_size);
                }
                auto enable_randomize_batch_size = GetParam();
                file_batch_reader->EnableRandomizeBatchSize(enable_randomize_batch_size);
                readers.push_back(std::move(file_batch_reader));
            }
            ASSERT_OK_AND_ASSIGN(
                auto data_evolution_file_reader,
                DataEvolutionFileReader::Create(std::move(readers), read_schema, batch_size,
                                                reader_offsets, field_offsets, pool_));
            // check metrics, data_evolution_file_reader collects all row of each
            // MockFileBatchReader
            auto metrics = data_evolution_file_reader->GetReaderMetrics();
            ASSERT_EQ(metrics->ToString(),
                      "{\"mock.number.of.rows\":" + std::to_string(total_row_count) + "}");

            // check result array
            ASSERT_OK_AND_ASSIGN(
                auto result_array,
                paimon::test::ReadResultCollector::CollectResult(data_evolution_file_reader.get()));
            data_evolution_file_reader->Close();
            auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_array);
            ASSERT_TRUE(result_array->Equals(expected_chunk_array));
        }
    }

    void CheckNextBatchForSingleReader(int32_t inner_batch_size, int32_t read_batch_size,
                                       const std::shared_ptr<arrow::Array>& src_array,
                                       const std::optional<RoaringBitmap32>& selection_bitmap,
                                       const std::shared_ptr<arrow::Array>& expected_array) const {
        std::unique_ptr<MockFileBatchReader> file_batch_reader;
        if (selection_bitmap) {
            file_batch_reader = std::make_unique<MockFileBatchReader>(
                src_array, src_array->type(), selection_bitmap.value(), inner_batch_size);
        } else {
            file_batch_reader = std::make_unique<MockFileBatchReader>(src_array, src_array->type(),
                                                                      inner_batch_size);
        }
        auto enable_randomize_batch_size = GetParam();
        file_batch_reader->EnableRandomizeBatchSize(enable_randomize_batch_size);
        std::vector<std::unique_ptr<BatchReader>> readers;
        readers.push_back(std::move(file_batch_reader));
        DataEvolutionFileReader fake_data_evolution_reader(
            std::move(readers), /*read_schema=*/arrow::schema({}), read_batch_size,
            /*reader_offsets=*/{}, /*field_offsets=*/{}, GetArrowPool(pool_));
        arrow::ArrayVector result_array_vec;
        while (true) {
            ASSERT_OK_AND_ASSIGN(auto result_array,
                                 fake_data_evolution_reader.NextBatchForSingleReader(0));
            if (result_array == nullptr) {
                break;
            }
            ASSERT_EQ(result_array->offset(), 0);
            result_array_vec.push_back(result_array);
        }
        ASSERT_EQ(result_array_vec.size(),
                  std::ceil(static_cast<double>(expected_array->length()) / read_batch_size));
        // except for last batch, the length each array is expected to be aligned to read_batch_size
        for (size_t i = 0; i < result_array_vec.size() - 1; i++) {
            ASSERT_EQ(result_array_vec[i]->length(), read_batch_size);
        }
        if (expected_array->length() % read_batch_size == 0) {
            ASSERT_EQ(result_array_vec.back()->length(), read_batch_size);
        } else {
            ASSERT_EQ(result_array_vec.back()->length(),
                      expected_array->length() % read_batch_size);
        }
        auto result_chunk_array = std::make_shared<arrow::ChunkedArray>(result_array_vec);
        auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_array);
        ASSERT_TRUE(result_chunk_array->Equals(expected_chunk_array));
    }

 private:
    std::shared_ptr<MemoryPool> pool_;
};

TEST_F(DataEvolutionFileReaderTest, TestInvalid) {
    {
        arrow::FieldVector read_fields;
        auto read_schema = arrow::schema(read_fields);
        ASSERT_NOK_WITH_MSG(
            DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10, {}, {}, pool_),
            "read schema must not be empty");
    }
    {
        arrow::FieldVector read_fields = {
            arrow::field("f0", arrow::int32()),
            arrow::field("f1", arrow::int32()),
            arrow::field("f2", arrow::utf8()),
            arrow::field("f3", arrow::int32()),
        };
        auto read_schema = arrow::schema(read_fields);
        std::vector<int32_t> reader_offsets = {0, 0, 1};
        std::vector<int32_t> field_offsets = {0, 1, 0};
        ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10,
                                                            reader_offsets, field_offsets, pool_),
                            "read schema, row offsets and field offsets must have the same size");
    }
    {
        arrow::FieldVector read_fields = {
            arrow::field("f0", arrow::int32()),
            arrow::field("f1", arrow::int32()),
            arrow::field("f2", arrow::utf8()),
            arrow::field("f3", arrow::int32()),
        };
        auto read_schema = arrow::schema(read_fields);
        std::vector<int32_t> reader_offsets = {0, 0, 1, 1};
        std::vector<int32_t> field_offsets = {0, 1, 1, 0};
        ASSERT_NOK_WITH_MSG(DataEvolutionFileReader::Create({}, read_schema, /*read_batch_size=*/10,
                                                            reader_offsets, field_offsets, pool_),
                            "readers must not be empty");
    }
    {
        std::vector<std::unique_ptr<BatchReader>> readers;
        readers.push_back(nullptr);

        arrow::FieldVector read_fields = {
            arrow::field("f0", arrow::int32()),
            arrow::field("f1", arrow::int32()),
            arrow::field("f2", arrow::utf8()),
            arrow::field("f3", arrow::int32()),
        };
        auto read_schema = arrow::schema(read_fields);
        std::vector<int32_t> reader_offsets = {0, 0, 1, 1};
        std::vector<int32_t> field_offsets = {0, 1, 1, 0};
        ASSERT_NOK_WITH_MSG(
            DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10,
                                            reader_offsets, field_offsets, pool_),
            "reader offset is out of range of readers");
    }
}

TEST_P(DataEvolutionFileReaderTest, TestNextBatchForSingleReader) {
    auto prepare_array = [](int64_t array_length) -> std::shared_ptr<arrow::Array> {
        auto array_builder = std::make_shared<arrow::Int32Builder>();
        for (int32_t i = 0; i < array_length; ++i) {
            EXPECT_TRUE(array_builder->Append(i).ok());
        }
        std::shared_ptr<arrow::Array> array;
        EXPECT_TRUE(array_builder->Finish(&array).ok());
        return array;
    };
    auto prepare_array_with_bitmap =
        [](const RoaringBitmap32& bitmap) -> std::shared_ptr<arrow::Array> {
        auto array_builder = std::make_shared<arrow::Int32Builder>();
        for (auto iter = bitmap.Begin(); iter != bitmap.End(); ++iter) {
            EXPECT_TRUE(array_builder->Append(*iter).ok());
        }
        std::shared_ptr<arrow::Array> array;
        EXPECT_TRUE(array_builder->Finish(&array).ok());
        return array;
    };
    {
        // src array length = 10, read batch size = 10
        auto src_array = prepare_array(10);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 10)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/10, src_array,
                                          /*selection_bitmap=*/std::nullopt,
                                          /*expected_array=*/src_array);
        }
    }
    {
        // src array length = 10, read batch size = 6
        auto src_array = prepare_array(10);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 6)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/6, src_array,
                                          /*selection_bitmap=*/std::nullopt,
                                          /*expected_array=*/src_array);
        }
    }
    {
        // src array length = 10, read batch size = 15
        auto src_array = prepare_array(10);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          /*selection_bitmap=*/std::nullopt,
                                          /*expected_array=*/src_array);
        }
    }
    {
        // test bulk data, src array length = 10000, read batch size = 1024
        auto src_array = prepare_array(10000);
        for (int32_t inner_batch_size : {1, 2, 8, 16, 20, 100, 1024}) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/1024, src_array,
                                          /*selection_bitmap=*/std::nullopt,
                                          /*expected_array=*/src_array);
        }
    }
    {
        // src array length = 10, selection bitmap = {1, 3, 5}
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({1, 3, 5});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          selected_bitmap, expected_array);
        }
    }
    {
        // src array length = 10, selection all
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          selected_bitmap, /*expected_array=*/src_array);
        }
    }
    {
        // src array length = 10, selection first
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          selected_bitmap, expected_array);
        }
    }
    {
        // src array length = 10, selection first
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({9});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          selected_bitmap, expected_array);
        }
    }
    {
        // src array length = 10, selection consecutive positions
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({2, 3, 4, 5});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/15, src_array,
                                          selected_bitmap, expected_array);
        }
    }
    {
        auto src_array = prepare_array(10);
        RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0, 1, 2, 3, 4, 5, 7, 8, 9});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        // inner batch: [0, 1, 2, 3] | [4, 5] [7, 8] | [9]
        CheckNextBatchForSingleReader(/*inner_batch_size=*/4, /*read_batch_size=*/5, src_array,
                                      selected_bitmap, expected_array);
    }
    {
        // test bulk data, src array length = 10000, read batch size = 1024
        auto src_array = prepare_array(10000);
        RoaringBitmap32 selected_bitmap =
            RoaringBitmap32::From({0, 10, 1000, 2333, 4566, 7838, 8787, 9999});
        auto expected_array = prepare_array_with_bitmap(selected_bitmap);
        for (int32_t inner_batch_size : {1, 2, 8, 16, 20, 100, 1024}) {
            CheckNextBatchForSingleReader(inner_batch_size, /*read_batch_size=*/1024, src_array,
                                          selected_bitmap, expected_array);
        }
    }
}

TEST_P(DataEvolutionFileReaderTest, TestSimple) {
    arrow::FieldVector read_fields = {
        arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32()),
        arrow::field("f2", arrow::utf8()),  arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::utf8()),  arrow::field("f5", arrow::int32()),
    };
    auto read_schema = arrow::schema(read_fields);

    std::vector<int32_t> reader_offsets = {0, 2, 0, 1, 2, 1};
    std::vector<int32_t> field_offsets = {0, 0, 1, 1, 1, 0};

    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[0], read_fields[2]}), R"([
        [0, "00"],
        [1, "01"],
        [2, "02"],
        [3, "03"],
        [4, "04"],
        [5, "05"]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[5], read_fields[3]}), R"([
        [10, 110],
        [11, 111],
        [12, 112],
        [13, 113],
        [14, 114],
        [15, 115]
])")
                      .ValueOrDie();
    auto array2 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[4]}), R"([
        [20, "20"],
        [21, "21"],
        [22, "22"],
        [23, "23"],
        [24, "24"],
        [25, "25"]
])")
                      .ValueOrDie();

    auto expected_array =
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_fields), R"([
        [0, 20, "00", 110, "20", 10],
        [1, 21, "01", 111, "21", 11],
        [2, 22, "02", 112, "22", 12],
        [3, 23, "03", 113, "23", 13],
        [4, 24, "04", 114, "24", 14],
        [5, 25, "05", 115, "25", 15]
])")
            .ValueOrDie();
    CheckResult({array0, array1, array2}, read_schema, reader_offsets, field_offsets,
                expected_array);
}

TEST_P(DataEvolutionFileReaderTest, TestWithNonExistField) {
    arrow::FieldVector read_fields = {
        arrow::field("f0", arrow::int32()),        arrow::field("f1", arrow::int32()),
        arrow::field("f2", arrow::utf8()),         arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::utf8()),         arrow::field("f5", arrow::int32()),
        arrow::field("non-field", arrow::int32()),
    };
    auto read_schema = arrow::schema(read_fields);

    std::vector<int32_t> reader_offsets = {0, 2, 0, 1, 2, 1, -1};
    std::vector<int32_t> field_offsets = {0, 0, 1, 1, 1, 0, -1};

    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[0], read_fields[2]}), R"([
        [0, "00"],
        [1, "01"],
        [2, "02"],
        [3, "03"],
        [4, "04"],
        [5, "05"]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[5], read_fields[3]}), R"([
        [10, 110],
        [11, 111],
        [12, 112],
        [13, 113],
        [14, 114],
        [15, 115]
])")
                      .ValueOrDie();
    auto array2 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[4]}), R"([
        [20, "20"],
        [21, "21"],
        [22, "22"],
        [23, "23"],
        [24, "24"],
        [25, "25"]
])")
                      .ValueOrDie();

    auto expected_array =
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_fields), R"([
        [0, 20, "00", 110, "20", 10, null],
        [1, 21, "01", 111, "21", 11, null],
        [2, 22, "02", 112, "22", 12, null],
        [3, 23, "03", 113, "23", 13, null],
        [4, 24, "04", 114, "24", 14, null],
        [5, 25, "05", 115, "25", 15, null]
])")
            .ValueOrDie();
    CheckResult({array0, array1, array2}, read_schema, reader_offsets, field_offsets,
                expected_array);
}

TEST_P(DataEvolutionFileReaderTest, TestReadFromPartialReaders) {
    arrow::FieldVector read_fields = {
        arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32()),
        arrow::field("f2", arrow::utf8()),  arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::utf8()),  arrow::field("f5", arrow::int32()),
    };
    auto read_schema = arrow::schema(read_fields);
    // simulate reader2 has no field to read
    std::vector<int32_t> reader_offsets = {0, 3, 0, 1, 3, 1};
    std::vector<int32_t> field_offsets = {0, 0, 1, 1, 1, 0};

    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[0], read_fields[2]}), R"([
        [0, "00"],
        [1, "01"],
        [2, "02"],
        [3, "03"],
        [4, "04"],
        [5, "05"]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[5], read_fields[3]}), R"([
        [10, 110],
        [11, 111],
        [12, 112],
        [13, 113],
        [14, 114],
        [15, 115]
])")
                      .ValueOrDie();
    auto array3 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[4]}), R"([
        [20, "20"],
        [21, "21"],
        [22, "22"],
        [23, "23"],
        [24, "24"],
        [25, "25"]
])")
                      .ValueOrDie();

    auto expected_array =
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_fields), R"([
        [0, 20, "00", 110, "20", 10],
        [1, 21, "01", 111, "21", 11],
        [2, 22, "02", 112, "22", 12],
        [3, 23, "03", 113, "23", 13],
        [4, 24, "04", 114, "24", 14],
        [5, 25, "05", 115, "25", 15]
])")
            .ValueOrDie();

    CheckResult({array0, array1, nullptr, array3}, read_schema, reader_offsets, field_offsets,
                expected_array);
}

TEST_P(DataEvolutionFileReaderTest, TestNestedType) {
    arrow::FieldVector read_fields = {
        arrow::field("f1", arrow::map(arrow::int8(), arrow::int16())),
        arrow::field("f2", arrow::list(arrow::float32())),
        arrow::field("f3", arrow::struct_({arrow::field("f0", arrow::boolean()),
                                           arrow::field("f1", arrow::int64())})),
        arrow::field("f4", arrow::timestamp(arrow::TimeUnit::NANO)),
        arrow::field("f5", arrow::date32()),
        arrow::field("f6", arrow::decimal128(2, 2))};
    auto read_schema = arrow::schema(read_fields);

    std::vector<int32_t> reader_offsets = {0, 1, 0, 1, 0, 1};
    std::vector<int32_t> field_offsets = {2, 0, 1, 1, 0, 2};

    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[4], read_fields[2], read_fields[0]}), R"([
        [2456, [true, 2], [[0, 0]]],
        [24, [true, 1], [[0, 1]]],
        [2456, [false, 12], [[10, 10]]],
        [245, [false, 2222], [[127, 32767], [-128, -32768]]],
        [24, [true, 2], [[1, 64], [2, 32]]],
        [24, [true, 2], [[11, 64], [12, 32]]]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[3], read_fields[5]}), R"([
        [[0.1, 0.2], "1970-01-01 00:02:03.123123", "0.22"],
        [[0.1, 0.3], "1970-01-01 00:02:03.999999", "0.28"],
        [[1.1, 1.2], "1970-01-01 00:02:03.123123", "0.22"],
        [[1.1, 1.2], "1970-01-01 00:02:03.123123", "0.12"],
        [[2.2, 3.2], "1970-01-01 00:00:00.0", "0.78"],
        [[2.2, 3.2], "1970-01-01 00:00:00.123123", "0.78"]
])")
                      .ValueOrDie();

    auto expected_array =
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_fields), R"([
        [[[0, 0]], [0.1, 0.2], [true, 2], "1970-01-01 00:02:03.123123", 2456, "0.22"],
        [[[0, 1]], [0.1, 0.3], [true, 1], "1970-01-01 00:02:03.999999", 24, "0.28"],
        [[[10, 10]], [1.1, 1.2], [false, 12], "1970-01-01 00:02:03.123123", 2456, "0.22"],
        [[[127, 32767], [-128, -32768]], [1.1, 1.2], [false, 2222], "1970-01-01 00:02:03.123123", 245, "0.12"],
        [[[1, 64], [2, 32]], [2.2, 3.2], [true, 2], "1970-01-01 00:00:00.0", 24, "0.78"],
        [[[11, 64], [12, 32]], [2.2, 3.2], [true, 2], "1970-01-01 00:00:00.123123", 24, "0.78"]
])")
            .ValueOrDie();
    CheckResult({array0, array1}, read_schema, reader_offsets, field_offsets, expected_array);
}

TEST_P(DataEvolutionFileReaderTest, TestWithBitmap) {
    arrow::FieldVector read_fields = {
        arrow::field("f0", arrow::int32()),       arrow::field("f1", arrow::int32()),
        arrow::field("f2", arrow::utf8()),        arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::utf8()),        arrow::field("f5", arrow::int32()),
        arrow::field("non-exist", arrow::int32())};
    auto read_schema = arrow::schema(read_fields);

    std::vector<int32_t> reader_offsets = {0, 2, 0, 1, 2, 1, -1};
    std::vector<int32_t> field_offsets = {0, 0, 1, 1, 1, 0, -1};

    RoaringBitmap32 selection_bitmap = RoaringBitmap32::From({1, 3, 5});
    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[0], read_fields[2]}), R"([
        [0, "00"],
        [1, "01"],
        [2, "02"],
        [3, "03"],
        [4, "04"],
        [5, "05"]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[5], read_fields[3]}), R"([
        [10, 110],
        [11, 111],
        [12, 112],
        [13, 113],
        [14, 114],
        [15, 115]
])")
                      .ValueOrDie();
    auto array2 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[4]}), R"([
        [20, "20"],
        [21, "21"],
        [22, "22"],
        [23, "23"],
        [24, "24"],
        [25, "25"]
])")
                      .ValueOrDie();

    auto expected_array =
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(read_fields), R"([
        [1, 21, "01", 111, "21", 11, null],
        [3, 23, "03", 113, "23", 13, null],
        [5, 25, "05", 115, "25", 15, null]
])")
            .ValueOrDie();
    CheckResult({array0, array1, array2}, read_schema, reader_offsets, field_offsets,
                expected_array, selection_bitmap);
}

TEST_P(DataEvolutionFileReaderTest, TestSingleReaderRowCountMismatch) {
    arrow::FieldVector read_fields = {
        arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32()),
        arrow::field("f2", arrow::utf8()), arrow::field("f3", arrow::int32())};
    auto read_schema = arrow::schema(read_fields);

    std::vector<int32_t> reader_offsets = {0, 1, 0, 1};
    std::vector<int32_t> field_offsets = {0, 0, 1, 1};

    auto array0 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[0], read_fields[2]}), R"([
        [0, "00"],
        [1, "01"],
        [2, "02"],
        [3, "03"],
        [4, "04"],
        [5, "05"]
])")
                      .ValueOrDie();
    auto array1 = arrow::ipc::internal::json::ArrayFromJSON(
                      arrow::struct_({read_fields[1], read_fields[3]}), R"([
        [10, 110],
        [11, 111],
        [12, 112],
        [13, 113],
        [14, 114]
])")
                      .ValueOrDie();
    std::vector<std::unique_ptr<BatchReader>> readers;
    for (const auto& array : {array0, array1}) {
        auto file_batch_reader =
            std::make_unique<MockFileBatchReader>(array, array->type(), /*read_batch_size=*/10);
        auto enable_randomize_batch_size = GetParam();
        file_batch_reader->EnableRandomizeBatchSize(enable_randomize_batch_size);
        readers.push_back(std::move(file_batch_reader));
    }
    ASSERT_OK_AND_ASSIGN(
        auto data_evolution_file_reader,
        DataEvolutionFileReader::Create(std::move(readers), read_schema, /*read_batch_size=*/10,
                                        reader_offsets, field_offsets, pool_));
    // array0 has 6 rows but array1 only has 5 rows
    ASSERT_NOK_WITH_MSG(
        paimon::test::ReadResultCollector::CollectResult(data_evolution_file_reader.get()),
        "array for single reader length mismatch others");
}

INSTANTIATE_TEST_SUITE_P(EnableRandomizeBatchSize, DataEvolutionFileReaderTest,
                         ::testing::ValuesIn({true, false}));

}  // namespace paimon::test
