src/iceberg/test/arrow_row_builder_test.cc

/*
 * 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.
 */

/// \file arrow_row_builder_test.cc
/// Unit tests for ArrowRowBuilder and its typed append helpers.

#include <memory>
#include <vector>

#include <arrow/array.h>
#include <arrow/c/bridge.h>
#include <arrow/record_batch.h>
#include <gtest/gtest.h>

#include "iceberg/arrow_row_builder_internal.h"
#include "iceberg/schema.h"
#include "iceberg/schema_field.h"
#include "iceberg/schema_internal.h"
#include "iceberg/test/matchers.h"
#include "iceberg/type.h"

namespace iceberg {
namespace {

/// \brief A schema exercising every append helper: int32, string, int64,
/// boolean, and map<string, string>.
std::shared_ptr<Schema> MakeTestSchema() {
  return std::make_shared<Schema>(std::vector<SchemaField>{
      SchemaField::MakeRequired(1, "id", int32()),
      SchemaField::MakeOptional(2, "name", string()),
      SchemaField::MakeOptional(3, "count", int64()),
      SchemaField::MakeOptional(4, "active", boolean()),
      SchemaField::MakeOptional(
          5, "props",
          std::make_shared<MapType>(SchemaField::MakeRequired(6, "key", string()),
                                    SchemaField::MakeRequired(7, "value", string())))});
}

/// \brief A schema exercising the list/map append helpers backing Iceberg's
/// optional DataFile fields (e.g. equality_ids, split_offsets, column_sizes,
/// lower_bounds): list<int32>, list<int64>, map<int32, int64>, and
/// map<int32, binary>.
std::shared_ptr<Schema> MakeListMapTestSchema() {
  return std::make_shared<Schema>(std::vector<SchemaField>{
      SchemaField::MakeOptional(1, "int32_list",
                                list(SchemaField::MakeRequired(
                                    2, std::string(ListType::kElementName), int32()))),
      SchemaField::MakeOptional(3, "int64_list",
                                list(SchemaField::MakeRequired(
                                    4, std::string(ListType::kElementName), int64()))),
      SchemaField::MakeOptional(
          5, "int_map",
          map(SchemaField::MakeRequired(6, std::string(MapType::kKeyName), int32()),
              SchemaField::MakeRequired(7, std::string(MapType::kValueName), int64()))),
      SchemaField::MakeOptional(
          8, "binary_map",
          map(SchemaField::MakeRequired(9, std::string(MapType::kKeyName), int32()),
              SchemaField::MakeRequired(10, std::string(MapType::kValueName),
                                        binary())))});
}

/// \brief Finish a builder and import the result into an Arrow RecordBatch.
std::shared_ptr<::arrow::RecordBatch> FinishAndImport(ArrowRowBuilder builder,
                                                      const Schema& schema) {
  auto array_result = std::move(builder).Finish();
  EXPECT_THAT(array_result, IsOk());

  ArrowSchema c_schema;
  EXPECT_THAT(ToArrowSchema(schema, &c_schema), IsOk());
  auto arrow_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie();

  // ImportRecordBatch takes ownership of the array and releases it.
  return ::arrow::ImportRecordBatch(&array_result.value(), arrow_schema).ValueOrDie();
}

}  // namespace

TEST(ArrowRowBuilderTest, BuildsRowsWithTypedValues) {
  auto schema = MakeTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  ASSERT_EQ(builder.num_columns(), 5);
  ASSERT_EQ(builder.num_rows(), 0);

  // Row 0
  ASSERT_THAT(AppendInt(builder.column(0), 1), IsOk());
  ASSERT_THAT(AppendString(builder.column(1), "alice"), IsOk());
  ASSERT_THAT(AppendInt(builder.column(2), 100), IsOk());
  ASSERT_THAT(AppendBoolean(builder.column(3), true), IsOk());
  ASSERT_THAT(AppendStringMap(builder.column(4), {{"k", "v"}}), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());
  ASSERT_EQ(builder.num_rows(), 1);

  // Row 1
  ASSERT_THAT(AppendInt(builder.column(0), 2), IsOk());
  ASSERT_THAT(AppendString(builder.column(1), "bob"), IsOk());
  ASSERT_THAT(AppendInt(builder.column(2), 200), IsOk());
  ASSERT_THAT(AppendBoolean(builder.column(3), false), IsOk());
  ASSERT_THAT(AppendStringMap(builder.column(4), {}), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());
  ASSERT_EQ(builder.num_rows(), 2);

  auto batch = FinishAndImport(std::move(builder), *schema);
  ASSERT_EQ(batch->num_rows(), 2);
  ASSERT_EQ(batch->num_columns(), 5);

  auto id = std::static_pointer_cast<::arrow::Int32Array>(batch->column(0));
  EXPECT_EQ(id->Value(0), 1);
  EXPECT_EQ(id->Value(1), 2);

  auto name = std::static_pointer_cast<::arrow::StringArray>(batch->column(1));
  EXPECT_EQ(name->GetString(0), "alice");
  EXPECT_EQ(name->GetString(1), "bob");

  auto count = std::static_pointer_cast<::arrow::Int64Array>(batch->column(2));
  EXPECT_EQ(count->Value(0), 100);
  EXPECT_EQ(count->Value(1), 200);

  auto active = std::static_pointer_cast<::arrow::BooleanArray>(batch->column(3));
  EXPECT_TRUE(active->Value(0));
  EXPECT_FALSE(active->Value(1));

  // props: one entry in row 0, empty (but non-null) map in row 1.
  auto props = std::static_pointer_cast<::arrow::MapArray>(batch->column(4));
  EXPECT_FALSE(props->IsNull(0));
  EXPECT_FALSE(props->IsNull(1));
  EXPECT_EQ(props->value_length(0), 1);
  EXPECT_EQ(props->value_length(1), 0);
}

TEST(ArrowRowBuilderTest, AppendsNullForOptionalColumns) {
  auto schema = MakeTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  ASSERT_THAT(AppendInt(builder.column(0), 42), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(AppendStringMap(builder.column(4), {}), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  auto batch = FinishAndImport(std::move(builder), *schema);
  ASSERT_EQ(batch->num_rows(), 1);

  auto id = std::static_pointer_cast<::arrow::Int32Array>(batch->column(0));
  EXPECT_FALSE(id->IsNull(0));
  EXPECT_EQ(id->Value(0), 42);

  EXPECT_TRUE(batch->column(1)->IsNull(0));
  EXPECT_TRUE(batch->column(2)->IsNull(0));
  EXPECT_TRUE(batch->column(3)->IsNull(0));
}

TEST(ArrowRowBuilderTest, AppendsMultiEntryStringMap) {
  auto schema = MakeTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  ASSERT_THAT(AppendInt(builder.column(0), 1), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(AppendStringMap(builder.column(4), {{"a", "1"}, {"b", "2"}, {"c", "3"}}),
              IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  auto batch = FinishAndImport(std::move(builder), *schema);
  auto props = std::static_pointer_cast<::arrow::MapArray>(batch->column(4));
  EXPECT_EQ(props->value_length(0), 3);
}

TEST(ArrowRowBuilderTest, EmptyBuilderProducesZeroRowBatch) {
  auto schema = MakeTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  auto batch = FinishAndImport(std::move(builder), *schema);
  EXPECT_EQ(batch->num_rows(), 0);
  EXPECT_EQ(batch->num_columns(), 5);
}

TEST(ArrowRowBuilderTest, ColumnIndexOutOfRangeReturnsNull) {
  auto schema = MakeTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  EXPECT_EQ(builder.num_columns(), 5);
  EXPECT_NE(builder.column(0), nullptr);
  EXPECT_NE(builder.column(4), nullptr);
  EXPECT_EQ(builder.column(-1), nullptr);
  EXPECT_EQ(builder.column(5), nullptr);
}

TEST(ArrowRowBuilderTest, AppendIntListWritesNullForEmptyInput) {
  auto schema = MakeListMapTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  // Row 0: non-empty lists.
  ASSERT_THAT(AppendIntList(builder.column(0), std::vector<int32_t>{1, 2}), IsOk());
  ASSERT_THAT(AppendIntList(builder.column(1), std::vector<int64_t>{3, 4}), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  // Row 1: empty lists must be encoded as null, not an empty list.
  ASSERT_THAT(AppendIntList(builder.column(0), std::vector<int32_t>{}), IsOk());
  ASSERT_THAT(AppendIntList(builder.column(1), std::vector<int64_t>{}), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  auto batch = FinishAndImport(std::move(builder), *schema);
  ASSERT_EQ(batch->num_rows(), 2);

  auto int32_list = std::static_pointer_cast<::arrow::ListArray>(batch->column(0));
  EXPECT_FALSE(int32_list->IsNull(0));
  EXPECT_TRUE(int32_list->IsNull(1));

  auto int64_list = std::static_pointer_cast<::arrow::ListArray>(batch->column(1));
  EXPECT_FALSE(int64_list->IsNull(0));
  EXPECT_TRUE(int64_list->IsNull(1));
}

TEST(ArrowRowBuilderTest, AppendIntMapWritesNullForEmptyInput) {
  auto schema = MakeListMapTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  // Row 0: a non-empty map.
  ASSERT_THAT(AppendNull(builder.column(0)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendIntMap(builder.column(2), {{1, 100}}), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  // Row 1: an empty map must be encoded as null, not an empty map.
  ASSERT_THAT(AppendNull(builder.column(0)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendIntMap(builder.column(2), {}), IsOk());
  ASSERT_THAT(AppendNull(builder.column(3)), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  auto batch = FinishAndImport(std::move(builder), *schema);
  ASSERT_EQ(batch->num_rows(), 2);

  auto int_map = std::static_pointer_cast<::arrow::MapArray>(batch->column(2));
  EXPECT_FALSE(int_map->IsNull(0));
  EXPECT_TRUE(int_map->IsNull(1));
}

TEST(ArrowRowBuilderTest, AppendBinaryMapWritesNullForEmptyInput) {
  auto schema = MakeListMapTestSchema();
  ICEBERG_UNWRAP_OR_FAIL(auto builder, ArrowRowBuilder::Make(*schema));

  // Row 0: a non-empty map.
  ASSERT_THAT(AppendNull(builder.column(0)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendBinaryMap(builder.column(3), {{1, {0x01, 0x02}}}), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  // Row 1: an empty map must be encoded as null, not an empty map.
  ASSERT_THAT(AppendNull(builder.column(0)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(1)), IsOk());
  ASSERT_THAT(AppendNull(builder.column(2)), IsOk());
  ASSERT_THAT(AppendBinaryMap(builder.column(3), {}), IsOk());
  ASSERT_THAT(builder.FinishRow(), IsOk());

  auto batch = FinishAndImport(std::move(builder), *schema);
  ASSERT_EQ(batch->num_rows(), 2);

  auto binary_map = std::static_pointer_cast<::arrow::MapArray>(batch->column(3));
  EXPECT_FALSE(binary_map->IsNull(0));
  EXPECT_TRUE(binary_map->IsNull(1));
}

}  // namespace iceberg