From a14a841f96ed9c79616417be11588b439d870afc Mon Sep 17 00:00:00 2001 From: "zhangchaoming.zcm" Date: Wed, 2 Sep 2026 15:53:05 +0800 Subject: [PATCH 1/2] feat(vector): support primary-key tables --- docs/source/user_guide/data_types.rst | 11 +- .../common/data/columnar/columnar_array.cpp | 8 + .../common/data/columnar/columnar_row.cpp | 9 + .../common/data/columnar/columnar_row_ref.cpp | 9 + src/paimon/common/data/internal_row.cpp | 3 +- .../serializer/binary_serializer_utils.cpp | 6 +- .../serializer/row_compacted_serializer.cpp | 6 +- .../io/key_value_in_memory_record_reader.cpp | 30 +- .../core/io/row_to_arrow_array_converter.h | 36 ++ .../aggregate/field_aggregate_utils.cpp | 8 +- .../compact/internal_row_equalizer.h | 7 +- .../core/mergetree/in_memory_sort_buffer.cpp | 5 + .../realtime/realtime_primary_key_writer.cpp | 21 +- src/paimon/core/schema/schema_validation.cpp | 17 +- .../core/schema/schema_validation_test.cpp | 57 ++- test/inte/write_and_read_inte_test.cpp | 373 ++++++++++++++++++ test/inte/write_inte_test.cpp | 147 +++++++ 17 files changed, 717 insertions(+), 36 deletions(-) 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 c3367b62a..3a07965b4 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -77,17 +77,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 c5d9b9d2e..6534fa2e8 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, @@ -723,6 +898,204 @@ TEST_P(WriteAndReadInteTest, TestPKListAggPreservesResultsAcrossKeys) { 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 27bcea674..5b7e36638 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 { @@ -4407,6 +4468,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 = { From 031313422e85ba4144dd97a4aba6eb68196be763 Mon Sep 17 00:00:00 2001 From: "zhangchaoming.zcm" Date: Mon, 7 Sep 2026 11:31:22 +0800 Subject: [PATCH 2/2] fix(vector): address PK table review feedback --- .../io/key_value_in_memory_record_reader.cpp | 30 +- ...key_value_in_memory_record_reader_test.cpp | 63 ++++ .../realtime/realtime_primary_key_writer.cpp | 21 +- test/inte/realtime_write_inte_test.cpp | 29 ++ test/inte/write_and_read_inte_test.cpp | 291 +++--------------- test/inte/write_inte_test.cpp | 97 +----- 6 files changed, 162 insertions(+), 369 deletions(-) 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 f4d3e8870..9918b328c 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,42 +111,22 @@ 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_) { - PAIMON_RETURN_NOT_OK(append_sort_key(name, arrow::compute::SortOrder::Ascending)); + sort_keys.emplace_back(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_) { - PAIMON_RETURN_NOT_OK(append_sort_key(name, sequence_sort_order)); + sort_keys.emplace_back(name, sequence_sort_order); } auto sort_options = arrow::compute::SortOptions(sort_keys, arrow::compute::NullPlacement::AtStart); arrow::compute::ExecContext exec_context(arrow_pool_.get()); - // 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)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr sorted_indices, + arrow::compute::SortIndices(arrow::Datum(value_struct_array_), + 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/key_value_in_memory_record_reader_test.cpp b/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp index 7604064ee..ca6f331e9 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp +++ b/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp @@ -26,6 +26,7 @@ #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/data/internal_array.h" #include "paimon/common/types/data_field.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/fields_comparator.h" @@ -349,6 +350,68 @@ TEST_F(KeyValueInMemoryRecordReaderTest, TestStableSortWithDuplicateKeys) { ASSERT_FALSE(eof_iter); } +TEST_F(KeyValueInMemoryRecordReaderTest, TestSortWithVectorValues) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto src_type = arrow::struct_( + {arrow::field("pk", arrow::int32()), arrow::field("seq", arrow::int32()), + arrow::field("embedding", vector_type), + arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_type)}))}); + auto src_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(src_type, R"([ + [2, 20, [0, 1, 2], [[0, 1, 2]]], + [1, 20, null, null], + [2, 10, [2, 3, 4], [null]], + [1, null, [3, 4, 5], [[3, 4, 5]]], + [2, 20, [4, 5, 6], [[4, 5, 6]]] + ])") + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr key_comparator, + FieldsComparator::Create({DataField(0, arrow::field("pk", arrow::int32()))}, + /*is_ascending_order=*/true)); + + auto check_sort = [&](const std::vector& sequence_fields, bool ascending, + const std::vector& expected_indices) { + KeyValueInMemoryRecordReader reader( + /*last_sequence_num=*/0, src_array, std::vector{}, + std::vector{"pk"}, sequence_fields, ascending, key_comparator, pool_); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(&reader))); + ASSERT_EQ(expected_indices.size(), results.size()); + for (size_t i = 0; i < results.size(); ++i) { + const int64_t source_index = expected_indices[i]; + SCOPED_TRACE(source_index); + ASSERT_EQ(source_index, results[i].sequence_number); + const auto& value = results[i].value; + ASSERT_EQ(source_index == 1 || source_index == 3 ? 1 : 2, results[i].key->GetInt(0)); + ASSERT_EQ(source_index == 1, value->IsNullAt(2)); + ASSERT_EQ(source_index == 1, value->IsNullAt(3)); + if (source_index == 1) { + continue; + } + const std::vector expected_vector = {static_cast(source_index), + static_cast(source_index + 1), + static_cast(source_index + 2)}; + ASSERT_OK_AND_ASSIGN(std::vector vector, value->GetArray(2)->ToFloatArray()); + ASSERT_EQ(expected_vector, vector); + auto payload = value->GetRow(3, 1); + ASSERT_EQ(source_index == 2, payload->IsNullAt(0)); + if (source_index != 2) { + ASSERT_OK_AND_ASSIGN(std::vector nested_vector, + payload->GetArray(0)->ToFloatArray()); + ASSERT_EQ(expected_vector, nested_vector); + } + } + }; + // Equal sort keys retain input order regardless of the non-sortable value columns. + check_sort({}, /*ascending=*/true, {1, 3, 0, 2, 4}); + check_sort({"seq"}, /*ascending=*/true, {3, 1, 2, 0, 4}); + check_sort({"seq"}, /*ascending=*/false, {3, 1, 0, 4, 2}); +} + TEST_F(KeyValueInMemoryRecordReaderTest, TestVariantType) { // test null, repeated key, sequence fields, out of order in variant data type // precondition: fields[0] is key fields, fields[1] is sequence field, src_array are like: diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 3a07965b4..c3367b62a 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -77,36 +77,17 @@ 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(sort_batch), options, &context)); + arrow::compute::SortIndices(arrow::Datum(transport), options, &context)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( arrow::Datum sorted, arrow::compute::Take(arrow::Datum(transport), indices, diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 95adf74fc..03e55c24f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1164,6 +1164,35 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { read_array.reset(); } +TEST_F(RealtimeWriteInteTest, TestPkVector) { + options_[Options::FILE_FORMAT] = "parquet"; + std::shared_ptr vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + fields_ = {arrow::field("id", arrow::int64()), arrow::field("embedding", vector_type)}; + schema_ = arrow::schema(fields_); + CreatePkTable(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeUnpartitionedBatchFromJson(R"([ + [0, 2, [2.0, 2.0, 2.0]], + [1, 3, null], + [2, 1, [1.0, 1.0, 1.0]], + [3, 2, [3.0, 3.0, 3.0]] + ])")); + ASSERT_OK(writer->Write(std::move(batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadPlanWithSchemaAndCheck(plan, realtime_context, schema_, + R"([[0, 1, [1.0, 1.0, 1.0]], + [0, 2, [3.0, 3.0, 3.0]], + [0, 3, null]])"); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkRealtimeReadOptimizedScanUnsupported) { CreatePkTable(); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index 6534fa2e8..7ce6f2e62 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -17,7 +17,6 @@ */ #include -#include #include #include #include @@ -70,179 +69,6 @@ #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 @@ -911,8 +737,9 @@ TEST_P(WriteAndReadInteTest, TestPKVector) { 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_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, + {Options::BUCKET, "1"}, {Options::FILE_SYSTEM, file_system}, }; if (file_system == "jindo") { @@ -923,17 +750,18 @@ TEST_P(WriteAndReadInteTest, TestPKVector) { 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_AND_ASSIGN( + std::unique_ptr initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([[1, [1.0, 2.0, 3.0]], [2, null]])", + /*partition_map=*/{}, /*bucket=*/0, {})); 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_AND_ASSIGN( + std::unique_ptr update_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([[1, [4.0, 5.0, 6.0]], [3, [7.0, 8.0, 9.0]]])", + /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/1, /*expected_commit_messages=*/std::nullopt)); @@ -944,16 +772,11 @@ TEST_P(WriteAndReadInteTest, TestPKVector) { 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)); + bool success, helper->ReadAndCheckResult( + arrow::struct_(result_fields), data_splits, + R"([[0, 1, [4.0, 5.0, 6.0]], [0, 2, null], [0, 3, [7.0, 8.0, 9.0]]])")); + ASSERT_TRUE(success); } TEST_P(WriteAndReadInteTest, TestPKNestedVector) { @@ -970,8 +793,9 @@ TEST_P(WriteAndReadInteTest, TestPKNestedVector) { 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_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, + {Options::BUCKET, "1"}, {Options::FILE_SYSTEM, file_system}, }; if (file_system == "jindo") { @@ -984,21 +808,18 @@ TEST_P(WriteAndReadInteTest, TestPKNestedVector) { 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}})); + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([[1, [[1.0, 2.0, 3.0], "initial"]], [2, [null, "null-vector"]], [3, null]])", + /*partition_map=*/{}, /*bucket=*/0, {})); 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_AND_ASSIGN(std::unique_ptr update_batch, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([[1, [[4.0, 5.0, 6.0], "updated"]], [2, [[7.0, 8.0, 9.0], null]]])", + /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/1, /*expected_commit_messages=*/std::nullopt)); @@ -1009,19 +830,13 @@ TEST_P(WriteAndReadInteTest, TestPKNestedVector) { 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)); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(arrow::struct_(result_fields), data_splits, R"([ + [0, 1, [[4.0, 5.0, 6.0], "updated"]], + [0, 2, [[7.0, 8.0, 9.0], null]], + [0, 3, null] + ])")); + ASSERT_TRUE(success); } TEST_P(WriteAndReadInteTest, TestPKVectorWithListagg) { @@ -1038,7 +853,6 @@ TEST_P(WriteAndReadInteTest, TestPKVectorWithListagg) { 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"}, @@ -1057,23 +871,24 @@ TEST_P(WriteAndReadInteTest, TestPKVectorWithListagg) { 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")}})); + TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([[1, [1.0, 2.0, 3.0], "alpha"], [2, null, "one"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); 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_AND_ASSIGN(std::unique_ptr update_batch, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([[1, [4.0, 5.0, 6.0], "beta"], [2, [7.0, 8.0, 9.0], "two"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); 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")}})); + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([[1, null, "gamma"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); ASSERT_OK(helper->WriteAndCommit(std::move(null_vector_batch), /*commit_identifier=*/2, /*expected_commit_messages=*/std::nullopt)); @@ -1084,16 +899,12 @@ TEST_P(WriteAndReadInteTest, TestPKVectorWithListagg) { 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)); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(arrow::struct_(result_fields), data_splits, R"([ + [0, 1, [4.0, 5.0, 6.0], "alpha,beta,gamma"], + [0, 2, [7.0, 8.0, 9.0], "one,two"] + ])")); + ASSERT_TRUE(success); } TEST_P(WriteAndReadInteTest, TestInputChangelogStreamRead) { diff --git a/test/inte/write_inte_test.cpp b/test/inte/write_inte_test.cpp index 5b7e36638..3c6d2111d 100644 --- a/test/inte/write_inte_test.cpp +++ b/test/inte/write_inte_test.cpp @@ -15,7 +15,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include #include #include #include @@ -107,65 +106,6 @@ 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: @@ -4482,6 +4422,7 @@ TEST_P(WriteInteTest, TestPkSpillableVector) { arrow::field("pt", arrow::int32()), arrow::field("embedding", vector_type), }; + auto data_type = arrow::struct_(fields); std::map options = { {Options::FILE_FORMAT, file_format}, {Options::BUCKET, "1"}, @@ -4513,15 +4454,15 @@ TEST_P(WriteInteTest, TestPkSpillableVector) { 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)); + auto batch1 = + arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([["Alice", 10, [1.0, 2.0, 3.0]]])") + .ValueOrDie(); + auto batch2 = + arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([["Bob", 10, [4.0, 5.0, 6.0]]])") + .ValueOrDie(); + auto batch3 = + arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([["Alice", 10, [7.0, 8.0, 9.0]]])") + .ValueOrDie(); ASSERT_OK(write_array(file_store_write.get(), batch1)); ASSERT_EQ(1, TestHelper::CountChannelFiles(file_system_, tmp_dir)); @@ -4537,21 +4478,9 @@ TEST_P(WriteInteTest, TestPkSpillableVector) { 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)); + ASSERT_OK(ScanAndVerifyResult(table_path, fields, + R"([[0, "Alice", 10, [7.0, 8.0, 9.0]], + [0, "Bob", 10, [4.0, 5.0, 6.0]]])")); } TEST_P(WriteInteTest, TestPkSpillableMultiBucketMultiRoundDataCorrectness) {