Skip to content

Commit 05cffcc

Browse files
ChaomingZhangCNChaomingZhangCN
committed
feat(vector): support primary-key tables
1 parent 404e5d9 commit 05cffcc

17 files changed

Lines changed: 717 additions & 36 deletions

‎docs/source/user_guide/data_types.rst‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -194,12 +194,13 @@ and `Arrow DataTypes <https://arrow.apache.org/docs/format/Columnar.html#data-ty
194194
``SMALLINT``, ``INT``, ``BIGINT``, ``FLOAT``, or ``DOUBLE``. A VECTOR
195195
value may be NULL, but its elements cannot be NULL.
196196

197-
Paimon C++ currently supports VECTOR columns only in append-only tables
198-
backed by Parquet data files. They use the standard Parquet LIST
197+
Paimon C++ supports VECTOR value columns in append-only and primary-key
198+
tables backed by Parquet data files. They use the standard Parquet LIST
199199
representation on disk and are restored as Arrow ``FixedSizeList``
200-
values on read. Primary-key tables and data-evolution tables containing
201-
VECTOR fields are rejected. VECTOR columns also cannot be partition or
202-
bucket keys. Dedicated vector storage is not included yet.
200+
values on read. VECTOR is not supported in data-evolution tables. VECTOR
201+
columns cannot be primary, partition, or bucket keys, nor comparator-based
202+
ordering fields such as sequence and sequence-group fields. Dedicated
203+
vector storage is not included yet.
203204

204205
Paimon C++ also reads Parquet files written by Paimon Rust or Python whose
205206
embedded Arrow schema restores VECTOR columns as ``FixedSizeList``,

‎src/paimon/common/data/columnar/columnar_array.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,14 @@ Timestamp ColumnarArray::GetTimestamp(int32_t pos, int32_t precision) const {
6767
}
6868

6969
std::shared_ptr<InternalArray> ColumnarArray::GetArray(int32_t pos) const {
70+
if (array_->type_id() == arrow::Type::FIXED_SIZE_LIST) {
71+
auto fixed_size_list_array = checked_cast<const arrow::FixedSizeListArray*>(array_);
72+
auto fixed_size_list_type =
73+
checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
74+
int32_t offset = static_cast<int32_t>(fixed_size_list_array->value_offset(offset_ + pos));
75+
return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), pool_, offset,
76+
fixed_size_list_type->list_size());
77+
}
7078
auto list_array = checked_cast<const arrow::ListArray*>(array_);
7179
int32_t offset = list_array->value_offset(offset_ + pos);
7280
int32_t length = list_array->value_length(offset_ + pos);

‎src/paimon/common/data/columnar/columnar_row.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,15 @@ std::shared_ptr<InternalRow> ColumnarRow::GetRow(int32_t pos, int32_t num_fields
6666
}
6767

6868
std::shared_ptr<InternalArray> ColumnarRow::GetArray(int32_t pos) const {
69+
if (array_vec_[pos]->type_id() == arrow::Type::FIXED_SIZE_LIST) {
70+
auto fixed_size_list_array =
71+
checked_cast<const arrow::FixedSizeListArray*>(array_vec_[pos]);
72+
auto fixed_size_list_type =
73+
checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
74+
int32_t offset = static_cast<int32_t>(fixed_size_list_array->value_offset(row_id_));
75+
return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), pool_, offset,
76+
fixed_size_list_type->list_size());
77+
}
6978
auto list_array = checked_cast<const arrow::ListArray*>(array_vec_[pos]);
7079
int32_t offset = list_array->value_offset(row_id_);
7180
int32_t length = list_array->value_length(row_id_);

‎src/paimon/common/data/columnar/columnar_row_ref.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,15 @@ std::shared_ptr<InternalRow> ColumnarRowRef::GetRow(int32_t pos, int32_t num_fie
6161
}
6262

6363
std::shared_ptr<InternalArray> ColumnarRowRef::GetArray(int32_t pos) const {
64+
if (ctx_->array_vec[pos]->type_id() == arrow::Type::FIXED_SIZE_LIST) {
65+
auto fixed_size_list_array =
66+
checked_cast<const arrow::FixedSizeListArray*>(ctx_->array_vec[pos].get());
67+
auto fixed_size_list_type =
68+
checked_pointer_cast<arrow::FixedSizeListType>(fixed_size_list_array->type());
69+
int32_t offset = static_cast<int32_t>(fixed_size_list_array->value_offset(row_id_));
70+
return std::make_shared<ColumnarArray>(fixed_size_list_array->values().get(), ctx_->pool,
71+
offset, fixed_size_list_type->list_size());
72+
}
6473
auto list_array = checked_cast<const arrow::ListArray*>(ctx_->array_vec[pos].get());
6574
int32_t offset = list_array->value_offset(row_id_);
6675
int32_t length = list_array->value_length(row_id_);

‎src/paimon/common/data/internal_row.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,8 @@ Result<InternalRow::FieldGetterFunc> InternalRow::CreateFieldGetter(
128128
};
129129
break;
130130
}
131-
case arrow::Type::type::LIST: {
131+
case arrow::Type::type::LIST:
132+
case arrow::Type::type::FIXED_SIZE_LIST: {
132133
field_getter = [field_idx](const InternalRow& row) -> VariantType {
133134
return row.GetArray(field_idx);
134135
};

‎src/paimon/common/data/serializer/binary_serializer_utils.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,7 @@ Result<std::shared_ptr<BinaryArray>> BinarySerializerUtils::WriteBinaryArray(
3131
return binary_array;
3232
}
3333
auto binary_array = std::make_shared<BinaryArray>();
34-
auto list_type = checked_pointer_cast<arrow::ListType>(type);
35-
auto value_type = list_type->value_type();
34+
auto value_type = type->field(0)->type();
3635
// TODO(xinyu.lxy): reuse BinaryWriter
3736
BinaryArrayWriter binary_writer(binary_array.get(), value->Size(),
3837
BinaryArrayWriter::GetElementSize(value_type->id()), pool);
@@ -183,7 +182,8 @@ Status BinarySerializerUtils::WriteBinaryData(const std::shared_ptr<arrow::DataT
183182
}
184183
break;
185184
}
186-
case arrow::Type::type::LIST: {
185+
case arrow::Type::type::LIST:
186+
case arrow::Type::type::FIXED_SIZE_LIST: {
187187
auto internal_array = getter->GetArray(pos);
188188
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<BinaryArray> binary_array,
189189
WriteBinaryArray(internal_array, type, pool));

‎src/paimon/common/data/serializer/row_compacted_serializer.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,8 @@ Result<RowCompactedSerializer::FieldReader> RowCompactedSerializer::CreateFieldR
285285
};
286286
break;
287287
}
288-
case arrow::Type::type::LIST: {
288+
case arrow::Type::type::LIST:
289+
case arrow::Type::type::FIXED_SIZE_LIST: {
289290
field_reader = [](int32_t pos, RowReader* reader) -> Result<VariantType> {
290291
PAIMON_ASSIGN_OR_RAISE(VariantType value, reader->ReadArray());
291292
return value;
@@ -414,7 +415,8 @@ Result<RowCompactedSerializer::FieldWriter> RowCompactedSerializer::CreateFieldW
414415
};
415416
break;
416417
}
417-
case arrow::Type::type::LIST: {
418+
case arrow::Type::type::LIST:
419+
case arrow::Type::type::FIXED_SIZE_LIST: {
418420
field_writer = [field_type](int32_t pos, const VariantType& field,
419421
RowWriter* writer) -> Status {
420422
return writer->WriteArray(

‎src/paimon/core/io/key_value_in_memory_record_reader.cpp‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -111,22 +111,42 @@ void KeyValueInMemoryRecordReader::Close() {
111111
Result<std::shared_ptr<arrow::NumericArray<arrow::UInt64Type>>>
112112
KeyValueInMemoryRecordReader::SortBatch() const {
113113
std::vector<arrow::compute::SortKey> sort_keys;
114+
arrow::FieldVector sort_fields;
115+
arrow::ArrayVector sort_columns;
114116
sort_keys.reserve(primary_keys_.size() + user_defined_sequence_fields_.size());
117+
sort_fields.reserve(primary_keys_.size() + user_defined_sequence_fields_.size());
118+
sort_columns.reserve(primary_keys_.size() + user_defined_sequence_fields_.size());
119+
const arrow::StructType* value_type = value_struct_array_->struct_type();
120+
auto append_sort_key = [&](const std::string& name, arrow::compute::SortOrder order) -> Status {
121+
int32_t field_index = value_type->GetFieldIndex(name);
122+
if (field_index < 0) {
123+
return Status::Invalid(fmt::format("cannot find field {} in data batch", name));
124+
}
125+
sort_keys.emplace_back(name, order);
126+
sort_fields.push_back(value_type->field(field_index));
127+
sort_columns.push_back(value_struct_array_->field(field_index));
128+
return Status::OK();
129+
};
115130
for (const auto& name : primary_keys_) {
116-
sort_keys.emplace_back(name, arrow::compute::SortOrder::Ascending);
131+
PAIMON_RETURN_NOT_OK(append_sort_key(name, arrow::compute::SortOrder::Ascending));
117132
}
118133
const auto sequence_sort_order = sequence_fields_ascending_
119134
? arrow::compute::SortOrder::Ascending
120135
: arrow::compute::SortOrder::Descending;
121136
for (const auto& name : user_defined_sequence_fields_) {
122-
sort_keys.emplace_back(name, sequence_sort_order);
137+
PAIMON_RETURN_NOT_OK(append_sort_key(name, sequence_sort_order));
123138
}
124139
auto sort_options =
125140
arrow::compute::SortOptions(sort_keys, arrow::compute::NullPlacement::AtStart);
126141
arrow::compute::ExecContext exec_context(arrow_pool_.get());
127-
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> sorted_indices,
128-
arrow::compute::SortIndices(arrow::Datum(value_struct_array_),
129-
sort_options, &exec_context));
142+
// Arrow's StructArray sorting path may inspect value columns outside the sort keys. Restrict
143+
// the batch to the requested fields so non-sortable values, such as VECTOR, are never compared.
144+
std::shared_ptr<arrow::RecordBatch> sort_batch =
145+
arrow::RecordBatch::Make(arrow::schema(std::move(sort_fields)),
146+
value_struct_array_->length(), std::move(sort_columns));
147+
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
148+
std::shared_ptr<arrow::Array> sorted_indices,
149+
arrow::compute::SortIndices(arrow::Datum(sort_batch), sort_options, &exec_context));
130150
if (!sorted_indices || sorted_indices->type_id() != arrow::Type::UINT64) {
131151
return Status::Invalid("cannot cast sorted indices to UInt64Array");
132152
}

‎src/paimon/core/io/row_to_arrow_array_converter.h‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,12 @@ Status RowToArrowArrayConverter<T, R>::Reserve(arrow::ArrayBuilder* array_builde
158158
PAIMON_RETURN_NOT_OK(Reserve(list_builder->value_builder(), idx));
159159
break;
160160
}
161+
case arrow::Type::type::FIXED_SIZE_LIST: {
162+
PAIMON_ASSIGN_OR_RAISE(auto* list_builder,
163+
CastToTypedBuilder<arrow::FixedSizeListBuilder>(array_builder));
164+
PAIMON_RETURN_NOT_OK(Reserve(list_builder->value_builder(), idx));
165+
break;
166+
}
161167
case arrow::Type::type::MAP: {
162168
PAIMON_ASSIGN_OR_RAISE(auto* map_builder,
163169
CastToTypedBuilder<arrow::MapBuilder>(array_builder));
@@ -224,6 +230,11 @@ Status RowToArrowArrayConverter<T, R>::Accumulate(const arrow::Array* array, int
224230
PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx));
225231
break;
226232
}
233+
case arrow::Type::type::FIXED_SIZE_LIST: {
234+
auto list_array = checked_cast<const arrow::FixedSizeListArray*>(array);
235+
PAIMON_RETURN_NOT_OK(Accumulate(list_array->values().get(), idx));
236+
break;
237+
}
227238
case arrow::Type::type::MAP: {
228239
auto map_array = checked_cast<const arrow::MapArray*>(array);
229240
PAIMON_RETURN_NOT_OK(Accumulate(map_array->keys().get(), idx));
@@ -432,6 +443,31 @@ RowToArrowArrayConverter<T, R>::AppendField(bool use_view, arrow::ArrayBuilder*
432443
return arrow::Status::OK();
433444
});
434445
}
446+
case arrow::Type::type::FIXED_SIZE_LIST: {
447+
PAIMON_ASSIGN_OR_RAISE(auto* list_builder,
448+
CastToTypedBuilder<arrow::FixedSizeListBuilder>(array_builder));
449+
std::shared_ptr<arrow::FixedSizeListType> list_type =
450+
checked_pointer_cast<arrow::FixedSizeListType>(list_builder->type());
451+
int32_t list_size = list_type->list_size();
452+
PAIMON_ASSIGN_OR_RAISE(AppendValueFunc value_func,
453+
(RowToArrowArrayConverter<T, R>::AppendField(
454+
use_view, list_builder->value_builder(), reserve_count)));
455+
return RowToArrowArrayConverter<T, R>::AppendValueFunc(
456+
[list_builder, list_size, value_func](const DataGetters& data_getter,
457+
int32_t pos) -> arrow::Status {
458+
CHECK_AND_APPEND_NULL(data_getter, list_builder, pos);
459+
std::shared_ptr<InternalArray> sub_array = data_getter.GetArray(pos);
460+
if (!sub_array || sub_array->Size() != list_size) {
461+
return arrow::Status::Invalid(
462+
"VECTOR length does not match its declared dimension");
463+
}
464+
ARROW_RETURN_NOT_OK(list_builder->Append());
465+
for (int32_t i = 0; i < list_size; ++i) {
466+
ARROW_RETURN_NOT_OK(value_func(*sub_array, i));
467+
}
468+
return arrow::Status::OK();
469+
});
470+
}
435471
case arrow::Type::type::MAP: {
436472
PAIMON_ASSIGN_OR_RAISE(auto* map_builder,
437473
CastToTypedBuilder<arrow::MapBuilder>(array_builder));

‎src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,12 +73,12 @@ Result<bool> EqualRows(const std::shared_ptr<InternalRow>& lhs,
7373

7474
Result<bool> EqualArrays(const std::shared_ptr<InternalArray>& lhs,
7575
const std::shared_ptr<InternalArray>& rhs,
76-
const std::shared_ptr<arrow::ListType>& type) {
76+
const std::shared_ptr<arrow::DataType>& type) {
7777
if (!lhs || !rhs || lhs->Size() != rhs->Size()) {
7878
return lhs == rhs;
7979
}
8080
for (int32_t i = 0; i < lhs->Size(); ++i) {
81-
PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->value_type()));
81+
PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->field(0)->type()));
8282
if (!equal) {
8383
return false;
8484
}
@@ -174,6 +174,7 @@ Result<VariantType> FieldAggregateUtils::GetValue(const DataGetters& getters, in
174174
getters.GetDecimal(pos, decimal_type->precision(), decimal_type->scale()));
175175
}
176176
case arrow::Type::LIST:
177+
case arrow::Type::FIXED_SIZE_LIST:
177178
return VariantType(getters.GetArray(pos));
178179
case arrow::Type::MAP:
179180
return VariantType(getters.GetMap(pos));
@@ -229,9 +230,10 @@ Result<bool> FieldAggregateUtils::Equals(const VariantType& lhs, const VariantTy
229230
DataDefine::GetVariantValue<std::shared_ptr<InternalRow>>(rhs),
230231
checked_pointer_cast<arrow::StructType>(type));
231232
case arrow::Type::LIST:
233+
case arrow::Type::FIXED_SIZE_LIST:
232234
return EqualArrays(DataDefine::GetVariantValue<std::shared_ptr<InternalArray>>(lhs),
233235
DataDefine::GetVariantValue<std::shared_ptr<InternalArray>>(rhs),
234-
checked_pointer_cast<arrow::ListType>(type));
236+
type);
235237
case arrow::Type::MAP:
236238
return EqualMaps(DataDefine::GetVariantValue<std::shared_ptr<InternalMap>>(lhs),
237239
DataDefine::GetVariantValue<std::shared_ptr<InternalMap>>(rhs),

0 commit comments

Comments
 (0)