Skip to content

Commit 32e919c

Browse files
lxy-9602lxy-9602
authored
feat(lumina): support null values in vector column during index building (alibaba#310)
1 parent 433487e commit 32e919c

4 files changed

Lines changed: 285 additions & 28 deletions

File tree

‎src/paimon/global_index/lumina/lumina_global_index.cpp‎

Lines changed: 51 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -177,8 +177,12 @@ Result<std::shared_ptr<GlobalIndexReader>> LuminaGlobalIndex::CreateReader(
177177
class LuminaDataset : public ::lumina::api::Dataset {
178178
public:
179179
LuminaDataset(int64_t element_count, uint32_t dimension,
180-
const std::vector<std::shared_ptr<arrow::FloatArray>>& array_vec)
181-
: element_count_(element_count), dimension_(dimension), array_vec_(array_vec) {}
180+
const std::vector<std::shared_ptr<arrow::FloatArray>>& array_vec,
181+
const std::vector<int64_t>& start_ids)
182+
: element_count_(element_count),
183+
dimension_(dimension),
184+
array_vec_(array_vec),
185+
start_ids_(start_ids) {}
182186

183187
uint32_t Dim() const noexcept override {
184188
return dimension_;
@@ -200,8 +204,8 @@ class LuminaDataset : public ::lumina::api::Dataset {
200204
vector_buffer.resize(value_array_length);
201205
memcpy(vector_buffer.data(), value_ptr, sizeof(float) * value_array_length);
202206
id_buffer.resize(element_count);
203-
std::iota(id_buffer.begin(), id_buffer.end(), id_);
204-
id_ += element_count;
207+
std::iota(id_buffer.begin(), id_buffer.end(),
208+
static_cast<::lumina::core::vector_id_t>(start_ids_[cursor_]));
205209

206210
// release the array when copy to vector_buffer
207211
value_array.reset();
@@ -213,8 +217,8 @@ class LuminaDataset : public ::lumina::api::Dataset {
213217
int64_t element_count_;
214218
uint32_t dimension_;
215219
std::vector<std::shared_ptr<arrow::FloatArray>> array_vec_;
220+
std::vector<int64_t> start_ids_;
216221
size_t cursor_ = 0;
217-
::lumina::core::vector_id_t id_ = 0;
218222
};
219223

220224
LuminaIndexWriter::LuminaIndexWriter(const std::string& field_name,
@@ -254,35 +258,62 @@ Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array,
254258
auto list_field_array = std::dynamic_pointer_cast<arrow::ListArray>(field_array);
255259
CHECK_NOT_NULL(list_field_array,
256260
"invalid input array in LuminaIndexWriter, field array must be list array");
257-
auto value_array = std::dynamic_pointer_cast<arrow::FloatArray>(list_field_array->values());
258-
CHECK_NOT_NULL(
259-
value_array,
260-
"invalid input array in LuminaIndexWriter, field value array must be float array");
261-
if (value_array->null_count() != 0) {
262-
return Status::Invalid("field value array in LuminaIndexWriter is invalid, must not null");
263-
}
264-
if (value_array->length() != field_length * dimension_) {
265-
return Status::Invalid(fmt::format(
266-
"invalid input array in LuminaIndexWriter, length of field array [{}] multiplied "
267-
"dimension [{}] must match length of field value array [{}]",
268-
field_length, dimension_, value_array->length()));
261+
262+
// Split into contiguous non-null segments, skipping null rows in the list field.
263+
int64_t segment_start = -1;
264+
for (int64_t i = 0; i <= field_length; i++) {
265+
bool is_null = (i < field_length) && list_field_array->IsNull(i);
266+
bool is_end = (i == field_length);
267+
268+
if (!is_null && !is_end && segment_start == -1) {
269+
segment_start = i;
270+
}
271+
272+
if ((is_null || is_end) && segment_start != -1) {
273+
int64_t segment_len = i - segment_start;
274+
// Use value_offset to precisely locate the float range for this segment
275+
auto value_start_offset = list_field_array->value_offset(segment_start);
276+
auto value_end_offset = list_field_array->value_offset(segment_start + segment_len);
277+
int64_t value_length = value_end_offset - value_start_offset;
278+
auto sliced_values = std::dynamic_pointer_cast<arrow::FloatArray>(
279+
list_field_array->values()->Slice(value_start_offset, value_length));
280+
CHECK_NOT_NULL(sliced_values,
281+
"invalid sliced value array in LuminaIndexWriter, must be float array");
282+
if (sliced_values->null_count() != 0) {
283+
return Status::Invalid(
284+
"field value array in LuminaIndexWriter is invalid, must not null");
285+
}
286+
if (sliced_values->length() != segment_len * static_cast<int64_t>(dimension_)) {
287+
return Status::Invalid(fmt::format(
288+
"invalid input array in LuminaIndexWriter, length of field array [{}] "
289+
"multiplied dimension [{}] must match length of field value array [{}]",
290+
segment_len, dimension_, sliced_values->length()));
291+
}
292+
array_vec_.push_back(std::move(sliced_values));
293+
array_start_ids_.push_back(count_ + segment_start);
294+
indexed_count_ += segment_len;
295+
segment_start = -1;
296+
}
269297
}
298+
270299
count_ += array->length();
271-
array_vec_.push_back(std::move(value_array));
272300
return Status::OK();
273301
}
274302

275303
Result<std::vector<GlobalIndexIOMeta>> LuminaIndexWriter::Finish() {
304+
if (indexed_count_ == 0) {
305+
return std::vector<GlobalIndexIOMeta>();
306+
}
276307
::lumina::core::MemoryResourceConfig memory_resource(pool_.get());
277308
PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA(
278309
::lumina::api::LuminaBuilder builder,
279310
::lumina::api::LuminaBuilder::Create(builder_options_, memory_resource));
280311
// pretrain
281-
LuminaDataset dataset1(count_, dimension_, array_vec_);
312+
LuminaDataset dataset1(indexed_count_, dimension_, array_vec_, array_start_ids_);
282313
PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.PretrainFrom(dataset1));
283314

284315
// insert data
285-
LuminaDataset dataset2(count_, dimension_, array_vec_);
316+
LuminaDataset dataset2(indexed_count_, dimension_, array_vec_, array_start_ids_);
286317
std::vector<std::shared_ptr<arrow::FloatArray>>().swap(array_vec_);
287318
PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.InsertFrom(dataset2));
288319

‎src/paimon/global_index/lumina/lumina_global_index.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ class LuminaIndexWriter : public GlobalIndexWriter {
9090

9191
private:
9292
int64_t count_ = 0;
93+
int64_t indexed_count_ = 0;
9394
std::shared_ptr<LuminaMemoryPool> pool_;
9495
std::string field_name_;
9596
std::shared_ptr<arrow::DataType> arrow_type_;
@@ -99,6 +100,7 @@ class LuminaIndexWriter : public GlobalIndexWriter {
99100
::lumina::api::IOOptions io_options_;
100101
std::map<std::string, std::string> lumina_options_;
101102
std::vector<std::shared_ptr<arrow::FloatArray>> array_vec_;
103+
std::vector<int64_t> array_start_ids_;
102104
};
103105

104106
class LuminaIndexReader : public GlobalIndexReader {

‎src/paimon/global_index/lumina/lumina_global_index_test.cpp‎

Lines changed: 198 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ class LuminaGlobalIndexTest : public ::testing::Test {
158158
return struct_array;
159159
}
160160

161-
private:
161+
protected:
162162
std::shared_ptr<MemoryPool> pool_ = GetDefaultPool();
163163
std::shared_ptr<FileSystem> fs_ = std::make_shared<LocalFileSystem>();
164164
std::map<std::string, std::string> options_ = {{"lumina.index.dimension", "4"},
@@ -470,4 +470,201 @@ TEST_F(LuminaGlobalIndexTest, TestHighCardinalityAndMultiThreadSearch) {
470470
}
471471
}
472472

473+
TEST_F(LuminaGlobalIndexTest, TestWriteWithNullRows) {
474+
auto test_root_dir = paimon::test::UniqueTestDirectory::Create();
475+
ASSERT_TRUE(test_root_dir);
476+
std::string test_root = test_root_dir->Str();
477+
478+
// Array with null at row 1 (middle): rows 0,2,3 are valid, row 1 is null
479+
// This should split into two segments: [0,0] and [2,3]
480+
std::shared_ptr<arrow::Array> array_with_null =
481+
arrow::ipc::internal::json::ArrayFromJSON(data_type_,
482+
R"([
483+
[[0.0, 0.0, 0.0, 0.0]],
484+
[null],
485+
[[1.0, 0.0, 1.0, 0.0]],
486+
[[1.0, 1.0, 1.0, 1.0]]
487+
])")
488+
.ValueOrDie();
489+
490+
ASSERT_OK_AND_ASSIGN(
491+
auto meta, WriteGlobalIndex(test_root, data_type_, options_, array_with_null, Range(0, 3)));
492+
ASSERT_OK_AND_ASSIGN(auto reader,
493+
CreateGlobalIndexReader(test_root, data_type_, options_, meta));
494+
{
495+
// Search should return ids 0, 2, 3 (skipping null row 1)
496+
ASSERT_OK_AND_ASSIGN(
497+
auto scored_result,
498+
reader->VisitVectorSearch(std::make_shared<VectorSearch>(
499+
/*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr,
500+
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_)));
501+
// Only 3 vectors indexed (row 1 is null), so limit=4 returns 3
502+
CheckResult(scored_result, {3l, 2l, 0l}, {0.01f, 2.21f, 4.21f});
503+
}
504+
}
505+
506+
TEST_F(LuminaGlobalIndexTest, TestWriteWithMultipleNullSegments) {
507+
auto test_root_dir = paimon::test::UniqueTestDirectory::Create();
508+
ASSERT_TRUE(test_root_dir);
509+
std::string test_root = test_root_dir->Str();
510+
511+
// Nulls at rows 0, 2, 5: valid rows are 1, 3, 4
512+
// Splits into segments: [1,1], [3,4]
513+
std::shared_ptr<arrow::Array> array_with_nulls =
514+
arrow::ipc::internal::json::ArrayFromJSON(data_type_,
515+
R"([
516+
[null],
517+
[[0.0, 1.0, 0.0, 1.0]],
518+
[null],
519+
[[1.0, 0.0, 1.0, 0.0]],
520+
[[1.0, 1.0, 1.0, 1.0]],
521+
[null]
522+
])")
523+
.ValueOrDie();
524+
525+
ASSERT_OK_AND_ASSIGN(auto meta, WriteGlobalIndex(test_root, data_type_, options_,
526+
array_with_nulls, Range(0, 5)));
527+
ASSERT_OK_AND_ASSIGN(auto reader,
528+
CreateGlobalIndexReader(test_root, data_type_, options_, meta));
529+
{
530+
ASSERT_OK_AND_ASSIGN(
531+
auto scored_result,
532+
reader->VisitVectorSearch(std::make_shared<VectorSearch>(
533+
/*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr,
534+
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_)));
535+
// Only 3 vectors indexed at ids 1, 3, 4
536+
CheckResult(scored_result, {4l, 1l, 3l}, {0.01f, 2.01f, 2.21f});
537+
}
538+
}
539+
540+
TEST_F(LuminaGlobalIndexTest, TestWriteWithAllNullRows) {
541+
auto test_root_dir = paimon::test::UniqueTestDirectory::Create();
542+
ASSERT_TRUE(test_root_dir);
543+
std::string test_root = test_root_dir->Str();
544+
545+
// All rows are null — no vectors to index
546+
std::shared_ptr<arrow::Array> all_null_array =
547+
arrow::ipc::internal::json::ArrayFromJSON(data_type_,
548+
R"([
549+
[null],
550+
[null],
551+
[null]
552+
])")
553+
.ValueOrDie();
554+
555+
auto global_index = std::make_shared<LuminaGlobalIndex>(options_);
556+
auto path_factory = std::make_shared<FakeIndexPathFactory>(test_root);
557+
auto file_writer = std::make_shared<GlobalIndexFileManager>(fs_, path_factory);
558+
559+
ASSERT_OK_AND_ASSIGN(
560+
std::shared_ptr<GlobalIndexWriter> global_writer,
561+
global_index->CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_writer, pool_));
562+
563+
ArrowArray c_array;
564+
ASSERT_TRUE(arrow::ExportArray(*all_null_array, &c_array).ok());
565+
std::vector<int64_t> row_ids = {0, 1, 2};
566+
ASSERT_OK(global_writer->AddBatch(&c_array, std::move(row_ids)));
567+
// Finish with zero indexed vectors — returns empty metas
568+
ASSERT_OK_AND_ASSIGN(auto result_metas, global_writer->Finish());
569+
ASSERT_TRUE(result_metas.empty());
570+
}
571+
572+
TEST_F(LuminaGlobalIndexTest, TestWriteWithNullAndFilter) {
573+
auto test_root_dir = paimon::test::UniqueTestDirectory::Create();
574+
ASSERT_TRUE(test_root_dir);
575+
std::string test_root = test_root_dir->Str();
576+
577+
// Null at row 2: valid rows are 0, 1, 3
578+
std::shared_ptr<arrow::Array> array_with_null =
579+
arrow::ipc::internal::json::ArrayFromJSON(data_type_,
580+
R"([
581+
[[0.0, 0.0, 0.0, 0.0]],
582+
[[0.0, 1.0, 0.0, 1.0]],
583+
[null],
584+
[[1.0, 1.0, 1.0, 1.0]]
585+
])")
586+
.ValueOrDie();
587+
588+
ASSERT_OK_AND_ASSIGN(
589+
auto meta, WriteGlobalIndex(test_root, data_type_, options_, array_with_null, Range(0, 3)));
590+
ASSERT_OK_AND_ASSIGN(auto reader,
591+
CreateGlobalIndexReader(test_root, data_type_, options_, meta));
592+
{
593+
// Filter: only allow ids < 3 (filters out id=3), so only ids 0, 1 remain
594+
auto filter = [](int64_t id) -> bool { return id < 3; };
595+
ASSERT_OK_AND_ASSIGN(
596+
auto scored_result,
597+
reader->VisitVectorSearch(std::make_shared<VectorSearch>(
598+
/*field_name=*/"f0", /*limit=*/4, query_, filter,
599+
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_)));
600+
CheckResult(scored_result, {1l, 0l}, {2.01f, 4.21f});
601+
}
602+
}
603+
604+
TEST_F(LuminaGlobalIndexTest, TestWriteWithNullAcrossMultipleBatches) {
605+
auto test_root_dir = paimon::test::UniqueTestDirectory::Create();
606+
ASSERT_TRUE(test_root_dir);
607+
std::string test_root = test_root_dir->Str();
608+
609+
// Batch 1: rows 0-2, null at row 1 → indexed ids: {0, 2}
610+
std::shared_ptr<arrow::Array> batch1 = arrow::ipc::internal::json::ArrayFromJSON(data_type_,
611+
R"([
612+
[[0.0, 0.0, 0.0, 0.0]],
613+
[null],
614+
[[1.0, 0.0, 1.0, 0.0]]
615+
])")
616+
.ValueOrDie();
617+
618+
// Batch 2: rows 3-5, null at row 3 → indexed ids: {4, 5}
619+
std::shared_ptr<arrow::Array> batch2 = arrow::ipc::internal::json::ArrayFromJSON(data_type_,
620+
R"([
621+
[null],
622+
[[1.0, 1.0, 1.0, 1.0]],
623+
[[0.0, 1.0, 0.0, 1.0]]
624+
])")
625+
.ValueOrDie();
626+
627+
auto global_index = std::make_shared<LuminaGlobalIndex>(options_);
628+
auto path_factory = std::make_shared<FakeIndexPathFactory>(test_root);
629+
auto file_writer = std::make_shared<GlobalIndexFileManager>(fs_, path_factory);
630+
631+
ASSERT_OK_AND_ASSIGN(
632+
std::shared_ptr<GlobalIndexWriter> global_writer,
633+
global_index->CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_writer, pool_));
634+
635+
// AddBatch 1: row_ids {0, 1, 2}
636+
{
637+
ArrowArray c_array;
638+
ASSERT_TRUE(arrow::ExportArray(*batch1, &c_array).ok());
639+
std::vector<int64_t> row_ids = {0, 1, 2};
640+
ASSERT_OK(global_writer->AddBatch(&c_array, std::move(row_ids)));
641+
}
642+
// AddBatch 2: row_ids {3, 4, 5}
643+
{
644+
ArrowArray c_array;
645+
ASSERT_TRUE(arrow::ExportArray(*batch2, &c_array).ok());
646+
std::vector<int64_t> row_ids = {3, 4, 5};
647+
ASSERT_OK(global_writer->AddBatch(&c_array, std::move(row_ids)));
648+
}
649+
650+
ASSERT_OK_AND_ASSIGN(auto result_metas, global_writer->Finish());
651+
ASSERT_EQ(result_metas.size(), 1);
652+
653+
ASSERT_OK_AND_ASSIGN(auto reader,
654+
CreateGlobalIndexReader(test_root, data_type_, options_, result_metas[0]));
655+
{
656+
// Search all: should return ids {0, 2, 4, 5}, never {1, 3}
657+
ASSERT_OK_AND_ASSIGN(
658+
auto scored_result,
659+
reader->VisitVectorSearch(std::make_shared<VectorSearch>(
660+
/*field_name=*/"f0", /*limit=*/10, query_, /*filter=*/nullptr,
661+
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_)));
662+
// id 0: [0,0,0,0] → L2 dist to [1,1,1,1.1] = 4.21
663+
// id 2: [1,0,1,0] → L2 dist = 2.21
664+
// id 4: [1,1,1,1] → L2 dist = 0.01
665+
// id 5: [0,1,0,1] → L2 dist = 2.01
666+
CheckResult(scored_result, {4l, 5l, 2l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f});
667+
}
668+
}
669+
473670
} // namespace paimon::lumina::test

0 commit comments

Comments
 (0)