diff --git a/docs/source/user_guide/data_types.rst b/docs/source/user_guide/data_types.rst index add537adb..57f0eafd2 100644 --- a/docs/source/user_guide/data_types.rst +++ b/docs/source/user_guide/data_types.rst @@ -194,12 +194,13 @@ and `Arrow DataTypes ColumnarArray::GetArray(int32_t pos) const { + if (array_->type_id() == arrow::Type::FIXED_SIZE_LIST) { + auto fixed_size_list_array = checked_cast(array_); + auto fixed_size_list_type = + checked_pointer_cast(fixed_size_list_array->type()); + int32_t offset = static_cast(fixed_size_list_array->value_offset(offset_ + pos)); + return std::make_shared(fixed_size_list_array->values().get(), pool_, offset, + fixed_size_list_type->list_size()); + } auto list_array = checked_cast(array_); int32_t offset = list_array->value_offset(offset_ + pos); int32_t length = list_array->value_length(offset_ + pos); diff --git a/src/paimon/common/data/columnar/columnar_row.cpp b/src/paimon/common/data/columnar/columnar_row.cpp index 5099b42b2..40d17309d 100644 --- a/src/paimon/common/data/columnar/columnar_row.cpp +++ b/src/paimon/common/data/columnar/columnar_row.cpp @@ -66,6 +66,15 @@ std::shared_ptr ColumnarRow::GetRow(int32_t pos, int32_t num_fields } std::shared_ptr ColumnarRow::GetArray(int32_t pos) const { + if (array_vec_[pos]->type_id() == arrow::Type::FIXED_SIZE_LIST) { + auto fixed_size_list_array = + checked_cast(array_vec_[pos]); + auto fixed_size_list_type = + checked_pointer_cast(fixed_size_list_array->type()); + int32_t offset = static_cast(fixed_size_list_array->value_offset(row_id_)); + return std::make_shared(fixed_size_list_array->values().get(), pool_, offset, + fixed_size_list_type->list_size()); + } auto list_array = checked_cast(array_vec_[pos]); int32_t offset = list_array->value_offset(row_id_); int32_t length = list_array->value_length(row_id_); diff --git a/src/paimon/common/data/columnar/columnar_row_ref.cpp b/src/paimon/common/data/columnar/columnar_row_ref.cpp index d73e6b774..224d77a92 100644 --- a/src/paimon/common/data/columnar/columnar_row_ref.cpp +++ b/src/paimon/common/data/columnar/columnar_row_ref.cpp @@ -61,6 +61,15 @@ std::shared_ptr ColumnarRowRef::GetRow(int32_t pos, int32_t num_fie } std::shared_ptr ColumnarRowRef::GetArray(int32_t pos) const { + if (ctx_->array_vec[pos]->type_id() == arrow::Type::FIXED_SIZE_LIST) { + auto fixed_size_list_array = + checked_cast(ctx_->array_vec[pos].get()); + auto fixed_size_list_type = + checked_pointer_cast(fixed_size_list_array->type()); + int32_t offset = static_cast(fixed_size_list_array->value_offset(row_id_)); + return std::make_shared(fixed_size_list_array->values().get(), ctx_->pool, + offset, fixed_size_list_type->list_size()); + } auto list_array = checked_cast(ctx_->array_vec[pos].get()); int32_t offset = list_array->value_offset(row_id_); int32_t length = list_array->value_length(row_id_); diff --git a/src/paimon/common/data/internal_row.cpp b/src/paimon/common/data/internal_row.cpp index 24b6b5bee..b7c011b2e 100644 --- a/src/paimon/common/data/internal_row.cpp +++ b/src/paimon/common/data/internal_row.cpp @@ -128,7 +128,8 @@ Result InternalRow::CreateFieldGetter( }; break; } - case arrow::Type::type::LIST: { + case arrow::Type::type::LIST: + case arrow::Type::type::FIXED_SIZE_LIST: { field_getter = [field_idx](const InternalRow& row) -> VariantType { return row.GetArray(field_idx); }; diff --git a/src/paimon/common/data/serializer/binary_serializer_utils.cpp b/src/paimon/common/data/serializer/binary_serializer_utils.cpp index c0248cd60..d72a2b0b1 100644 --- a/src/paimon/common/data/serializer/binary_serializer_utils.cpp +++ b/src/paimon/common/data/serializer/binary_serializer_utils.cpp @@ -31,8 +31,7 @@ Result> BinarySerializerUtils::WriteBinaryArray( return binary_array; } auto binary_array = std::make_shared(); - auto list_type = checked_pointer_cast(type); - auto value_type = list_type->value_type(); + auto value_type = type->field(0)->type(); // TODO(xinyu.lxy): reuse BinaryWriter BinaryArrayWriter binary_writer(binary_array.get(), value->Size(), BinaryArrayWriter::GetElementSize(value_type->id()), pool); @@ -183,7 +182,8 @@ Status BinarySerializerUtils::WriteBinaryData(const std::shared_ptrGetArray(pos); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr binary_array, WriteBinaryArray(internal_array, type, pool)); diff --git a/src/paimon/common/data/serializer/row_compacted_serializer.cpp b/src/paimon/common/data/serializer/row_compacted_serializer.cpp index 411dae687..1fd618a39 100644 --- a/src/paimon/common/data/serializer/row_compacted_serializer.cpp +++ b/src/paimon/common/data/serializer/row_compacted_serializer.cpp @@ -285,7 +285,8 @@ Result RowCompactedSerializer::CreateFieldR }; break; } - case arrow::Type::type::LIST: { + case arrow::Type::type::LIST: + case arrow::Type::type::FIXED_SIZE_LIST: { field_reader = [](int32_t pos, RowReader* reader) -> Result { PAIMON_ASSIGN_OR_RAISE(VariantType value, reader->ReadArray()); return value; @@ -414,7 +415,8 @@ Result RowCompactedSerializer::CreateFieldW }; break; } - case arrow::Type::type::LIST: { + case arrow::Type::type::LIST: + case arrow::Type::type::FIXED_SIZE_LIST: { field_writer = [field_type](int32_t pos, const VariantType& field, RowWriter* writer) -> Status { return writer->WriteArray( diff --git a/src/paimon/core/io/key_value_in_memory_record_reader.cpp b/src/paimon/core/io/key_value_in_memory_record_reader.cpp index 9918b328c..f4d3e8870 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader.cpp +++ b/src/paimon/core/io/key_value_in_memory_record_reader.cpp @@ -111,22 +111,42 @@ void KeyValueInMemoryRecordReader::Close() { Result>> KeyValueInMemoryRecordReader::SortBatch() const { std::vector sort_keys; + arrow::FieldVector sort_fields; + arrow::ArrayVector sort_columns; sort_keys.reserve(primary_keys_.size() + user_defined_sequence_fields_.size()); + sort_fields.reserve(primary_keys_.size() + user_defined_sequence_fields_.size()); + sort_columns.reserve(primary_keys_.size() + user_defined_sequence_fields_.size()); + const arrow::StructType* value_type = value_struct_array_->struct_type(); + auto append_sort_key = [&](const std::string& name, arrow::compute::SortOrder order) -> Status { + int32_t field_index = value_type->GetFieldIndex(name); + if (field_index < 0) { + return Status::Invalid(fmt::format("cannot find field {} in data batch", name)); + } + sort_keys.emplace_back(name, order); + sort_fields.push_back(value_type->field(field_index)); + sort_columns.push_back(value_struct_array_->field(field_index)); + return Status::OK(); + }; for (const auto& name : primary_keys_) { - sort_keys.emplace_back(name, arrow::compute::SortOrder::Ascending); + PAIMON_RETURN_NOT_OK(append_sort_key(name, arrow::compute::SortOrder::Ascending)); } const auto sequence_sort_order = sequence_fields_ascending_ ? arrow::compute::SortOrder::Ascending : arrow::compute::SortOrder::Descending; for (const auto& name : user_defined_sequence_fields_) { - sort_keys.emplace_back(name, sequence_sort_order); + PAIMON_RETURN_NOT_OK(append_sort_key(name, sequence_sort_order)); } auto sort_options = arrow::compute::SortOptions(sort_keys, arrow::compute::NullPlacement::AtStart); arrow::compute::ExecContext exec_context(arrow_pool_.get()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr sorted_indices, - arrow::compute::SortIndices(arrow::Datum(value_struct_array_), - sort_options, &exec_context)); + // Arrow's StructArray sorting path may inspect value columns outside the sort keys. Restrict + // the batch to the requested fields so non-sortable values, such as VECTOR, are never compared. + std::shared_ptr sort_batch = + arrow::RecordBatch::Make(arrow::schema(std::move(sort_fields)), + value_struct_array_->length(), std::move(sort_columns)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr sorted_indices, + arrow::compute::SortIndices(arrow::Datum(sort_batch), sort_options, &exec_context)); if (!sorted_indices || sorted_indices->type_id() != arrow::Type::UINT64) { return Status::Invalid("cannot cast sorted indices to UInt64Array"); } diff --git a/src/paimon/core/io/row_to_arrow_array_converter.h b/src/paimon/core/io/row_to_arrow_array_converter.h index 6caa9f618..336d9857f 100644 --- a/src/paimon/core/io/row_to_arrow_array_converter.h +++ b/src/paimon/core/io/row_to_arrow_array_converter.h @@ -158,6 +158,12 @@ Status RowToArrowArrayConverter::Reserve(arrow::ArrayBuilder* array_builde PAIMON_RETURN_NOT_OK(Reserve(list_builder->value_builder(), idx)); break; } + case arrow::Type::type::FIXED_SIZE_LIST: { + PAIMON_ASSIGN_OR_RAISE(auto* list_builder, + CastToTypedBuilder(array_builder)); + PAIMON_RETURN_NOT_OK(Reserve(list_builder->value_builder(), idx)); + break; + } case arrow::Type::type::MAP: { PAIMON_ASSIGN_OR_RAISE(auto* map_builder, CastToTypedBuilder(array_builder)); @@ -224,6 +230,11 @@ Status RowToArrowArrayConverter::Accumulate(const arrow::Array* array, int PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx)); break; } + case arrow::Type::type::FIXED_SIZE_LIST: { + auto list_array = checked_cast(array); + PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx)); + break; + } case arrow::Type::type::MAP: { auto map_array = checked_cast(array); PAIMON_RETURN_NOT_OK(Accumulate(map_array->keys().get(), idx)); @@ -432,6 +443,31 @@ RowToArrowArrayConverter::AppendField(bool use_view, arrow::ArrayBuilder* return arrow::Status::OK(); }); } + case arrow::Type::type::FIXED_SIZE_LIST: { + PAIMON_ASSIGN_OR_RAISE(auto* list_builder, + CastToTypedBuilder(array_builder)); + std::shared_ptr list_type = + checked_pointer_cast(list_builder->type()); + int32_t list_size = list_type->list_size(); + PAIMON_ASSIGN_OR_RAISE(AppendValueFunc value_func, + (RowToArrowArrayConverter::AppendField( + use_view, list_builder->value_builder(), reserve_count))); + return RowToArrowArrayConverter::AppendValueFunc( + [list_builder, list_size, value_func](const DataGetters& data_getter, + int32_t pos) -> arrow::Status { + CHECK_AND_APPEND_NULL(data_getter, list_builder, pos); + std::shared_ptr sub_array = data_getter.GetArray(pos); + if (!sub_array || sub_array->Size() != list_size) { + return arrow::Status::Invalid( + "VECTOR length does not match its declared dimension"); + } + ARROW_RETURN_NOT_OK(list_builder->Append()); + for (int32_t i = 0; i < list_size; ++i) { + ARROW_RETURN_NOT_OK(value_func(*sub_array, i)); + } + return arrow::Status::OK(); + }); + } case arrow::Type::type::MAP: { PAIMON_ASSIGN_OR_RAISE(auto* map_builder, CastToTypedBuilder(array_builder)); diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp index 9397ca4ec..22e02f108 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp @@ -73,12 +73,12 @@ Result EqualRows(const std::shared_ptr& lhs, Result EqualArrays(const std::shared_ptr& lhs, const std::shared_ptr& rhs, - const std::shared_ptr& type) { + const std::shared_ptr& type) { if (!lhs || !rhs || lhs->Size() != rhs->Size()) { return lhs == rhs; } for (int32_t i = 0; i < lhs->Size(); ++i) { - PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->value_type())); + PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->field(0)->type())); if (!equal) { return false; } @@ -174,6 +174,7 @@ Result FieldAggregateUtils::GetValue(const DataGetters& getters, in getters.GetDecimal(pos, decimal_type->precision(), decimal_type->scale())); } case arrow::Type::LIST: + case arrow::Type::FIXED_SIZE_LIST: return VariantType(getters.GetArray(pos)); case arrow::Type::MAP: return VariantType(getters.GetMap(pos)); @@ -229,9 +230,10 @@ Result FieldAggregateUtils::Equals(const VariantType& lhs, const VariantTy DataDefine::GetVariantValue>(rhs), checked_pointer_cast(type)); case arrow::Type::LIST: + case arrow::Type::FIXED_SIZE_LIST: return EqualArrays(DataDefine::GetVariantValue>(lhs), DataDefine::GetVariantValue>(rhs), - checked_pointer_cast(type)); + type); case arrow::Type::MAP: return EqualMaps(DataDefine::GetVariantValue>(lhs), DataDefine::GetVariantValue>(rhs), diff --git a/src/paimon/core/mergetree/compact/internal_row_equalizer.h b/src/paimon/core/mergetree/compact/internal_row_equalizer.h index 22f15d8bd..a0a0b4c55 100644 --- a/src/paimon/core/mergetree/compact/internal_row_equalizer.h +++ b/src/paimon/core/mergetree/compact/internal_row_equalizer.h @@ -144,11 +144,10 @@ class InternalRowEqualizer { .CompareTo(rhs.GetDecimal(rhs_pos, precision, scale)) == 0; }); } - case arrow::Type::LIST: { - std::shared_ptr list_type = - checked_pointer_cast(type); + case arrow::Type::LIST: + case arrow::Type::FIXED_SIZE_LIST: { PAIMON_ASSIGN_OR_RAISE(ValueEqualizer element_equalizer, - CreateValueEqualizer(list_type->value_type())); + CreateValueEqualizer(type->field(0)->type())); return ValueEqualizer([element_equalizer = std::move(element_equalizer)]( const DataGetters& lhs, int32_t lhs_pos, const DataGetters& rhs, int32_t rhs_pos) { diff --git a/src/paimon/core/mergetree/in_memory_sort_buffer.cpp b/src/paimon/core/mergetree/in_memory_sort_buffer.cpp index 5f53d5984..272d34b95 100644 --- a/src/paimon/core/mergetree/in_memory_sort_buffer.cpp +++ b/src/paimon/core/mergetree/in_memory_sort_buffer.cpp @@ -156,6 +156,11 @@ Result InMemorySortBuffer::EstimateMemoryUse(const std::shared_ptrvalues())); return null_bits_size_in_bytes + value_mem; } + case arrow::Type::type::FIXED_SIZE_LIST: { + auto list_array = checked_cast(array.get()); + PAIMON_ASSIGN_OR_RAISE(int64_t value_mem, EstimateMemoryUse(list_array->values())); + return null_bits_size_in_bytes + value_mem; + } case arrow::Type::type::MAP: { auto map_array = checked_cast(array.get()); PAIMON_ASSIGN_OR_RAISE(int64_t key_mem, EstimateMemoryUse(map_array->keys())); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 0b2849a63..1df5ddacf 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -87,17 +87,36 @@ Result> CreateRealtimePrimaryKeyTransportBat arrow::StructArray::Make(std::move(columns), transport_schema->fields())); std::vector sort_keys; + std::vector sort_field_names; sort_keys.reserve(trimmed_primary_keys.size() + 1); + sort_field_names.reserve(trimmed_primary_keys.size() + 1); for (const std::string& key : trimmed_primary_keys) { sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + sort_field_names.push_back(key); } sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), arrow::compute::SortOrder::Ascending); + sort_field_names.push_back(SpecialFields::SequenceNumber().Name()); arrow::compute::ExecContext context(arrow_pool); arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + arrow::FieldVector sort_fields; + arrow::ArrayVector sort_columns; + sort_fields.reserve(sort_keys.size()); + sort_columns.reserve(sort_keys.size()); + const arrow::StructType* transport_type = transport->struct_type(); + for (const std::string& name : sort_field_names) { + int32_t field_index = transport_type->GetFieldIndex(name); + if (field_index < 0) { + return Status::Invalid("PK sort field is missing from transport batch: ", name); + } + sort_fields.push_back(transport_type->field(field_index)); + sort_columns.push_back(transport->field(field_index)); + } + std::shared_ptr sort_batch = arrow::RecordBatch::Make( + arrow::schema(std::move(sort_fields)), transport->length(), std::move(sort_columns)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( arrow::Datum indices, - arrow::compute::SortIndices(arrow::Datum(transport), options, &context)); + arrow::compute::SortIndices(arrow::Datum(sort_batch), options, &context)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( arrow::Datum sorted, arrow::compute::Take(arrow::Datum(transport), indices, diff --git a/src/paimon/core/schema/schema_validation.cpp b/src/paimon/core/schema/schema_validation.cpp index 5e524a376..26d4100c8 100644 --- a/src/paimon/core/schema/schema_validation.cpp +++ b/src/paimon/core/schema/schema_validation.cpp @@ -127,6 +127,16 @@ Status ValidatePerLevelOption( return Status::OK(); } +Status ValidateVectorComparatorField(const TableSchema& schema, const std::string& field_name, + const std::string& role) { + PAIMON_ASSIGN_OR_RAISE(DataField field, schema.GetField(field_name)); + if (VectorUtils::ContainsVectorField(field.ArrowField())) { + return Status::Invalid( + fmt::format("VECTOR field '{}' cannot be used as {}.", field_name, role)); + } + return Status::OK(); +} + } // namespace bool SchemaValidation::IsComplexType(const std::shared_ptr& field) { @@ -378,6 +388,8 @@ Status SchemaValidation::ValidateSequenceGroup(const TableSchema& schema, fmt::format("The sequence field group: {} can not be found in table schema.", sequence_field_name)); } + PAIMON_RETURN_NOT_OK(ValidateVectorComparatorField(schema, sequence_field_name, + "a sequence-group ordering field")); } for (const auto& field : StringUtils::Split(v, Options::FIELDS_SEPARATOR)) { @@ -445,6 +457,7 @@ Status SchemaValidation::ValidateSequenceField(const TableSchema& schema, PAIMON_RETURN_NOT_OK(Preconditions::CheckState( std::find(field_names.begin(), field_names.end(), field) != field_names.end(), fmt::format("Sequence field: '{}' cannot be found in table schema.", field))); + PAIMON_RETURN_NOT_OK(ValidateVectorComparatorField(schema, field, "a sequence field")); PAIMON_ASSIGN_OR_RAISE(std::optional agg_func, options.GetFieldAggFunc(field)); @@ -776,10 +789,6 @@ Status SchemaValidation::ValidateVectorFields(const TableSchema& schema, if (!has_vector) { return Status::OK(); } - if (!schema.PrimaryKeys().empty()) { - return Status::NotImplemented( - "VECTOR fields in primary-key tables are not implemented yet."); - } if (options.DataEvolutionEnabled()) { return Status::NotImplemented( "VECTOR fields in data-evolution tables are not implemented yet."); diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 054f76da3..092b2ee25 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -86,12 +86,54 @@ TEST(SchemaValidationTest, TestVectorType) { ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), "in primary key field embedding is unsupported"); + ASSERT_OK_AND_ASSIGN( + table_schema, TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{"embedding"}, + /*primary_keys=*/{}, parquet_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "in partition field embedding is unsupported"); + + std::map bucket_key_options = { + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "embedding"}, + {Options::FILE_FORMAT, "parquet"}, + }; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, bucket_key_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "Nested type cannot be in bucket-key"); + + std::map sequence_field_options = { + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::SEQUENCE_FIELD, "embedding"}, + }; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"id"}, sequence_field_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR field 'embedding' cannot be used as a sequence field."); + + std::map sequence_group_options = { + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::MERGE_ENGINE, "partial-update"}, + {"fields.embedding.sequence-group", "id"}, + }; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"id"}, sequence_group_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR field 'embedding' cannot be used as a sequence-group ordering " + "field."); + primary_key_options[Options::FILE_FORMAT] = "parquet"; ASSERT_OK_AND_ASSIGN(table_schema, TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, /*primary_keys=*/{"id"}, primary_key_options)); - ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), - "VECTOR fields in primary-key tables are not implemented yet."); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); auto nested_schema = arrow::schema({ arrow::field("id", arrow::int64()), @@ -101,8 +143,7 @@ TEST(SchemaValidationTest, TestVectorType) { table_schema, TableSchema::Create(/*schema_id=*/0, nested_schema, /*partition_keys=*/{}, /*primary_keys=*/{"id"}, primary_key_options)); - ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), - "VECTOR fields in primary-key tables are not implemented yet."); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); std::map data_evolution_options = { {Options::BUCKET, "-1"}, @@ -116,10 +157,10 @@ TEST(SchemaValidationTest, TestVectorType) { /*primary_keys=*/{}, data_evolution_options)); ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), "VECTOR fields in data-evolution tables are not implemented yet."); - ASSERT_OK_AND_ASSIGN(table_schema, - TableSchema::Create(/*schema_id=*/0, nested_schema, - /*partition_keys=*/{}, - /*primary_keys=*/{}, data_evolution_options)); + ASSERT_OK_AND_ASSIGN( + table_schema, + TableSchema::Create(/*schema_id=*/0, nested_schema, + /*partition_keys=*/{}, /*primary_keys=*/{}, data_evolution_options)); ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), "VECTOR fields in data-evolution tables are not implemented yet."); } diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index 72bc37659..d93f12056 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -17,6 +17,7 @@ */ #include +#include #include #include #include @@ -69,6 +70,180 @@ #include "rapidjson/writer.h" namespace paimon::test { +namespace { + +struct PrimaryKeyVectorRow { + int64_t primary_key; + std::optional> embedding; + std::optional tag; +}; + +struct NestedPrimaryKeyVectorPayload { + std::optional> embedding; + std::optional tag; +}; + +struct NestedPrimaryKeyVectorRow { + int64_t primary_key; + std::optional payload; +}; + +Result> MakePrimaryKeyVectorArray( + const arrow::FieldVector& fields, const std::vector& rows, + bool include_row_kind) { + const int32_t data_field_offset = include_row_kind ? 1 : 0; + const bool has_tags = fields.size() == static_cast(data_field_offset + 3); + if (fields.size() != static_cast(data_field_offset + (has_tags ? 3 : 2))) { + return Status::Invalid("unexpected primary-key VECTOR test schema"); + } + + arrow::Int8Builder row_kind_builder; + arrow::Int64Builder primary_key_builder; + std::shared_ptr embedding_value_builder = + std::make_shared(); + arrow::FixedSizeListBuilder embedding_builder(arrow::default_memory_pool(), + embedding_value_builder, + fields[data_field_offset + 1]->type()); + arrow::StringBuilder tag_builder; + + for (const PrimaryKeyVectorRow& row : rows) { + if (include_row_kind) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Append(0)); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Append(row.primary_key)); + if (row.embedding.has_value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder.Append()); + for (float value : row.embedding.value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_value_builder->Append(value)); + } + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder.AppendNull()); + } + if (has_tags) { + if (row.tag.has_value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(tag_builder.Append(row.tag.value())); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(tag_builder.AppendNull()); + } + } + } + + std::vector> arrays; + arrays.reserve(fields.size()); + if (include_row_kind) { + std::shared_ptr row_kind_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); + arrays.push_back(std::move(row_kind_array)); + } + std::shared_ptr primary_key_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Finish(&primary_key_array)); + arrays.push_back(std::move(primary_key_array)); + std::shared_ptr embedding_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder.Finish(&embedding_array)); + arrays.push_back(std::move(embedding_array)); + if (has_tags) { + std::shared_ptr tag_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(tag_builder.Finish(&tag_array)); + arrays.push_back(std::move(tag_array)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::StructArray::Make(arrays, fields)); + return array; +} + +Result> MakePrimaryKeyVectorRecordBatch( + const arrow::FieldVector& fields, const std::vector& rows) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr array, + MakePrimaryKeyVectorArray(fields, rows, /*include_row_kind=*/false)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder batch_builder(&c_array); + return batch_builder.SetBucket(0).Finish(); +} + +Result> MakeNestedPrimaryKeyVectorArray( + const arrow::FieldVector& fields, const std::vector& rows, + bool include_row_kind) { + const int32_t data_field_offset = include_row_kind ? 1 : 0; + if (fields.size() != static_cast(data_field_offset + 2) || + fields[data_field_offset + 1]->type()->id() != arrow::Type::STRUCT) { + return Status::Invalid("unexpected nested primary-key VECTOR test schema"); + } + std::shared_ptr payload_type = + checked_pointer_cast(fields[data_field_offset + 1]->type()); + if (payload_type->num_fields() != 2) { + return Status::Invalid("unexpected nested primary-key VECTOR payload schema"); + } + + arrow::Int8Builder row_kind_builder; + arrow::Int64Builder primary_key_builder; + std::shared_ptr embedding_value_builder = + std::make_shared(); + std::shared_ptr embedding_builder = + std::make_shared( + arrow::default_memory_pool(), embedding_value_builder, payload_type->field(0)->type()); + std::shared_ptr tag_builder = std::make_shared(); + arrow::StructBuilder payload_builder(payload_type, arrow::default_memory_pool(), + {embedding_builder, tag_builder}); + + for (const NestedPrimaryKeyVectorRow& row : rows) { + if (include_row_kind) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Append(0)); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Append(row.primary_key)); + if (!row.payload.has_value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(payload_builder.AppendNull()); + continue; + } + + PAIMON_RETURN_NOT_OK_FROM_ARROW(payload_builder.Append()); + const NestedPrimaryKeyVectorPayload& payload = row.payload.value(); + if (payload.embedding.has_value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder->Append()); + for (float value : payload.embedding.value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_value_builder->Append(value)); + } + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder->AppendNull()); + } + if (payload.tag.has_value()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(tag_builder->Append(payload.tag.value())); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(tag_builder->AppendNull()); + } + } + + std::vector> arrays; + arrays.reserve(fields.size()); + if (include_row_kind) { + std::shared_ptr row_kind_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); + arrays.push_back(std::move(row_kind_array)); + } + std::shared_ptr primary_key_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Finish(&primary_key_array)); + arrays.push_back(std::move(primary_key_array)); + std::shared_ptr payload_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(payload_builder.Finish(&payload_array)); + arrays.push_back(std::move(payload_array)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::StructArray::Make(arrays, fields)); + return array; +} + +Result> MakeNestedPrimaryKeyVectorRecordBatch( + const arrow::FieldVector& fields, const std::vector& rows) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr array, + MakeNestedPrimaryKeyVectorArray(fields, rows, /*include_row_kind=*/false)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder batch_builder(&c_array); + return batch_builder.SetBucket(0).Finish(); +} + +} // namespace + // This is a sdk end-to-end test demo that supports write, commit, scan, and read operations. class WriteAndReadInteTest : public ::testing::Test, @@ -678,6 +853,204 @@ TEST_P(WriteAndReadInteTest, TestPKSimple) { ASSERT_TRUE(success); } +TEST_P(WriteAndReadInteTest, TestPKVector) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = { + arrow::field("pk", arrow::int64()), + arrow::field("embedding", vector_type), + }; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, file_system}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN( + auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, /*is_streaming_mode=*/true)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_batch, + MakePrimaryKeyVectorRecordBatch( + fields, {{1, std::array{1.0F, 2.0F, 3.0F}, std::nullopt}, + {2, std::nullopt, std::nullopt}})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + MakePrimaryKeyVectorRecordBatch( + fields, {{1, std::array{4.0F, 5.0F, 6.0F}, std::nullopt}, + {3, std::array{7.0F, 8.0F, 9.0F}, std::nullopt}})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/1, + /*expected_commit_messages=*/std::nullopt)); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/2)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + helper->ReadResult(data_splits)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + MakePrimaryKeyVectorArray(result_fields, + {{1, std::array{4.0F, 5.0F, 6.0F}, std::nullopt}, + {2, std::nullopt, std::nullopt}, + {3, std::array{7.0F, 8.0F, 9.0F}, std::nullopt}}, + /*include_row_kind=*/true)); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)); +} + +TEST_P(WriteAndReadInteTest, TestPKNestedVector) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = { + arrow::field("pk", arrow::int64()), + arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_type), + arrow::field("tag", arrow::utf8())})), + }; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, file_system}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN( + auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, /*is_streaming_mode=*/true)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr initial_batch, + MakeNestedPrimaryKeyVectorRecordBatch( + fields, {{1, NestedPrimaryKeyVectorPayload{std::array{1.0F, 2.0F, 3.0F}, + std::string("initial")}}, + {2, NestedPrimaryKeyVectorPayload{std::nullopt, std::string("null-vector")}}, + {3, std::nullopt}})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr update_batch, + MakeNestedPrimaryKeyVectorRecordBatch( + fields, {{1, NestedPrimaryKeyVectorPayload{std::array{4.0F, 5.0F, 6.0F}, + std::string("updated")}}, + {2, NestedPrimaryKeyVectorPayload{std::array{7.0F, 8.0F, 9.0F}, + std::nullopt}}})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/1, + /*expected_commit_messages=*/std::nullopt)); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/2)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + helper->ReadResult(data_splits)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + MakeNestedPrimaryKeyVectorArray( + result_fields, + {{1, NestedPrimaryKeyVectorPayload{std::array{4.0F, 5.0F, 6.0F}, + std::string("updated")}}, + {2, + NestedPrimaryKeyVectorPayload{std::array{7.0F, 8.0F, 9.0F}, std::nullopt}}, + {3, std::nullopt}}, + /*include_row_kind=*/true)); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)); +} + +TEST_P(WriteAndReadInteTest, TestPKVectorWithListagg) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = { + arrow::field("pk", arrow::int64()), + arrow::field("embedding", vector_type), + arrow::field("tags", arrow::utf8()), + }; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, + {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, file_system}, + {Options::MERGE_ENGINE, "aggregation"}, + {"fields.embedding.aggregate-function", "last_non_null_value"}, + {"fields.tags.aggregate-function", "listagg"}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN( + auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, /*is_streaming_mode=*/true)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr initial_batch, + MakePrimaryKeyVectorRecordBatch( + fields, {{1, std::array{1.0F, 2.0F, 3.0F}, std::string("alpha")}, + {2, std::nullopt, std::string("one")}})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr update_batch, + MakePrimaryKeyVectorRecordBatch( + fields, {{1, std::array{4.0F, 5.0F, 6.0F}, std::string("beta")}, + {2, std::array{7.0F, 8.0F, 9.0F}, std::string("two")}})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/1, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr null_vector_batch, + MakePrimaryKeyVectorRecordBatch(fields, {{1, std::nullopt, std::string("gamma")}})); + ASSERT_OK(helper->WriteAndCommit(std::move(null_vector_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + helper->ReadResult(data_splits)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + MakePrimaryKeyVectorArray( + result_fields, + {{1, std::array{4.0F, 5.0F, 6.0F}, std::string("alpha,beta,gamma")}, + {2, std::array{7.0F, 8.0F, 9.0F}, std::string("one,two")}}, + /*include_row_kind=*/true)); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)); +} + TEST_P(WriteAndReadInteTest, TestInputChangelogStreamRead) { arrow::FieldVector fields = { arrow::field("pk", arrow::utf8()), diff --git a/test/inte/write_inte_test.cpp b/test/inte/write_inte_test.cpp index 63768416c..48fca8ac8 100644 --- a/test/inte/write_inte_test.cpp +++ b/test/inte/write_inte_test.cpp @@ -15,6 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include #include #include #include @@ -106,6 +107,66 @@ class TableSchema; } // namespace paimon namespace paimon::test { +namespace { + +struct SpillableVectorRow { + std::string primary_key; + int32_t partition; + std::array embedding; +}; + +Result> MakeSpillableVectorArray( + const arrow::FieldVector& fields, const std::vector& rows, + bool include_row_kind) { + const int32_t data_field_offset = include_row_kind ? 1 : 0; + if (fields.size() != static_cast(data_field_offset + 3)) { + return Status::Invalid("unexpected spillable VECTOR test schema"); + } + + arrow::Int8Builder row_kind_builder; + arrow::StringBuilder primary_key_builder; + arrow::Int32Builder partition_builder; + std::shared_ptr embedding_value_builder = + std::make_shared(); + arrow::FixedSizeListBuilder embedding_builder(arrow::default_memory_pool(), + embedding_value_builder, + fields[data_field_offset + 2]->type()); + + for (const SpillableVectorRow& row : rows) { + if (include_row_kind) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Append(0)); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Append(row.primary_key)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(partition_builder.Append(row.partition)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder.Append()); + for (float value : row.embedding) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_value_builder->Append(value)); + } + } + + std::vector> arrays; + arrays.reserve(fields.size()); + if (include_row_kind) { + std::shared_ptr row_kind_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); + arrays.push_back(std::move(row_kind_array)); + } + std::shared_ptr primary_key_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(primary_key_builder.Finish(&primary_key_array)); + arrays.push_back(std::move(primary_key_array)); + std::shared_ptr partition_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(partition_builder.Finish(&partition_array)); + arrays.push_back(std::move(partition_array)); + std::shared_ptr embedding_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(embedding_builder.Finish(&embedding_array)); + arrays.push_back(std::move(embedding_array)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::StructArray::Make(arrays, fields)); + return array; +} + +} // namespace + class WriteInteTest : public testing::Test, public ::testing::WithParamInterface { public: void SetUp() override { @@ -4393,6 +4454,92 @@ TEST_P(WriteInteTest, TestPkSpillableIntermediateMergeWithTempFileTracking) { ASSERT_OK(ScanAndVerifyResult(table_path, fields, expected)); } +TEST_P(WriteInteTest, TestPkSpillableVector) { + auto file_format = GetParam(); + if (file_format != "parquet") { + return; + } + + auto dir = UniqueTestDirectory::Create(); + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = { + arrow::field("f0", arrow::utf8()), + arrow::field("pt", arrow::int32()), + arrow::field("embedding", vector_type), + }; + std::map options = { + {Options::FILE_FORMAT, file_format}, + {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, "local"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + {Options::WRITE_BUFFER_SPILLABLE, "true"}, + {Options::LOCAL_SORT_MAX_NUM_FILE_HANDLES, "2"}, + {Options::WRITE_ONLY, "true"}, + }; + auto schema = arrow::schema(fields); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(auto table_path, CreateTestTable(dir->Str(), "db", "tbl", &c_schema, + /*partition_keys=*/{"pt"}, + /*primary_keys=*/{"pt", "f0"}, options)); + + std::string tmp_dir = PathUtil::JoinPath(dir->Str(), "tmp"); + WriteContextBuilder write_builder(table_path, "commit_user_1"); + write_builder.WithStreamingMode(true).WithTempDirectory(tmp_dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, write_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); + + auto write_array = [](FileStoreWrite* writer, const std::shared_ptr& array) { + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + auto batch = std::make_unique(std::map{{"pt", "10"}}, + /*bucket=*/0, + std::vector{}, &c_array); + return writer->Write(std::move(batch)); + }; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch1, + MakeSpillableVectorArray(fields, {{"Alice", 10, {1.0F, 2.0F, 3.0F}}}, + /*include_row_kind=*/false)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch2, + MakeSpillableVectorArray(fields, {{"Bob", 10, {4.0F, 5.0F, 6.0F}}}, + /*include_row_kind=*/false)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch3, + MakeSpillableVectorArray(fields, {{"Alice", 10, {7.0F, 8.0F, 9.0F}}}, + /*include_row_kind=*/false)); + + ASSERT_OK(write_array(file_store_write.get(), batch1)); + ASSERT_EQ(1, TestHelper::CountChannelFiles(file_system_, tmp_dir)); + ASSERT_OK(write_array(file_store_write.get(), batch2)); + ASSERT_EQ(1, TestHelper::CountChannelFiles(file_system_, tmp_dir)); + ASSERT_OK(write_array(file_store_write.get(), batch3)); + ASSERT_EQ(2, TestHelper::CountChannelFiles(file_system_, tmp_dir)); + + ASSERT_OK_AND_ASSIGN(auto commit_messages, + file_store_write->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/0)); + ASSERT_EQ(0, TestHelper::CountChannelFiles(file_system_, tmp_dir)); + ASSERT_OK(CommitMessages(table_path, commit_messages)); + ASSERT_OK(file_store_write->Close()); + + std::map scan_options = {{Options::FILE_SYSTEM, "local"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(table_path, scan_options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + helper->ReadResult(data_splits)); + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + MakeSpillableVectorArray( + result_fields, {{"Alice", 10, {7.0F, 8.0F, 9.0F}}, {"Bob", 10, {4.0F, 5.0F, 6.0F}}}, + /*include_row_kind=*/true)); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)); +} + TEST_P(WriteInteTest, TestPkSpillableMultiBucketMultiRoundDataCorrectness) { auto dir = UniqueTestDirectory::Create(); arrow::FieldVector fields = {