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()); + auto 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..db732d908 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()); + auto 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..8ee635de7 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()); + auto 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_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/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/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 3e708f87a..f34b25fa9 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -87,12 +87,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()), @@ -102,8 +144,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"}, @@ -117,10 +158,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/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 c5d9b9d2e..7ce6f2e62 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -69,6 +69,7 @@ #include "rapidjson/writer.h" namespace paimon::test { + // This is a sdk end-to-end test demo that supports write, commit, scan, and read operations. class WriteAndReadInteTest : public ::testing::Test, @@ -723,6 +724,189 @@ 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::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, + 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, + 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)); + + 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( + 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) { + 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::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, + 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, + 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)); + + 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(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) { + 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::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, + 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, + 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, + 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)); + + 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(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) { 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..3c6d2111d 100644 --- a/test/inte/write_inte_test.cpp +++ b/test/inte/write_inte_test.cpp @@ -106,6 +106,7 @@ class TableSchema; } // namespace paimon namespace paimon::test { + class WriteInteTest : public testing::Test, public ::testing::WithParamInterface { public: void SetUp() override { @@ -4407,6 +4408,81 @@ 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), + }; + auto data_type = arrow::struct_(fields); + 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)); + }; + + 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)); + 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()); + + 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) { auto dir = UniqueTestDirectory::Create(); arrow::FieldVector fields = {