feat(vector): support VECTOR as a non-key value column in Parquet-backed primary-key tables. (#275)
diff --git a/docs/source/user_guide/data_types.rst b/docs/source/user_guide/data_types.rst
index add537a..57f0eaf 100644
--- a/docs/source/user_guide/data_types.rst
+++ b/docs/source/user_guide/data_types.rst
@@ -194,12 +194,13 @@
``SMALLINT``, ``INT``, ``BIGINT``, ``FLOAT``, or ``DOUBLE``. A VECTOR
value may be NULL, but its elements cannot be NULL.
- Paimon C++ currently supports VECTOR columns only in append-only tables
- backed by Parquet data files. They use the standard Parquet LIST
+ Paimon C++ supports VECTOR value columns in append-only and primary-key
+ tables backed by Parquet data files. They use the standard Parquet LIST
representation on disk and are restored as Arrow ``FixedSizeList``
- values on read. Primary-key tables and data-evolution tables containing
- VECTOR fields are rejected. VECTOR columns also cannot be partition or
- bucket keys. Dedicated vector storage is not included yet.
+ values on read. VECTOR is not supported in data-evolution tables. VECTOR
+ columns cannot be primary, partition, or bucket keys, nor comparator-based
+ ordering fields such as sequence and sequence-group fields. Dedicated
+ vector storage is not included yet.
Paimon C++ also reads Parquet files written by Paimon Rust or Python whose
embedded Arrow schema restores VECTOR columns as ``FixedSizeList``,
diff --git a/src/paimon/common/data/columnar/columnar_array.cpp b/src/paimon/common/data/columnar/columnar_array.cpp
index 6566aea..11a80df 100644
--- a/src/paimon/common/data/columnar/columnar_array.cpp
+++ b/src/paimon/common/data/columnar/columnar_array.cpp
@@ -67,6 +67,14 @@
}
std::shared_ptr<InternalArray> ColumnarArray::GetArray(int32_t pos) const {
+ if (array_->type_id() == arrow::Type::FIXED_SIZE_LIST) {
+ auto fixed_size_list_array = checked_cast<const arrow::FixedSizeListArray*>(array_);
+ auto fixed_size_list_type =
+ checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
+ auto offset = static_cast<int32_t>(fixed_size_list_array->value_offset(offset_ + pos));
+ return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), pool_, offset,
+ fixed_size_list_type->list_size());
+ }
auto list_array = checked_cast<const arrow::ListArray*>(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 5099b42..db732d9 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<InternalArray> ColumnarRow::GetArray(int32_t pos) const {
+ if (array_vec_[pos]->type_id() == arrow::Type::FIXED_SIZE_LIST) {
+ auto fixed_size_list_array =
+ checked_cast<const arrow::FixedSizeListArray*>(array_vec_[pos]);
+ auto fixed_size_list_type =
+ checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
+ auto offset = static_cast<int32_t>(fixed_size_list_array->value_offset(row_id_));
+ return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), pool_, offset,
+ fixed_size_list_type->list_size());
+ }
auto list_array = checked_cast<const arrow::ListArray*>(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 d73e6b7..8ee635d 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<InternalArray> 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<const arrow::FixedSizeListArray*>(ctx_->array_vec[pos].get());
+ auto fixed_size_list_type =
+ checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
+ auto offset = static_cast<int32_t>(fixed_size_list_array->value_offset(row_id_));
+ return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), ctx_->pool,
+ offset, fixed_size_list_type->list_size());
+ }
auto list_array = checked_cast<const arrow::ListArray*>(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 24b6b5b..b7c011b 100644
--- a/src/paimon/common/data/internal_row.cpp
+++ b/src/paimon/common/data/internal_row.cpp
@@ -128,7 +128,8 @@
};
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 c0248cd..d72a2b0 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 @@
return binary_array;
}
auto binary_array = std::make_shared<BinaryArray>();
- auto list_type = checked_pointer_cast<arrow::ListType>(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 @@
}
break;
}
- case arrow::Type::type::LIST: {
+ case arrow::Type::type::LIST:
+ case arrow::Type::type::FIXED_SIZE_LIST: {
auto internal_array = getter->GetArray(pos);
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<BinaryArray> 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 411dae6..1fd618a 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 @@
};
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<VariantType> {
PAIMON_ASSIGN_OR_RAISE(VariantType value, reader->ReadArray());
return value;
@@ -414,7 +415,8 @@
};
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 7604064..ca6f331 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 @@
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::StructArray>(
+ 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<FieldsComparator> key_comparator,
+ FieldsComparator::Create({DataField(0, arrow::field("pk", arrow::int32()))},
+ /*is_ascending_order=*/true));
+
+ auto check_sort = [&](const std::vector<std::string>& sequence_fields, bool ascending,
+ const std::vector<int64_t>& expected_indices) {
+ KeyValueInMemoryRecordReader reader(
+ /*last_sequence_num=*/0, src_array, std::vector<RecordBatch::RowKind>{},
+ std::vector<std::string>{"pk"}, sequence_fields, ascending, key_comparator, pool_);
+ ASSERT_OK_AND_ASSIGN(
+ std::vector<KeyValue> results,
+ (ReadResultCollector::CollectKeyValueResult<KeyValueInMemoryRecordReader,
+ KeyValueRecordReader::Iterator>(&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<float> expected_vector = {static_cast<float>(source_index),
+ static_cast<float>(source_index + 1),
+ static_cast<float>(source_index + 2)};
+ ASSERT_OK_AND_ASSIGN(std::vector<float> 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<float> 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 6caa9f6..336d985 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 @@
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<arrow::FixedSizeListBuilder>(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<arrow::MapBuilder>(array_builder));
@@ -224,6 +230,11 @@
PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx));
break;
}
+ case arrow::Type::type::FIXED_SIZE_LIST: {
+ auto list_array = checked_cast<const arrow::FixedSizeListArray*>(array);
+ PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx));
+ break;
+ }
case arrow::Type::type::MAP: {
auto map_array = checked_cast<const arrow::MapArray*>(array);
PAIMON_RETURN_NOT_OK(Accumulate(map_array->keys().get(), idx));
@@ -432,6 +443,31 @@
return arrow::Status::OK();
});
}
+ case arrow::Type::type::FIXED_SIZE_LIST: {
+ PAIMON_ASSIGN_OR_RAISE(auto* list_builder,
+ CastToTypedBuilder<arrow::FixedSizeListBuilder>(array_builder));
+ std::shared_ptr<arrow::FixedSizeListType> list_type =
+ checked_pointer_cast<arrow::FixedSizeListType>(list_builder->type());
+ int32_t list_size = list_type->list_size();
+ PAIMON_ASSIGN_OR_RAISE(AppendValueFunc value_func,
+ (RowToArrowArrayConverter<T, R>::AppendField(
+ use_view, list_builder->value_builder(), reserve_count)));
+ return RowToArrowArrayConverter<T, R>::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<InternalArray> 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<arrow::MapBuilder>(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 9397ca4..22e02f1 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<bool> EqualArrays(const std::shared_ptr<InternalArray>& lhs,
const std::shared_ptr<InternalArray>& rhs,
- const std::shared_ptr<arrow::ListType>& type) {
+ const std::shared_ptr<arrow::DataType>& 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 @@
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 @@
DataDefine::GetVariantValue<std::shared_ptr<InternalRow>>(rhs),
checked_pointer_cast<arrow::StructType>(type));
case arrow::Type::LIST:
+ case arrow::Type::FIXED_SIZE_LIST:
return EqualArrays(DataDefine::GetVariantValue<std::shared_ptr<InternalArray>>(lhs),
DataDefine::GetVariantValue<std::shared_ptr<InternalArray>>(rhs),
- checked_pointer_cast<arrow::ListType>(type));
+ type);
case arrow::Type::MAP:
return EqualMaps(DataDefine::GetVariantValue<std::shared_ptr<InternalMap>>(lhs),
DataDefine::GetVariantValue<std::shared_ptr<InternalMap>>(rhs),
diff --git a/src/paimon/core/mergetree/compact/internal_row_equalizer.h b/src/paimon/core/mergetree/compact/internal_row_equalizer.h
index 22f15d8..a0a0b4c 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 @@
.CompareTo(rhs.GetDecimal(rhs_pos, precision, scale)) == 0;
});
}
- case arrow::Type::LIST: {
- std::shared_ptr<arrow::ListType> list_type =
- checked_pointer_cast<arrow::ListType>(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 5f53d59..272d34b 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 @@
PAIMON_ASSIGN_OR_RAISE(int64_t value_mem, EstimateMemoryUse(list_array->values()));
return null_bits_size_in_bytes + value_mem;
}
+ case arrow::Type::type::FIXED_SIZE_LIST: {
+ auto list_array = checked_cast<const arrow::FixedSizeListArray*>(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<const arrow::MapArray*>(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 5e524a3..26d4100 100644
--- a/src/paimon/core/schema/schema_validation.cpp
+++ b/src/paimon/core/schema/schema_validation.cpp
@@ -127,6 +127,16 @@
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<arrow::Field>& field) {
@@ -378,6 +388,8 @@
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 @@
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<std::string> agg_func,
options.GetFieldAggFunc(field));
@@ -776,10 +789,6 @@
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 3e708f8..f34b25f 100644
--- a/src/paimon/core/schema/schema_validation_test.cpp
+++ b/src/paimon/core/schema/schema_validation_test.cpp
@@ -87,12 +87,54 @@
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<std::string, std::string> 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<std::string, std::string> 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<std::string, std::string> 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 @@
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<std::string, std::string> data_evolution_options = {
{Options::BUCKET, "-1"},
@@ -117,10 +158,10 @@
/*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 95adf74..03e55c2 100644
--- a/test/inte/realtime_write_inte_test.cpp
+++ b/test/inte/realtime_write_inte_test.cpp
@@ -1164,6 +1164,35 @@
read_array.reset();
}
+TEST_F(RealtimeWriteInteTest, TestPkVector) {
+ options_[Options::FILE_FORMAT] = "parquet";
+ std::shared_ptr<arrow::DataType> 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<RealtimeContext> realtime_context,
+ RealtimeContext::Create());
+ ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
+ CreateRealtimeWriter(realtime_context));
+ ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> 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> 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<RealtimeContext> realtime_context,
diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp
index c5d9b9d..7ce6f2e 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 @@
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<std::string, std::string> 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<RecordBatch> 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<RecordBatch> 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<std::shared_ptr<Split>> 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<std::string, std::string> 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<RecordBatch> 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<RecordBatch> 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<std::shared_ptr<Split>> 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<std::string, std::string> 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<RecordBatch> 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<RecordBatch> 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<RecordBatch> 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<std::shared_ptr<Split>> 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 27bcea6..3c6d211 100644
--- a/test/inte/write_inte_test.cpp
+++ b/test/inte/write_inte_test.cpp
@@ -106,6 +106,7 @@
} // namespace paimon
namespace paimon::test {
+
class WriteInteTest : public testing::Test, public ::testing::WithParamInterface<std::string> {
public:
void SetUp() override {
@@ -4407,6 +4408,81 @@
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<std::string, std::string> 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<WriteContext> 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<arrow::Array>& array) {
+ ArrowArray c_array;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array));
+ auto batch = std::make_unique<RecordBatch>(std::map<std::string, std::string>{{"pt", "10"}},
+ /*bucket=*/0,
+ std::vector<RecordBatch::RowKind>{}, &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 = {
