Google Git Google Git
Sign in
apache/paimon-cpp/refs/heads/release-0.4/./test/inte/global_index_test.cpp
blob: 3c24f83f8ef11fac5374ab2c3aaf676fff45f315 [file]
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <algorithm>
#include <atomic>
#include <limits>
#include "arrow/c/bridge.h"
#include "arrow/type.h"
#include "fmt/format.h"
#include "fmt/ranges.h"
#include "gtest/gtest.h"
#include "paimon/common/factories/io_hook.h"
#include "paimon/common/global_index/bitmap/bitmap_global_index_factory.h"
#include "paimon/common/global_index/bitmap/bitmap_global_index_options.h"
#include "paimon/common/global_index/sorted_index_file_meta.h"
#include "paimon/common/global_index/union_global_index_reader.h"
#include "paimon/common/table/special_fields.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/common/utils/scope_guard.h"
#include "paimon/core/global_index/global_index_scan_impl.h"
#include "paimon/core/global_index/indexed_split_impl.h"
#include "paimon/core/table/source/data_split_impl.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/core/utils/tag_manager.h"
#include "paimon/defs.h"
#include "paimon/executor.h"
#include "paimon/fs/file_system.h"
#include "paimon/fs/local/local_file_system.h"
#include "paimon/global_index/bitmap_global_index_result.h"
#include "paimon/global_index/bitmap_scored_global_index_result.h"
#include "paimon/global_index/global_index_reader.h"
#include "paimon/global_index/global_index_result.h"
#include "paimon/global_index/global_index_scan.h"
#include "paimon/global_index/global_index_write_task.h"
#include "paimon/global_index/global_indexer.h"
#include "paimon/global_index/global_indexer_factory.h"
#include "paimon/predicate/literal.h"
#include "paimon/predicate/predicate_builder.h"
#include "paimon/result.h"
#include "paimon/status.h"
#include "paimon/table/source/full_text_search_builder.h"
#include "paimon/table/source/full_text_search_split.h"
#include "paimon/testing/utils/binary_row_generator.h"
#include "paimon/testing/utils/counting_cache_test_utils.h"
#include "paimon/testing/utils/deletion_vector_test_helper.h"
#include "paimon/testing/utils/io_exception_helper.h"
#include "paimon/testing/utils/test_helper.h"
#include "paimon/testing/utils/testharness.h"
namespace paimon::test {
namespace {
class CountingGlobalIndexFileManager : public GlobalIndexFileManager {
public:
CountingGlobalIndexFileManager(const std::shared_ptr<FileSystem>& fs,
const std::shared_ptr<IndexPathFactory>& path_factory)
: GlobalIndexFileManager(fs, path_factory, /*checkpoint_path_factory=*/nullptr) {}
Result<std::unique_ptr<InputStream>> GetInputStream(
const std::string& file_path) const override {
open_count_.fetch_add(1, std::memory_order_relaxed);
return GlobalIndexFileManager::GetInputStream(file_path);
}
int32_t OpenCount() const {
return open_count_.load(std::memory_order_relaxed);
}
private:
mutable std::atomic<int32_t> open_count_ = 0;
};
class DeleteAfterListFileSystem final : public LocalFileSystem {
public:
explicit DeleteAfterListFileSystem(std::string deleted_path)
: deleted_path_(std::move(deleted_path)),
listed_directory_name_(PathUtil::GetName(PathUtil::GetParentDirPath(deleted_path_))) {}
Status ListDir(const std::string& directory,
std::vector<BasicFileStatus>* file_status_list) const override {
PAIMON_RETURN_NOT_OK(LocalFileSystem::ListDir(directory, file_status_list));
if (PathUtil::GetName(directory) == listed_directory_name_ && !deleted_.exchange(true)) {
return LocalFileSystem::Delete(deleted_path_);
}
return Status::OK();
}
private:
std::string deleted_path_;
std::string listed_directory_name_;
mutable std::atomic<bool> deleted_ = false;
};
#ifdef PAIMON_ENABLE_ORC
class NestedFieldsIndexWriter : public GlobalIndexWriter {
public:
explicit NestedFieldsIndexWriter(std::shared_ptr<arrow::DataType> type)
: type_(std::move(type)) {}
Status AddBatch(::ArrowArray* array, std::vector<int64_t>&&) override {
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> decoded,
arrow::ImportArray(array, type_));
return ToPaimonStatus(decoded->ValidateFull());
}
Result<std::vector<GlobalIndexIOMeta>> Finish() override {
return std::vector<GlobalIndexIOMeta>{};
}
private:
std::shared_ptr<arrow::DataType> type_;
};
class NestedFieldsIndexer : public GlobalIndexer {
public:
Result<std::optional<std::vector<std::string>>> GetExtraFieldNames() const override {
return std::optional<std::vector<std::string>>(
{"details", "attributes", "plain_details", "plain_attributes"});
}
Result<std::shared_ptr<GlobalIndexWriter>> CreateWriter(
const std::string& field_name, ::ArrowSchema* arrow_schema,
const std::shared_ptr<GlobalIndexFileWriter>& file_writer,
const std::shared_ptr<MemoryPool>& pool) const override {
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::DataType> type,
arrow::ImportType(arrow_schema));
return std::shared_ptr<GlobalIndexWriter>(
std::make_shared<NestedFieldsIndexWriter>(std::move(type)));
}
Result<std::shared_ptr<GlobalIndexReader>> CreateReader(
::ArrowSchema* arrow_schema, const std::shared_ptr<GlobalIndexFileReader>& file_reader,
const std::vector<GlobalIndexIOMeta>& files, int64_t total_row_count,
const std::optional<std::vector<Range>>& row_ranges,
const std::shared_ptr<MemoryPool>& pool) const override {
return Status::NotImplemented("test index has no reader");
}
};
class NestedFieldsIndexFactory : public GlobalIndexerFactory {
public:
const char* Identifier() const override {
return "nested-fields-test-global";
}
Result<std::unique_ptr<GlobalIndexer>> Create(
const std::map<std::string, std::string>& options) const override {
return std::unique_ptr<GlobalIndexer>(std::make_unique<NestedFieldsIndexer>());
}
};
REGISTER_PAIMON_FACTORY(NestedFieldsIndexFactory);
#endif
} // namespace
using ScoredRows = std::vector<std::pair<int64_t, float>>;
using PlannedShards = std::vector<std::pair<Range, std::vector<Range>>>;
// string: FileFormat, bool: UseSpecificFileSystem
using ParamType = std::tuple<std::string, bool>;
/// This is a sdk end-to-end test for global index.
class GlobalIndexTest : public ::testing::Test, public ::testing::WithParamInterface<ParamType> {
void SetUp() override {
file_format_ = std::get<0>(GetParam());
dir_ = UniqueTestDirectory::Create("local");
if (std::get<1>(GetParam())) {
fs_ = dir_->GetFileSystem();
}
int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs();
std::srand(seed);
}
void TearDown() override {
dir_.reset();
}
void CreateTable(const std::vector<std::string>& partition_keys,
const std::shared_ptr<arrow::Schema>& schema,
const std::map<std::string, std::string>& options) const {
::ArrowSchema c_schema;
ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok());
ASSERT_OK_AND_ASSIGN(auto catalog, Catalog::Create(dir_->Str(), {}, fs_));
ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false));
ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), &c_schema, partition_keys,
/*primary_keys=*/{}, options,
/*ignore_if_exists=*/false));
}
void CreateTable(const std::vector<std::string>& partition_keys) const {
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
return CreateTable(partition_keys, arrow::schema(fields_), options);
}
void CreateTable() const {
return CreateTable(/*partition_keys=*/{});
}
Result<std::vector<std::shared_ptr<CommitMessage>>> WriteArray(
const std::string& table_path, const std::map<std::string, std::string>& partition,
const std::vector<std::string>& write_cols,
const std::shared_ptr<arrow::Array>& write_array) const {
// write
WriteContextBuilder write_builder(table_path, "commit_user_1");
write_builder.WithWriteSchema(write_cols);
write_builder.WithFileSystem(fs_);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<WriteContext> write_context, write_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(auto file_store_write,
FileStoreWrite::Create(std::move(write_context)));
ArrowArray c_array;
EXPECT_TRUE(arrow::ExportArray(*write_array, &c_array).ok());
auto record_batch = std::make_unique<RecordBatch>(
partition, /*bucket=*/0,
/*row_kinds=*/std::vector<RecordBatch::RowKind>(), &c_array);
PAIMON_RETURN_NOT_OK(file_store_write->Write(std::move(record_batch)));
PAIMON_ASSIGN_OR_RAISE(auto commit_msgs,
file_store_write->PrepareCommit(
/*wait_compaction=*/false, /*commit_identifier=*/0));
PAIMON_RETURN_NOT_OK(file_store_write->Close());
return commit_msgs;
}
Result<std::vector<std::shared_ptr<CommitMessage>>> WriteArray(
const std::string& table_path, const std::vector<std::string>& write_cols,
const std::shared_ptr<arrow::Array>& write_array) const {
return WriteArray(table_path, /*partition=*/{}, write_cols, write_array);
}
Status Commit(const std::string& table_path,
const std::vector<std::shared_ptr<CommitMessage>>& commit_msgs) const {
// commit
CommitContextBuilder commit_builder(table_path, "commit_user_1");
commit_builder.WithFileSystem(fs_);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<CommitContext> commit_context,
commit_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileStoreCommit> file_store_commit,
FileStoreCommit::Create(std::move(commit_context)));
return file_store_commit->Commit(commit_msgs);
}
Result<std::shared_ptr<DataSplitImpl>> ScanData(
const std::string& table_path,
const std::vector<std::map<std::string, std::string>>& partition_filters) const {
ScanContextBuilder scan_context_builder(table_path);
scan_context_builder.SetPartitionFilter(partition_filters);
scan_context_builder.WithFileSystem(fs_);
PAIMON_ASSIGN_OR_RAISE(auto scan_context, scan_context_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(auto table_scan, TableScan::Create(std::move(scan_context)));
PAIMON_ASSIGN_OR_RAISE(auto result_plan, table_scan->CreatePlan());
EXPECT_EQ(result_plan->Splits().size(), 1);
return std::dynamic_pointer_cast<DataSplitImpl>(result_plan->Splits()[0]);
}
Status WriteIndex(const std::string& table_path,
const std::vector<std::map<std::string, std::string>>& partition_filters,
const std::string& index_field_name, const std::string& index_type,
const std::map<std::string, std::string>& options, const Range& range) {
PAIMON_ASSIGN_OR_RAISE(auto split, ScanData(table_path, partition_filters));
PAIMON_ASSIGN_OR_RAISE(
auto index_commit_msg,
GlobalIndexWriteTask::WriteIndex(
table_path, index_field_name, index_type,
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>({range})), options,
/*task_id=*/std::nullopt, pool_, fs_));
return Commit(table_path, {index_commit_msg});
}
static Result<ScoredRows> SearchFullText(
const std::shared_ptr<GlobalIndexReader>& reader, const std::string& query, int32_t limit,
const std::optional<RoaringBitmap64>& include_row_ids) {
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<ScoredGlobalIndexResult> result,
reader->VisitFullTextSearch(std::make_shared<FullTextSearch>(
"f0", query, limit, include_row_ids)));
if (!result) {
return Status::Invalid("full-text search returned no result");
}
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScoredGlobalIndexResult::ScoredIterator> iter,
result->CreateScoredIterator());
ScoredRows rows;
while (iter->HasNext()) {
rows.push_back(iter->NextWithScore());
}
return rows;
}
static void CheckFullTextTopK(const std::shared_ptr<GlobalIndexReader>& reader,
const std::string& query,
const std::optional<RoaringBitmap64>& include_row_ids,
const std::vector<int64_t>& expected_row_ids) {
ASSERT_OK_AND_ASSIGN(ScoredRows all_rows,
SearchFullText(reader, query, /*limit=*/100, include_row_ids));
std::vector<int64_t> row_ids;
for (const auto& [row_id, score] : all_rows) {
row_ids.push_back(row_id);
}
ASSERT_EQ(row_ids, expected_row_ids);
ScoredRows ranked_rows = all_rows;
std::sort(
ranked_rows.begin(), ranked_rows.end(),
[](const std::pair<int64_t, float>& left, const std::pair<int64_t, float>& right) {
return left.second != right.second ? left.second > right.second
: left.first < right.first;
});
for (size_t limit = 1; limit < ranked_rows.size(); ++limit) {
ScoredRows expected_rows(ranked_rows.begin(), ranked_rows.begin() + limit);
std::sort(expected_rows.begin(), expected_rows.end());
ASSERT_OK_AND_ASSIGN(
ScoredRows rows,
SearchFullText(reader, query, static_cast<int32_t>(limit), include_row_ids));
ASSERT_EQ(rows, expected_rows) << "limit=" << limit;
}
}
std::unique_ptr<FullTextSearchBuilder> NewFullTextSearchBuilder(
const std::string& table_path, const std::string& query, int32_t limit,
const std::map<std::string, std::string>& options = {}) const {
auto builder = std::make_unique<FullTextSearchBuilder>(table_path);
builder->SetQuery("f0", query).SetLimit(limit).SetOptions(options).WithFileSystem(fs_);
return builder;
}
static Result<ScoredRows> ToScoredRows(const std::shared_ptr<ScoredGlobalIndexResult>& result) {
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScoredGlobalIndexResult::ScoredIterator> iter,
result->CreateScoredIterator());
ScoredRows rows;
while (iter->HasNext()) {
rows.push_back(iter->NextWithScore());
}
return rows;
}
Result<ScoredRows> SearchTable(
const std::string& table_path, const std::string& query, int32_t limit,
const std::vector<std::map<std::string, std::string>>& partition_filters = {},
const std::map<std::string, std::string>& options = {}) const {
std::unique_ptr<FullTextSearchBuilder> builder =
NewFullTextSearchBuilder(table_path, query, limit, options);
builder->SetPartitionFilter(partition_filters);
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<ScoredGlobalIndexResult> result,
builder->ExecuteLocal());
return ToScoredRows(result);
}
Status TagSnapshot(const std::string& table_path, int64_t snapshot_id,
const std::string& tag_name) const {
std::shared_ptr<FileSystem> fs = dir_->GetFileSystem();
std::string snapshot_json;
PAIMON_RETURN_NOT_OK(fs->ReadFile(SnapshotManager(fs, table_path).SnapshotPath(snapshot_id),
&snapshot_json));
return fs->WriteFile(TagManager(fs, table_path).TagPath(tag_name), snapshot_json,
/*overwrite=*/false);
}
Result<PlannedShards> PlanShards(const std::string& table_path,
const std::string& query) const {
PAIMON_ASSIGN_OR_RAISE(
std::unique_ptr<FullTextScan> scan,
NewFullTextSearchBuilder(table_path, query, /*limit=*/1)->NewFullTextScan());
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FullTextScan::Plan> plan, scan->CreatePlan());
PlannedShards shards;
for (const auto& split : plan->Splits()) {
auto index_split = std::dynamic_pointer_cast<IndexFullTextSearchSplit>(split);
if (!index_split || index_split->ColumnName() != "f0" ||
index_split->SnapshotId() != plan->SnapshotId()) {
return Status::Invalid("unexpected full-text search split");
}
shards.emplace_back(Range(index_split->RowRangeStart(), index_split->RowRangeEnd()),
index_split->SearchRowRanges());
}
return shards;
}
void CheckTableSearchMatchesReader(const std::string& table_path, const std::string& query,
const std::shared_ptr<GlobalIndexReader>& reader,
int32_t max_limit) const {
for (int32_t limit = 1; limit <= max_limit; ++limit) {
ASSERT_OK_AND_ASSIGN(ScoredRows expected_rows,
SearchFullText(reader, query, limit, std::nullopt));
ASSERT_OK_AND_ASSIGN(ScoredRows rows, SearchTable(table_path, query, limit));
ASSERT_EQ(rows, expected_rows) << "limit=" << limit;
}
}
void CheckSearchOverlappingIndexRanges(const std::string& index_type,
const std::map<std::string, std::string>& index_options,
const std::string& query) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["This is an test document.", 0],
["This is an new document document document.", 1],
["Document document document document test.", 2],
["Document document document document document.", 3]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, schema->field_names(), src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(0, 2)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(1, 1)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(1, 3)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(
auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers.size(), 1u);
CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2, 3});
ASSERT_OK_AND_ASSIGN(ScoredRows top_rows,
SearchFullText(readers[0], query, /*limit=*/1, std::nullopt));
ASSERT_EQ(top_rows.size(), 1u);
ASSERT_EQ(top_rows[0].first, 3);
CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2, 3});
ASSERT_OK_AND_ASSIGN(PlannedShards shards, PlanShards(table_path, query));
ASSERT_EQ(shards,
PlannedShards({{Range(0, 2), {Range(0, 2)}}, {Range(1, 3), {Range(3, 3)}}}));
CheckTableSearchMatchesReader(table_path, query, readers[0], /*max_limit=*/5);
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(3, 3)}));
ASSERT_OK_AND_ASSIGN(readers, global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(readers.size(), 1u);
CheckFullTextTopK(readers[0], query, std::nullopt, {3});
}
void CheckSearchSkipsDeletedRows(const std::string& index_type,
const std::map<std::string, std::string>& index_options) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::DELETION_VECTORS_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["apple apple apple apple apple apple apple apple apple apple", 0],
["apple apple apple apple apple apple apple apple apple", 1],
["apple apple apple apple apple apple apple apple", 2],
["apple apple apple apple apple apple apple", 3],
["apple apple apple apple apple apple", 4]
])")
.ValueOrDie();
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["apple apple apple apple apple", 5],
["apple apple apple apple", 6],
["apple apple apple", 7],
["apple apple", 8],
["apple", 9]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, schema->field_names(), src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, schema->field_names(), src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(0, 4)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(5, 9)));
auto create_scan = [&]() {
return GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_);
};
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> scan_before_deletion, create_scan());
std::string query = R"({"match":{"query":"apple"}})";
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextScan> full_text_scan,
NewFullTextSearchBuilder(table_path, query, /*limit=*/100)->NewFullTextScan());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<FullTextScan::Plan> plan_before_deletion,
full_text_scan->CreatePlan());
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
DeletionVectorTestHelper::RetrieveAnchorFileNames(split->DataFiles()));
ASSERT_EQ(anchor_file_names.size(), 2u);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<CommitMessage> dv_commit_msg,
DeletionVectorTestHelper::CreateDeletionVectorCommitMessage(
dir_->GetFileSystem(), table_path, file_format_, commit_msgs1[0],
{{anchor_file_names[0], /*deleted_positions=*/{0, 1, 2, 3}},
{anchor_file_names[1], /*deleted_positions=*/{0}}},
pool_));
ASSERT_OK(Commit(table_path, {dv_commit_msg}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> scan, create_scan());
ASSERT_OK_AND_ASSIGN(
auto readers_before_deletion,
scan_before_deletion->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers_before_deletion.size(), 1u);
CheckFullTextTopK(readers_before_deletion[0], query, std::nullopt,
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
ASSERT_OK_AND_ASSIGN(
ScoredRows rows_before_deletion,
SearchFullText(readers_before_deletion[0], query, /*limit=*/100, std::nullopt));
RoaringBitmap64 deleted_rows = RoaringBitmap64::From({0, 1, 2, 3, 5});
ScoredRows expected_live_rows;
for (const auto& row : rows_before_deletion) {
if (!deleted_rows.Contains(row.first)) {
expected_live_rows.push_back(row);
}
}
ASSERT_OK_AND_ASSIGN(auto readers,
scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(ScoredRows live_rows,
SearchFullText(readers[0], query, /*limit=*/100, std::nullopt));
ASSERT_EQ(live_rows, expected_live_rows);
CheckFullTextTopK(readers[0], query, std::nullopt, {4, 6, 7, 8, 9});
CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({0, 5, 8}), {8});
CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({0, 1}), {});
CheckTableSearchMatchesReader(table_path, query, readers[0], /*max_limit=*/6);
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextRead> full_text_read,
NewFullTextSearchBuilder(table_path, query, /*limit=*/100)->NewFullTextRead());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> pinned_result,
full_text_read->Read(*plan_before_deletion));
ASSERT_OK_AND_ASSIGN(ScoredRows pinned_rows, ToScoredRows(pinned_result));
ASSERT_EQ(pinned_rows, rows_before_deletion);
std::shared_ptr<FileSystem> file_system = dir_->GetFileSystem();
int64_t snapshot_before_deletion = plan_before_deletion->SnapshotId().value();
std::string snapshot_path =
SnapshotManager(file_system, table_path).SnapshotPath(snapshot_before_deletion);
std::string tag_path = TagManager(file_system, table_path).TagPath("before-deletion");
ASSERT_OK(TagSnapshot(table_path, snapshot_before_deletion, "before-deletion"));
ASSERT_OK(file_system->Delete(snapshot_path));
ASSERT_OK_AND_ASSIGN(pinned_result, full_text_read->Read(*plan_before_deletion));
ASSERT_OK_AND_ASSIGN(pinned_rows, ToScoredRows(pinned_result));
ASSERT_EQ(pinned_rows, rows_before_deletion);
std::string unrelated_tag_path = TagManager(file_system, table_path).TagPath("a-unrelated");
ASSERT_OK(file_system->WriteFile(unrelated_tag_path, "", /*overwrite=*/false));
std::unique_ptr<FullTextSearchBuilder> racing_builder =
NewFullTextSearchBuilder(table_path, query, /*limit=*/100);
racing_builder->WithFileSystem(
std::make_shared<DeleteAfterListFileSystem>(unrelated_tag_path));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<FullTextRead> racing_read,
racing_builder->NewFullTextRead());
ASSERT_OK_AND_ASSIGN(pinned_result, racing_read->Read(*plan_before_deletion));
ASSERT_OK_AND_ASSIGN(pinned_rows, ToScoredRows(pinned_result));
ASSERT_EQ(pinned_rows, rows_before_deletion);
ASSERT_OK_AND_ASSIGN(bool unrelated_tag_exists, file_system->Exists(unrelated_tag_path));
ASSERT_FALSE(unrelated_tag_exists);
ASSERT_OK(file_system->Delete(tag_path));
ASSERT_NOK_WITH_MSG(full_text_read->Read(*plan_before_deletion),
"planned from no longer exists");
ASSERT_OK(file_system->WriteFile(snapshot_path, "not a snapshot", /*overwrite=*/false));
Result<std::shared_ptr<ScoredGlobalIndexResult>> unreadable_snapshot_result =
full_text_read->Read(*plan_before_deletion);
ASSERT_FALSE(unreadable_snapshot_result.ok());
ASSERT_EQ(unreadable_snapshot_result.status().ToString().find("no longer exists"),
std::string::npos);
ASSERT_OK(file_system->Delete(snapshot_path));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> index_result,
readers[0]->VisitFullTextSearch(
std::make_shared<FullTextSearch>("f0", query, /*limit=*/5)));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({SpecialFields::ValueKind().ArrowField(), fields[1]}),
R"([
[0, 4],
[0, 6],
[0, 7],
[0, 8],
[0, 9]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, {"f1"}, expected_array, /*predicate=*/nullptr, plan));
}
static std::shared_ptr<arrow::Array> AppleRows(const arrow::FieldVector& fields, int32_t from,
int32_t to) {
std::vector<std::string> rows;
for (int32_t row = from; row < to; ++row) {
std::vector<std::string> terms(static_cast<size_t>(row + 1), "apple");
rows.push_back(fmt::format(R"(["{}", {}])", fmt::join(terms, " "), row));
}
return arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_(fields), fmt::format("[{}]", fmt::join(rows, ",\n")))
.ValueOrDie();
}
void CheckTableFullTextSearch(const std::string& index_type,
const std::map<std::string, std::string>& index_options) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, schema->field_names(),
AppleRows(fields, 0, 5)));
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, schema->field_names(),
AppleRows(fields, 5, 10)));
ASSERT_OK(Commit(table_path, commit_msgs2));
std::string query = R"({"match":{"query":"apple"}})";
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(0, 4)));
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextScan> first_shard_scan,
NewFullTextSearchBuilder(table_path, query, /*limit=*/10)->NewFullTextScan());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<FullTextScan::Plan> first_shard_plan,
first_shard_scan->CreatePlan());
ASSERT_TRUE(first_shard_plan->SnapshotId());
ASSERT_EQ(first_shard_plan->Splits().size(), 1u);
ASSERT_OK_AND_ASSIGN(ScoredRows indexed_rows,
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::FULL_TEXT_INDEX_SEARCH_MODE, "fast"},
{Options::GLOBAL_INDEX_SEARCH_MODE, "full"}}));
ASSERT_EQ(indexed_rows.size(), 5u);
for (const auto& [row_id, score] : indexed_rows) {
ASSERT_LT(row_id, 5) << score;
}
for (const char* mode : {"full", "detail"}) {
ASSERT_NOK_WITH_MSG(
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::FULL_TEXT_INDEX_SEARCH_MODE, mode}}),
"full and detail are not supported yet");
}
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options,
Range(5, 9)));
ASSERT_OK_AND_ASSIGN(PlannedShards shards, PlanShards(table_path, query));
ASSERT_EQ(shards,
PlannedShards({{Range(0, 4), {Range(0, 4)}}, {Range(5, 9), {Range(5, 9)}}}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(
auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers.size(), 1u);
CheckTableSearchMatchesReader(table_path, query, readers[0], /*max_limit=*/11);
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextScan> scan,
NewFullTextSearchBuilder(table_path, query, /*limit=*/3)->NewFullTextScan());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<FullTextScan::Plan> plan, scan->CreatePlan());
std::vector<std::shared_ptr<FullTextSearchSplit>> received_splits;
for (const auto& split : plan->Splits()) {
ASSERT_OK_AND_ASSIGN(std::string bytes, FullTextSearchSplit::Serialize(split, pool_));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<FullTextSearchSplit> received_split,
FullTextSearchSplit::Deserialize(bytes.data(), bytes.size(), pool_));
received_splits.push_back(std::move(received_split));
}
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextRead> read,
NewFullTextSearchBuilder(table_path, query, /*limit=*/3)->NewFullTextRead());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> top_result,
read->Read(received_splits));
ASSERT_OK_AND_ASSIGN(ScoredRows top_rows, ToScoredRows(top_result));
ASSERT_EQ(top_rows.size(), 3u);
ASSERT_OK_AND_ASSIGN(ScoredRows local_top_rows,
SearchTable(table_path, query, /*limit=*/3));
ASSERT_EQ(top_rows, local_top_rows);
std::string first_shard_snapshot_id =
std::to_string(first_shard_plan->SnapshotId().value());
ASSERT_OK_AND_ASSIGN(ScoredRows first_shard_rows,
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::SCAN_SNAPSHOT_ID, first_shard_snapshot_id}}));
ASSERT_EQ(first_shard_rows.size(), 5u);
for (const auto& [row_id, score] : first_shard_rows) {
ASSERT_LT(row_id, 5) << score;
}
ASSERT_OK_AND_ASSIGN(ScoredRows snapshot_mode_rows,
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::SCAN_MODE, "from-snapshot"},
{Options::SCAN_SNAPSHOT_ID, first_shard_snapshot_id}}));
ASSERT_EQ(snapshot_mode_rows, first_shard_rows);
ASSERT_NOK_WITH_MSG(SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::SCAN_MODE, "latest-full"},
{Options::SCAN_SNAPSHOT_ID, first_shard_snapshot_id}}),
"scan.snapshot-id must be null when you use latest-full for scan.mode");
ASSERT_OK(TagSnapshot(table_path, first_shard_plan->SnapshotId().value(), "first-shard"));
ASSERT_OK_AND_ASSIGN(ScoredRows tag_rows,
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::SCAN_TAG_NAME, "first-shard"}}));
ASSERT_EQ(tag_rows, first_shard_rows);
ASSERT_OK_AND_ASSIGN(ScoredRows latest_rows, SearchTable(table_path, query, /*limit=*/10));
ASSERT_EQ(latest_rows.size(), 10u);
ASSERT_OK_AND_ASSIGN(ScoredRows timestamp_rows,
SearchTable(table_path, query, /*limit=*/10, /*partition_filters=*/{},
{{Options::SCAN_TIMESTAMP_MILLIS,
std::to_string(std::numeric_limits<int64_t>::max())}}));
ASSERT_EQ(timestamp_rows, latest_rows);
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<FullTextRead> read_ten,
NewFullTextSearchBuilder(table_path, query, /*limit=*/10)->NewFullTextRead());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> first_shard_result,
read_ten->Read(*first_shard_plan));
ASSERT_OK_AND_ASSIGN(ScoredRows first_shard_plan_rows, ToScoredRows(first_shard_result));
ASSERT_EQ(first_shard_plan_rows, first_shard_rows);
std::vector<std::string> expected_rows;
for (const auto& [row_id, score] : top_rows) {
expected_rows.push_back(fmt::format("[0, {}, {}]", row_id, score));
}
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({SpecialFields::ValueKind().ArrowField(), fields[1],
SpecialFields::IndexScore().ArrowField()}),
fmt::format("[{}]", fmt::join(expected_rows, ",\n")))
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto data_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, top_result));
ASSERT_OK(ReadData(table_path, {"f1", "_INDEX_SCORE"}, expected_array,
/*predicate=*/nullptr, data_plan));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap", /*options=*/{},
Range(0, 9)));
ASSERT_OK_AND_ASSIGN(shards, PlanShards(table_path, query));
ASSERT_EQ(shards,
PlannedShards({{Range(0, 4), {Range(0, 4)}}, {Range(5, 9), {Range(5, 9)}}}));
ASSERT_OK_AND_ASSIGN(ScoredRows rows_with_bitmap,
SearchTable(table_path, query, /*limit=*/3));
ASSERT_EQ(rows_with_bitmap, top_rows);
}
void CheckTableSearchWithPartitionFilter(const std::string& table_path,
const std::string& query) const {
auto row_ids_of = [](const ScoredRows& rows) {
std::vector<int64_t> row_ids;
for (const auto& row : rows) {
row_ids.push_back(row.first);
}
return row_ids;
};
ASSERT_OK_AND_ASSIGN(ScoredRows rows,
SearchTable(table_path, query, /*limit=*/10, {{{"f1", "10"}}}));
ASSERT_EQ(row_ids_of(rows), std::vector<int64_t>({0, 1}));
ASSERT_OK_AND_ASSIGN(rows, SearchTable(table_path, query, /*limit=*/10, {{{"f1", "20"}}}));
ASSERT_EQ(row_ids_of(rows), std::vector<int64_t>({2}));
ASSERT_OK_AND_ASSIGN(
rows, SearchTable(table_path, query, /*limit=*/10, {{{"f1", "10"}}, {{"f1", "20"}}}));
ASSERT_EQ(row_ids_of(rows), std::vector<int64_t>({0, 1, 2}));
ASSERT_OK_AND_ASSIGN(rows, SearchTable(table_path, query, /*limit=*/10));
ASSERT_EQ(row_ids_of(rows), std::vector<int64_t>({0, 1, 2}));
}
Result<std::shared_ptr<Plan>> ScanGlobalIndexAndData(
const std::string& table_path, const std::shared_ptr<Predicate>& predicate,
const std::map<std::string, std::string>& options = {},
const std::shared_ptr<GlobalIndexResult>& index_result = nullptr,
const std::shared_ptr<Cache>& cache = nullptr) const {
ScanContextBuilder scan_context_builder(table_path);
scan_context_builder.SetPredicate(predicate)
.SetOptions(options)
.SetGlobalIndexResult(index_result)
.WithFileSystem(fs_);
if (cache) {
scan_context_builder.WithCache(cache);
}
PAIMON_ASSIGN_OR_RAISE(auto scan_context, scan_context_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(auto table_scan, TableScan::Create(std::move(scan_context)));
PAIMON_ASSIGN_OR_RAISE(auto result_plan, table_scan->CreatePlan());
return result_plan;
}
Result<std::shared_ptr<Plan>> ScanDataWithIndexResult(
const std::string& table_path, const std::vector<Range>& row_ranges,
const std::map<int64_t, float>& id_to_score) const {
std::shared_ptr<GlobalIndexResult> index_result;
if (id_to_score.empty()) {
index_result = BitmapGlobalIndexResult::FromRanges(row_ranges);
} else {
RoaringBitmap64 bitmap;
for (const auto& range : row_ranges) {
bitmap.AddRange(range.from, range.to + 1);
}
std::vector<float> scores;
for (auto iter = bitmap.Begin(); iter != bitmap.End(); ++iter) {
scores.push_back(id_to_score.at(*iter));
}
index_result = std::make_shared<BitmapScoredGlobalIndexResult>(std::move(bitmap),
std::move(scores));
}
return ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result);
}
Status ReadData(const std::string& table_path, const std::vector<std::string>& read_schema,
const std::shared_ptr<arrow::Array>& expected_array,
const std::shared_ptr<Predicate>& predicate,
const std::shared_ptr<Plan>& result_plan) const {
auto splits = result_plan->Splits();
ReadContextBuilder read_context_builder(table_path);
read_context_builder.SetReadFieldNames(read_schema)
.SetPredicate(predicate)
.WithFileSystem(fs_)
.EnableLateMaterializing(false);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context,
read_context_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context)));
PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits));
PAIMON_ASSIGN_OR_RAISE(auto read_result,
ReadResultCollector::CollectResult(std::move(batch_reader)));
if (!expected_array) {
if (read_result) {
return Status::Invalid("expected array is empty, but read result is not empty");
}
return Status::OK();
}
auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_array);
if (!expected_chunk_array->ApproxEquals(*read_result,
arrow::EqualOptions::Defaults().atol(1E-2))) {
std::cout << "result=" << read_result->ToString() << std::endl
<< "expected=" << expected_chunk_array->ToString() << std::endl;
return Status::Invalid("expected array and result array not equal");
}
return Status::OK();
}
private:
std::string file_format_;
std::unique_ptr<UniqueTestDirectory> dir_;
arrow::FieldVector fields_ = {
arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32()),
arrow::field("f2", arrow::int32()),
arrow::field("f3", arrow::float64()),
};
std::shared_ptr<MemoryPool> pool_ = GetDefaultPool();
std::shared_ptr<FileSystem> fs_ = nullptr;
};
#ifdef PAIMON_ENABLE_LUMINA
TEST_P(GlobalIndexTest, TestWriteLuminaIndex) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::list(arrow::float32()))};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["a", [0.0, 0.0, 0.0, 0.0]],
["b", [0.0, 1.0, 0.0, 1.0]],
["c", [1.0, 0.0, 1.0, 0.0]],
["d", [1.0, 1.0, 1.0, 1.0]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(auto index_commit_msg, GlobalIndexWriteTask::WriteIndex(
table_path, "f1", "lumina",
std::make_shared<IndexedSplitImpl>(
split, std::vector<Range>({Range(0, 3)})),
/*options=*/lumina_options,
/*task_id=*/std::nullopt, pool_));
auto index_commit_msg_impl = std::dynamic_pointer_cast<CommitMessageImpl>(index_commit_msg);
ASSERT_TRUE(index_commit_msg_impl);
// check commit message
std::string index_meta_json =
R"({"distance.metric":"l2","encoding.type":"rawf32","index.dimension":"4","index.type":"bruteforce","search.parallel_number":"10"})";
GlobalIndexMeta expected_global_index_meta(
/*row_range_start=*/0, /*row_range_end=*/3, /*index_field_id=*/1,
/*extra_field_ids=*/std::nullopt, std::make_shared<Bytes>(index_meta_json, pool_.get()));
auto expected_index_file_meta =
std::make_shared<IndexFileMeta>("lumina", /*file_name=*/"fake_index_file", /*file_size=*/10,
/*row_count=*/4, /*dv_ranges=*/std::nullopt,
/*external_path=*/std::nullopt, expected_global_index_meta);
DataIncrement expected_data_increment({expected_index_file_meta});
auto expected_commit_message = std::make_shared<CommitMessageImpl>(
/*partition=*/BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/std::nullopt,
expected_data_increment, CompactIncrement({}, {}, {}));
ASSERT_TRUE(expected_commit_message->TEST_Equal(*index_commit_msg_impl));
}
TEST_P(GlobalIndexTest, TestWriteLuminaIndexWithCheckpoint) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::list(arrow::float32()))};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> lumina_options = {
{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.extension.build.ckpt.count", "1"},
{"lumina.extension.build.ckpt.threshold", "1"},
{"lumina.search.parallel_number", "10"}};
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["a", [0.0, 0.0, 0.0, 0.0]],
["b", [0.0, 1.0, 0.0, 1.0]],
["c", [1.0, 0.0, 1.0, 0.0]],
["d", [1.0, 1.0, 1.0, 1.0]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
std::shared_ptr<FileSystem> checkpoint_fs = dir_->GetFileSystem();
std::string task_id = "task-1";
auto external_dir = UniqueTestDirectory::Create("local");
ASSERT_TRUE(external_dir);
for (bool use_external_path : {false, true}) {
std::map<std::string, std::string> checkpoint_options = lumina_options;
std::string checkpoint_dir = PathUtil::JoinPath(table_path, "index/checkpoint");
if (use_external_path) {
checkpoint_options["global-index.external-path"] = "FILE://" + external_dir->Str();
checkpoint_dir = PathUtil::JoinPath(external_dir->Str(), "checkpoint");
}
std::string task_checkpoint_dir = PathUtil::JoinPath(checkpoint_dir, task_id);
ASSERT_OK(checkpoint_fs->Mkdirs(task_checkpoint_dir));
ASSERT_OK(checkpoint_fs->Mkdirs(PathUtil::JoinPath(checkpoint_dir, "another-task")));
auto checkpoint_path = [&](const std::string& task, int64_t id) {
return PathUtil::JoinPath(PathUtil::JoinPath(checkpoint_dir, task),
fmt::format("lumina-global-index-f1-0-3-{}.index.ckpt", id));
};
std::string matching_checkpoint = checkpoint_path(task_id, 9);
ASSERT_OK(checkpoint_fs->WriteFile(matching_checkpoint, "invalid checkpoint",
/*overwrite=*/false));
std::string unrelated_checkpoint = checkpoint_path("another-task", 999);
ASSERT_OK(checkpoint_fs->WriteFile(unrelated_checkpoint, "another task",
/*overwrite=*/false));
ASSERT_OK(GlobalIndexWriteTask::WriteIndex(
table_path, "f1", "lumina",
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>{Range(0, 3)}),
checkpoint_options, task_id, pool_));
std::vector<BasicFileStatus> checkpoint_files;
ASSERT_OK(checkpoint_fs->ListDir(task_checkpoint_dir, &checkpoint_files));
ASSERT_TRUE(checkpoint_files.empty());
ASSERT_OK_AND_ASSIGN(bool unrelated_exists, checkpoint_fs->Exists(unrelated_checkpoint));
ASSERT_TRUE(unrelated_exists);
ASSERT_OK(
GlobalIndexWriteTask::DeleteCheckpoint(table_path, "another-task", checkpoint_options));
ASSERT_OK_AND_ASSIGN(unrelated_exists, checkpoint_fs->Exists(unrelated_checkpoint));
ASSERT_FALSE(unrelated_exists);
ASSERT_OK_AND_ASSIGN(
bool unrelated_directory_exists,
checkpoint_fs->Exists(PathUtil::JoinPath(checkpoint_dir, "another-task")));
ASSERT_FALSE(unrelated_directory_exists);
}
}
TEST_P(GlobalIndexTest, TestWriteLuminaIndexWithMismatchedDimension) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::list(arrow::float32()))};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "3"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::READ_BATCH_SIZE, "1"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["a", [0.0, 0.0, 0.0]],
["b", [0.0, 0.0, 0.0, 0.0]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_NOK_WITH_MSG(
WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina",
/*options=*/lumina_options, Range(0, 1)),
"invalid input array in LuminaIndexWriter, vector at row [0] has length [4], "
"expected dimension [3]");
}
TEST_P(GlobalIndexTest, TestWriteAndQueryLuminaIndexWithOrcDictionaryStringTags) {
if (file_format_ != "orc") {
GTEST_SKIP() << "ORC-only dictionary encoding case";
}
arrow::FieldVector fields = {arrow::field("embedding", arrow::list(arrow::float32())),
arrow::field("color", arrow::utf8()),
arrow::field("labels", arrow::list(arrow::utf8()))};
std::shared_ptr<arrow::Schema> schema = arrow::schema(fields);
std::map<std::string, std::string> lumina_options = {
{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"},
{"lumina.extension.build.tag.tag_schema",
R"([{"key_name":"color","type":"enum","value_type":"string"},)"
R"({"key_name":"labels","type":"enum","value_type":"string"}])"}};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, file_format_}, {Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"},
{"orc.dictionary-key-size-threshold", "1"}, {"orc.read.enable-lazy-decoding", "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::shared_ptr<arrow::Array> src_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[[0.0, 0.0, 0.0, 0.0], "red", ["hot", "common"]],
[[0.0, 1.0, 0.0, 1.0], "blue", ["cold", "common"]],
[[1.0, 0.0, 1.0, 0.0], "red", ["cold"]],
[[1.0, 1.0, 1.0, 1.0], "blue", ["hot"]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> commit_msgs,
WriteArray(table_path, schema->field_names(), src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "embedding", "lumina",
lumina_options, Range(0, 3)));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lumina_options, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<GlobalIndexReader>> lumina_readers,
global_index_scan->CreateReaders("embedding", std::nullopt));
ASSERT_EQ(lumina_readers.size(), 1u);
std::vector<float> query = {1.0f, 1.0f, 1.0f, 1.1f};
auto search_and_check = [&](const std::shared_ptr<Predicate>& predicate,
const std::string& expected) {
std::shared_ptr<VectorSearch> vector_search = std::make_shared<VectorSearch>(
"embedding", /*limit=*/4, query, /*filter=*/nullptr, predicate,
/*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> scored_result,
lumina_readers[0]->VisitVectorSearch(vector_search));
ASSERT_EQ(scored_result->ToString(), expected);
};
search_and_check(
PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "red", 3)),
"row ids: {0,2}, scores: {4.21,2.21}");
search_and_check(
PredicateBuilder::In(/*field_index=*/2, /*field_name=*/"labels", FieldType::STRING,
{Literal(FieldType::STRING, "hot", 3)}),
"row ids: {0,3}, scores: {4.21,0.01}");
}
TEST_P(GlobalIndexTest, TestWriteIndex) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Alice", 20, null, 18.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 10, 0, 14.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(
auto index_commit_msg,
GlobalIndexWriteTask::WriteIndex(
table_path, "f0", "bitmap",
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>({Range(0, 7)})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_));
auto index_commit_msg_impl = std::dynamic_pointer_cast<CommitMessageImpl>(index_commit_msg);
ASSERT_TRUE(index_commit_msg_impl);
// check commit message
std::shared_ptr<Bytes> index_meta =
SortedIndexFileMeta(std::make_shared<Bytes>("Alice", pool_.get()),
std::make_shared<Bytes>("Tony", pool_.get()), /*has_nulls=*/false)
.Serialize(pool_.get());
GlobalIndexMeta expected_global_index_meta(
/*row_range_start=*/0, /*row_range_end=*/7, /*index_field_id=*/0,
/*extra_field_ids=*/std::nullopt, index_meta);
auto expected_index_file_meta =
std::make_shared<IndexFileMeta>("bitmap", /*file_name=*/"fake_index_file", /*file_size=*/10,
/*row_count=*/8, /*dv_ranges=*/std::nullopt,
/*external_path=*/std::nullopt, expected_global_index_meta);
DataIncrement expected_data_increment({expected_index_file_meta});
auto expected_commit_message = std::make_shared<CommitMessageImpl>(
/*partition=*/BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/std::nullopt,
expected_data_increment, CompactIncrement({}, {}, {}));
ASSERT_TRUE(expected_commit_message->TEST_Equal(*index_commit_msg_impl));
{
// test invalid write task with none-registered index type
ASSERT_NOK_WITH_MSG(
GlobalIndexWriteTask::WriteIndex(
table_path, "f0", "invalid",
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>({Range(0, 7)})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_),
"Unknown index type invalid, may not registered");
}
{
// test invalid multiple ranges
ASSERT_NOK_WITH_MSG(GlobalIndexWriteTask::WriteIndex(
table_path, "f0", "bitmap",
std::make_shared<IndexedSplitImpl>(
split, std::vector<Range>({Range(0, 6), Range(7, 7)})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_),
"GlobalIndexWriteTask only supports a single contiguous range.");
}
}
#ifdef PAIMON_ENABLE_ORC
TEST_P(GlobalIndexTest, TestWriteIndexDecodesNestedOrcDictionaries) {
if (file_format_ != "orc") {
GTEST_SKIP() << "ORC-only dictionary encoding case";
}
arrow::FieldVector fields = {
arrow::field("key", arrow::int32()),
arrow::field("details", arrow::struct_({arrow::field("text", arrow::utf8()),
arrow::field("number", arrow::int32())})),
arrow::field("attributes", arrow::map(arrow::utf8(), arrow::utf8())),
arrow::field("plain_details", arrow::struct_({arrow::field("left", arrow::int32()),
arrow::field("right", arrow::int32())})),
arrow::field("plain_attributes", arrow::map(arrow::int32(), arrow::int32()))};
std::shared_ptr<arrow::Schema> schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, "orc"},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{"orc.dictionary-key-size-threshold", "1"},
{"orc.read.enable-lazy-decoding", "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::shared_ptr<arrow::Array> source =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[0, ["red", 10], [["color", "red"]], [1, 2], [[1, 2]]],
[1, ["blue", 20], [["color", "blue"]], [3, 4], [[3, 4]]],
[2, ["red", 30], [["color", "red"]], [5, 6], [[5, 6]]],
[3, ["blue", 40], [["color", "blue"]], [7, 8], [[7, 8]]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> messages,
WriteArray(table_path, schema->field_names(), source));
ASSERT_OK(Commit(table_path, messages));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> split,
ScanData(table_path, /*partition_filters=*/{}));
std::shared_ptr<IndexedSplitImpl> indexed_split =
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>{Range(0, 3)});
ReadContextBuilder read_builder(table_path);
read_builder.SetOptions(options).WithFileSystem(fs_).SetReadFieldNames(
{"key", "details", "attributes", "plain_details", "plain_attributes",
SpecialFields::RowId().Name()});
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, read_builder.Finish());
ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> table_read,
TableRead::Create(std::move(read_context)));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader,
table_read->CreateReader(indexed_split));
ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch());
ASSERT_FALSE(BatchReader::IsEofBatch(batch));
arrow::Result<std::shared_ptr<arrow::Array>> raw_result =
arrow::ImportArray(batch.first.get(), batch.second.get());
ASSERT_TRUE(raw_result.ok()) << raw_result.status().ToString();
std::shared_ptr<arrow::Array> raw = std::move(raw_result).ValueOrDie();
std::shared_ptr<arrow::StructArray> raw_rows =
std::dynamic_pointer_cast<arrow::StructArray>(raw);
ASSERT_TRUE(raw_rows);
std::shared_ptr<arrow::StructArray> raw_details =
std::dynamic_pointer_cast<arrow::StructArray>(raw_rows->GetFieldByName("details"));
std::shared_ptr<arrow::MapArray> raw_attributes =
std::dynamic_pointer_cast<arrow::MapArray>(raw_rows->GetFieldByName("attributes"));
ASSERT_TRUE(raw_details);
ASSERT_TRUE(raw_attributes);
std::shared_ptr<arrow::Array> raw_text = raw_details->GetFieldByName("text");
ASSERT_TRUE(raw_text);
ASSERT_EQ(raw_text->type_id(), arrow::Type::DICTIONARY);
ASSERT_EQ(raw_attributes->keys()->type_id(), arrow::Type::DICTIONARY);
ASSERT_EQ(raw_attributes->items()->type_id(), arrow::Type::DICTIONARY);
ASSERT_OK(
GlobalIndexWriteTask::WriteIndex(table_path, "key", "nested-fields-test", indexed_split,
/*options=*/{}, /*task_id=*/std::nullopt, pool_, fs_));
}
#endif
TEST_P(GlobalIndexTest, TestWriteIndexWithPartition) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Tony", 10, 0, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 18.1],
["Lucy", 20, 1, 15.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
auto build_index_and_check =
[&](const std::vector<std::map<std::string, std::string>>& partition,
const Range& expected_range, const BinaryRow& expected_partition_row) {
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, partition));
ASSERT_OK_AND_ASSIGN(
auto index_commit_msg,
GlobalIndexWriteTask::WriteIndex(
table_path, "f0", "bitmap",
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>({expected_range})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_));
auto index_commit_msg_impl =
std::dynamic_pointer_cast<CommitMessageImpl>(index_commit_msg);
ASSERT_TRUE(index_commit_msg_impl);
// check commit message
std::shared_ptr<Bytes> index_meta =
SortedIndexFileMeta(std::make_shared<Bytes>("Alice", pool_.get()),
std::make_shared<Bytes>("Tony", pool_.get()),
/*has_nulls=*/false)
.Serialize(pool_.get());
GlobalIndexMeta expected_global_index_meta(
/*row_range_start=*/expected_range.from, /*row_range_end=*/expected_range.to,
/*index_field_id=*/0,
/*extra_field_ids=*/std::nullopt, index_meta);
auto expected_index_file_meta = std::make_shared<IndexFileMeta>(
"bitmap", /*file_name=*/"fake_index_file", /*file_size=*/10,
/*row_count=*/expected_range.Count(), /*dv_ranges=*/std::nullopt,
/*external_path=*/std::nullopt, expected_global_index_meta);
DataIncrement expected_data_increment({expected_index_file_meta});
auto expected_commit_message = std::make_shared<CommitMessageImpl>(
/*partition=*/expected_partition_row,
/*bucket=*/0,
/*total_buckets=*/std::nullopt, expected_data_increment,
CompactIncrement({}, {}, {}));
ASSERT_TRUE(expected_commit_message->TEST_Equal(*index_commit_msg_impl));
};
// build index for f1=20 partition
build_index_and_check({{{"f1", "20"}}}, Range(5, 7),
BinaryRowGenerator::GenerateRow({20}, pool_.get()));
// build index for f1=10 partition
build_index_and_check({{{"f1", "10"}}}, Range(0, 4),
BinaryRowGenerator::GenerateRow({10}, pool_.get()));
}
#endif
TEST_P(GlobalIndexTest, TestScanIndex) {
if (file_format_ == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index.db/append_with_global_index";
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// test index reader
// test f0 field
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0,7}");
// test f0, f1, f2 fields
auto global_index_scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
{
// test with non predicate
ASSERT_OK_AND_ASSIGN(auto index_result,
global_index_scan_impl->Scan(/*predicate=*/nullptr));
ASSERT_FALSE(index_result);
}
{
// test equal predicate for f0
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test not equal predicate for f0
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{1,2,3,4,5,6}");
}
{
// test equal predicate for f1
auto predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{4,6,7}");
}
{
// test equal predicate for f2
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,1,4,5}");
}
{
// test is null predicate
auto predicate =
PredicateBuilder::IsNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT);
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{7}");
}
{
// test is not null predicate
auto predicate =
PredicateBuilder::IsNotNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT);
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,1,2,3,4,5,6}");
}
{
// test in predicate
auto predicate = PredicateBuilder::In(
/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
{Literal(FieldType::STRING, "Alice", 5), Literal(FieldType::STRING, "Bob", 3),
Literal(FieldType::STRING, "Lucy", 4)});
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,1,4,5,7}");
}
{
// test not in predicate
auto predicate = PredicateBuilder::NotIn(
/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
{Literal(FieldType::STRING, "Alice", 5), Literal(FieldType::STRING, "Bob", 3),
Literal(FieldType::STRING, "Lucy", 4)});
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{2,3,6}");
}
{
// test starts with predicate for f0
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Predicate> predicate,
PredicateBuilder::StartsWith(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Al", 2)));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test ends with predicate for f0
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Predicate> predicate,
PredicateBuilder::EndsWith(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "ce", 2)));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test contains predicate for f0
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Predicate> predicate,
PredicateBuilder::Contains(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "lic", 3)));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test like predicate for f0
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Predicate> predicate,
PredicateBuilder::Like(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "A_i%", 4)));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test and predicate
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_predicate, f1_predicate}));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{7}");
}
{
// test or predicate
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({f0_predicate, f1_predicate}));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,4,6,7}");
}
{
// test non-result
auto predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(30));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{}");
}
{
// test early stopping
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
auto f2_predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(6));
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto predicate,
PredicateBuilder::And({f1_predicate, f2_predicate, f0_predicate}));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{}");
}
{
// test greater than predicate
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{4,6,7}");
}
{
// test greater or equal predicate
auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,1,2,3,4,5,6,7}");
}
{
// test less than predicate
auto predicate = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{}");
}
{
// test less or equal predicate
auto predicate = PredicateBuilder::LessOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,1,2,3,5}");
}
{
// test a predicate for field with no index
auto f3_predicate = PredicateBuilder::Equal(/*field_index=*/3, /*field_name=*/"f3",
FieldType::DOUBLE, Literal(1.2));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(f3_predicate));
ASSERT_FALSE(index_result);
}
}
TEST_P(GlobalIndexTest, TestScanIndexWithSpecificSnapshot) {
if (file_format_ == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index.db/append_with_global_index";
// snapshot 2 has f0 index
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/2l,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// test index reader
// test f0 field
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0,7}");
// test f1 field
ASSERT_OK_AND_ASSIGN(auto index_readers2, global_index_scan->CreateReaders("f1", std::nullopt));
ASSERT_EQ(index_readers2.size(), 0u);
auto global_index_scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
{
// test and predicate
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_predicate, f1_predicate}));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(index_result->ToString(), "{0,7}");
}
{
// test or predicate
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({f0_predicate, f1_predicate}));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_FALSE(index_result);
}
}
TEST_P(GlobalIndexTest, TestScanIndexWithSpecificSnapshotWithNoIndex) {
if (file_format_ == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index.db/append_with_global_index";
// snapshot 1 has no index
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/1l,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// test index reader
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 0u);
auto global_index_scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_FALSE(index_result);
}
TEST_P(GlobalIndexTest, TestScanIndexWithRange) {
if (file_format_ == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index.db/append_with_global_index";
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
auto global_index_scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
{
// test index reader
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0,7}");
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto evaluator_result, global_index_scan_impl->Scan(predicate));
ASSERT_EQ(evaluator_result->ToString(), "{1,2,3,4,5,6}");
}
{
// invalid range
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(10, 13)}));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(index_readers.size(), 0u);
}
}
TEST_P(GlobalIndexTest, TestMultipleIndexFilesInOneRangePrunedByKeyRange) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto first_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["A", 10, 0, 1.0],
["B", 10, 1, 2.0]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto first_commit_messages,
WriteArray(table_path, {{"f1", "10"}}, write_cols, first_array));
ASSERT_OK(Commit(table_path, first_commit_messages));
auto second_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Y", 10, 0, 3.0],
["Z", 10, 1, 4.0]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto second_commit_messages,
WriteArray(table_path, {{"f1", "10"}}, write_cols, second_array));
ASSERT_OK(Commit(table_path, second_commit_messages));
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{{{"f1", "10"}}}));
ASSERT_EQ(split->DataFiles().size(), 2u);
std::shared_ptr<DataFileMeta> first_file;
std::shared_ptr<DataFileMeta> second_file;
for (const auto& file : split->DataFiles()) {
ASSERT_TRUE(file->first_row_id);
if (file->first_row_id.value() == 0) {
first_file = file;
} else if (file->first_row_id.value() == 2) {
second_file = file;
}
}
ASSERT_TRUE(first_file);
ASSERT_TRUE(second_file);
auto create_single_file_split = [&](const std::shared_ptr<DataFileMeta>& file) {
DataSplitImpl::Builder builder(split->Partition(), split->Bucket(), split->BucketPath(),
std::vector<std::shared_ptr<DataFileMeta>>{file});
return builder.WithTotalBuckets(split->TotalBuckets())
.WithSnapshot(split->SnapshotId())
.IsStreaming(split->IsStreaming())
.RawConvertible(split->RawConvertible())
.Build();
};
ASSERT_OK_AND_ASSIGN(auto first_split, create_single_file_split(first_file));
ASSERT_OK_AND_ASSIGN(auto second_split, create_single_file_split(second_file));
const Range index_range(0, 3);
for (const std::string index_type : {"bitmap", "btree"}) {
ASSERT_OK_AND_ASSIGN(
auto first_index_commit_message,
GlobalIndexWriteTask::WriteIndex(
table_path, "f0", index_type,
std::make_shared<IndexedSplitImpl>(first_split, std::vector<Range>({index_range})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_, fs_));
ASSERT_OK_AND_ASSIGN(
auto second_index_commit_message,
GlobalIndexWriteTask::WriteIndex(
table_path, "f0", index_type,
std::make_shared<IndexedSplitImpl>(second_split, std::vector<Range>({index_range})),
/*options=*/{}, /*task_id=*/std::nullopt, pool_, fs_));
ASSERT_OK(Commit(table_path, {first_index_commit_message, second_index_commit_message}));
}
auto check_pruning = [&](const std::string& index_type) {
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
auto scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scan_impl);
const auto& range_to_metas = scan_impl->index_metas_.at(0).at(index_type);
ASSERT_EQ(range_to_metas.size(), 1u);
ASSERT_EQ(range_to_metas.at(index_range).size(), 2u);
auto file_manager = scan_impl->index_file_manager_;
auto counting_file_manager = std::make_shared<CountingGlobalIndexFileManager>(
file_manager->fs_, file_manager->path_factory_);
scan_impl->index_file_manager_ = counting_file_manager;
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexReader> reader,
global_index_scan->CreateReader("f0", index_type, /*row_range_index=*/std::nullopt));
ASSERT_TRUE(reader);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> result,
reader->VisitEqual(Literal(FieldType::STRING, "Z", /*size=*/1)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{3}");
// "Z" is outside the first file's [A, B] key range, so first_key/last_key pruning should
// avoid opening it and read only the second [Y, Z] file.
ASSERT_EQ(counting_file_manager->OpenCount(), 1);
ASSERT_OK_AND_ASSIGN(result,
reader->VisitNotEqual(Literal(FieldType::STRING, "A", /*size=*/1)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2,3}");
// NOT EQUAL selects both physical files in the same logical row range and merges their
// range-local bitmaps.
ASSERT_EQ(counting_file_manager->OpenCount(), 2);
};
check_pruning("bitmap");
check_pruning("btree");
}
TEST_P(GlobalIndexTest, TestFallbackUnsupportedPropagatesAcrossRowRanges) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto first_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["A", 10, 0, 1.0],
["B", 10, 1, 2.0]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto first_commit_messages,
WriteArray(table_path, {{"f1", "10"}}, write_cols, first_array));
ASSERT_OK(Commit(table_path, first_commit_messages));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "bitmap",
/*options=*/{}, Range(0, 1)));
auto second_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Y", 20, 0, 3.0],
["Z", 20, 1, 4.0]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto second_commit_messages,
WriteArray(table_path, {{"f1", "20"}}, write_cols, second_array));
ASSERT_OK(Commit(table_path, second_commit_messages));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}}}, "f0", "bitmap",
/*options=*/{}, Range(2, 3)));
std::map<std::string, std::string> options = {
{BitmapGlobalIndexOptions::kBitmapIndexFallbackScanMaxSize, "1B"}};
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, options, fs_,
/*executor=*/nullptr, pool_));
Literal y(FieldType::STRING, "Y", /*size=*/1);
ASSERT_OK_AND_ASSIGN(RowRangeIndex first_range, RowRangeIndex::Create({Range(0, 1)}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> first_reader,
global_index_scan->CreateReader("f0", "bitmap", first_range));
ASSERT_TRUE(first_reader);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> first_result,
first_reader->VisitGreaterOrEqual(y));
ASSERT_TRUE(first_result);
ASSERT_EQ(first_result->ToString(), "{}");
ASSERT_OK_AND_ASSIGN(RowRangeIndex second_range, RowRangeIndex::Create({Range(2, 3)}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> second_reader,
global_index_scan->CreateReader("f0", "bitmap", second_range));
ASSERT_TRUE(second_reader);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> second_result,
second_reader->VisitGreaterOrEqual(y));
ASSERT_FALSE(second_result);
auto predicate = PredicateBuilder::GreaterOrEqual(
/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, y);
auto scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scan_impl);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> result, scan_impl->Scan(predicate));
// Range [A, B] is known empty while range [Y, Z] is unsupported because its selected file
// exceeds the fallback budget. The union must remain unsupported instead of becoming empty.
ASSERT_FALSE(result);
}
TEST_P(GlobalIndexTest, TestScanIndexWithPartition) {
if (file_format_ == "avro") {
return;
}
// only f1=10 has index
std::string table_path =
paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index_with_partition.db/append_with_global_index_with_partition";
auto check_result =
[&](const std::optional<std::vector<std::map<std::string, std::string>>>& partitions) {
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
partitions, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// test index reader
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index,
RowRangeIndex::Create({Range(0, 4)}));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result, index_readers[0]->VisitEqual(
Literal(FieldType::STRING, "Bob", 3)));
ASSERT_EQ(index_result->ToString(), "{1,4}");
auto global_index_scan_impl =
std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
{
// null result as f2 does not have index
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_FALSE(index_result);
}
{
// test not equal predicate for Bob
ASSERT_OK_AND_ASSIGN(auto index_result, index_readers[0]->VisitNotEqual(
Literal(FieldType::STRING, "Bob", 3)));
ASSERT_EQ(index_result->ToString(), "{0,2,3}");
}
{
// test equal predicate for Alice
ASSERT_OK_AND_ASSIGN(auto index_result, index_readers[0]->VisitEqual(Literal(
FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0}");
}
};
std::vector<std::map<std::string, std::string>> partitions = {{{"f1", "10"}}};
check_result(partitions);
check_result(std::nullopt);
}
TEST_P(GlobalIndexTest, TestScanIndexWithPartitionPredicate) {
if (file_format_ == "avro") {
return;
}
// Only f1=10 has an index in this table.
std::string table_path =
paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index_with_partition.db/append_with_global_index_with_partition";
auto check_result = [&](const std::shared_ptr<Predicate>& partitions,
const std::shared_ptr<MemoryPool>& pool, bool expect_index) {
ASSERT_OK_AND_ASSIGN(
std::unique_ptr<GlobalIndexScan> scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, partitions,
/*options=*/{}, fs_, /*executor=*/nullptr, pool));
ASSERT_OK_AND_ASSIGN(auto readers, scan->CreateReaders("f0", std::nullopt));
if (!expect_index) {
ASSERT_TRUE(readers.empty());
return;
}
ASSERT_EQ(readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto result,
readers[0]->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,4}");
};
check_result(/*partitions=*/nullptr, /*pool=*/nullptr, /*expect_index=*/true);
check_result(PredicateBuilder::Equal(/*field_index=*/0, "f1", FieldType::INT, Literal(10)),
pool_, /*expect_index=*/true);
check_result(PredicateBuilder::Equal(/*field_index=*/0, "f1", FieldType::INT, Literal(20)),
pool_, /*expect_index=*/false);
}
TEST_P(GlobalIndexTest, TestScanUnregisteredIndex) {
if (file_format_ == "avro") {
return;
}
auto factory_creator = FactoryCreator::GetInstance();
factory_creator->TEST_Unregister("bitmap-global");
ScopeGuard guard([&factory_creator]() {
factory_creator->Register("bitmap-global", (new BitmapGlobalIndexFactory));
});
std::string table_path = paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_global_index.db/append_with_global_index";
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 0u);
auto global_index_scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto index_result, global_index_scan_impl->Scan(predicate));
ASSERT_FALSE(index_result);
}
TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndex) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Alice", 20, null, 18.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 10, 0, 14.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap",
/*options=*/{}, Range(0, 7)));
ASSERT_OK_AND_ASSIGN(auto global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0,1}");
}
#ifdef PAIMON_ENABLE_LUMINA
TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndexWithPartition) {
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{"f2"}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto write_data_and_index = [&](const std::shared_ptr<arrow::Array>& src_array,
const std::map<std::string, std::string>& partition,
const Range& expected_range) {
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, partition, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write bitmap index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{partition}, "f0", "bitmap",
/*options=*/{}, expected_range));
// write and commit lumina index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{partition}, "f1", "lumina",
/*options=*/lumina_options, expected_range));
};
// write partition f2 = 10
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1]
])")
.ValueOrDie();
write_data_and_index(src_array1, {{"f2", "10"}}, Range(0, 3));
// write partition f2 = 20
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1],
["Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1],
["Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1],
["Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1],
["Tony", [11.0, 10.0, 11.0, 10.0], 20, 17.1]
])")
.ValueOrDie();
write_data_and_index(src_array2, {{"f2", "20"}}, Range(4, 8));
auto scan_and_check_result = [&](const std::map<std::string, std::string>& partition,
const std::optional<RowRangeIndex>& row_range_index,
VectorSearch::PreFilter filter, int32_t limit,
const std::string& bitmap_result,
const std::string& lumina_result,
const std::shared_ptr<arrow::Array>& expected_array,
const std::map<int64_t, float>& id_to_score) {
std::vector<std::map<std::string, std::string>> partitions = {partition};
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, partitions,
lumina_options, fs_, /*executor=*/nullptr, pool_));
// check bitmap index
ASSERT_OK_AND_ASSIGN(auto readers, global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto result1,
readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_OK_AND_ASSIGN(auto result2,
readers[0]->VisitEqual(Literal(FieldType::STRING, "Paul", 4)));
ASSERT_OK_AND_ASSIGN(auto index_result, result1->Or(result2));
ASSERT_EQ(index_result->ToString(), bitmap_result);
// check lumina index
ASSERT_OK_AND_ASSIGN(auto lumina_readers,
global_index_scan->CreateReaders("f1", row_range_index));
ASSERT_EQ(lumina_readers.size(), 1u);
auto lumina_reader = lumina_readers[0];
std::vector<float> query = {1.0f, 1.0f, 1.0f, 1.1f};
auto vector_search = std::make_shared<VectorSearch>(
"f1", limit, query, filter,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(auto scored_result, lumina_reader->VisitVectorSearch(vector_search));
ASSERT_EQ(scored_result->ToString(), lumina_result);
// check read array
std::vector<std::string> read_field_names = schema->field_names();
read_field_names.push_back("_INDEX_SCORE");
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, scored_result));
ASSERT_OK(ReadData(table_path, read_field_names, expected_array,
/*predicate=*/nullptr, plan));
};
auto result_fields = fields;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
result_fields.push_back(SpecialFields::IndexScore().ArrowField());
std::map<int64_t, float> id_to_score1 = {{0, 4.21f}, {1, 2.01f}, {2, 2.21f}, {3, 0.01f}};
std::map<int64_t, float> id_to_score2 = {
{0, 398.01f}, {1, 360.01f}, {2, 322.21f}, {3, 322.21f}, {4, 360.21f}};
{
// test scan and read for f2=10
auto filter = [](int64_t id) -> bool { return id == 0; };
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1, 4.21]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(0, 3)}));
scan_and_check_result({{"f2", "10"}}, row_range_index, filter, /*limit=*/2, "{0}",
"row ids: {0}, scores: {4.21}", expected_array, id_to_score1);
}
{
// test scan and read for f2=20
auto filter = [](int64_t id) -> bool { return id == 7 || id == 8; };
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1, 322.21]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(4, 8)}));
scan_and_check_result({{"f2", "20"}}, row_range_index, filter, /*limit=*/1, "{4,7}",
"row ids: {7}, scores: {322.21}", expected_array, id_to_score2);
}
{
// test invalid partition input
ASSERT_NOK_WITH_MSG(
GlobalIndexScan::Create(
table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::vector<std::map<std::string, std::string>>(), lumina_options,
fs_, /*executor=*/nullptr, pool_),
"invalid input partition, supposed to be null or at least one partition");
}
}
TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndexWithScore) {
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1],
["NullGuy1", null, 10, 20.0],
["Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1],
["Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1],
["NullGuy2", null, 20, 21.0],
["Tony", [11.0, 10.0, 11.0, 10.0], 20, 17.1],
["Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1],
["Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write and commit lumina index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina",
/*options=*/lumina_options, Range(0, 10)));
auto scan_and_check_result = [&](const std::vector<Range>& read_row_ranges,
const std::shared_ptr<arrow::Array>& expected_array,
const std::map<int64_t, float>& id_to_score) {
// check read array
std::vector<std::string> read_field_names = schema->field_names();
read_field_names.push_back("_INDEX_SCORE");
ASSERT_OK_AND_ASSIGN(auto plan,
ScanDataWithIndexResult(table_path, read_row_ranges, id_to_score));
ASSERT_OK(ReadData(table_path, read_field_names, expected_array,
/*predicate=*/nullptr, plan));
};
auto result_fields = fields;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
result_fields.push_back(SpecialFields::IndexScore().ArrowField());
std::map<int64_t, float> id_to_score = {{0, 4.21f}, {1, 2.01f}, {2, 2.21f},
{3, 0.01f}, {5, 322.21f}, {6, 360.01f},
{8, 360.21f}, {9, 398.01f}, {10, 322.21f}};
{
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1, 4.21],
[0, "Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1, 2.01],
[0, "Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1, 2.21],
[0, "Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1, 0.01],
[0, "Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1, 322.21],
[0, "Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1, 360.01],
[0, "Tony", [11.0, 10.0, 11.0, 10.0], 20, 17.1, 360.21],
[0, "Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1, 398.01],
[0, "Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1, 322.21]
])")
.ValueOrDie();
scan_and_check_result({Range(0, 3), Range(5, 6), Range(8, 10)}, expected_array,
id_to_score);
}
{
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1, 2.21],
[0, "Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1, 0.01],
[0, "Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1, 398.01],
[0, "Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1, 322.21]
])")
.ValueOrDie();
scan_and_check_result({Range(2, 3), Range(9, 10)}, expected_array, id_to_score);
}
{
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1, 360.01]
])")
.ValueOrDie();
scan_and_check_result({Range(6, 6)}, expected_array, id_to_score);
}
{
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1, null],
[0, "Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1, null],
[0, "Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1, null],
[0, "Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1, null]
])")
.ValueOrDie();
scan_and_check_result({Range(2, 3), Range(9, 10)}, expected_array, /*id_to_score=*/{});
}
{
// Verify null rows (id 4 and 7) are never recalled by vector search
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lumina_options, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto lumina_readers,
global_index_scan->CreateReaders("f1", std::nullopt));
ASSERT_EQ(lumina_readers.size(), 1u);
std::vector<float> query = {1.0f, 1.0f, 1.0f, 1.1f};
auto vector_search = std::make_shared<VectorSearch>(
"f1", /*limit=*/20, query, /*filter=*/nullptr,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(auto scored_result,
lumina_readers[0]->VisitVectorSearch(vector_search));
auto typed_result = std::dynamic_pointer_cast<BitmapScoredGlobalIndexResult>(scored_result);
ASSERT_TRUE(typed_result);
// Should recall 9 vectors (11 total rows - 2 null rows)
ASSERT_EQ(typed_result->bitmap_.Cardinality(), 9u);
// Null row ids 4 and 7 must not be in the result
ASSERT_FALSE(typed_result->bitmap_.Contains(4));
ASSERT_FALSE(typed_result->bitmap_.Contains(7));
}
}
TEST_P(GlobalIndexTest, TestWriteAndQueryLuminaIndexWithTagNullAndEmptyValues) {
arrow::FieldVector fields = {arrow::field("name", arrow::utf8()),
arrow::field("embedding", arrow::list(arrow::float32())),
arrow::field("color", arrow::utf8()),
arrow::field("labels", arrow::list(arrow::utf8())),
arrow::field("price", arrow::float32()),
arrow::field("scores", arrow::list(arrow::float32())),
arrow::field("category", arrow::int32()),
arrow::field("category_ids", arrow::list(arrow::int32()))};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> lumina_options = {
{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"},
{"lumina.extension.build.tag.tag_schema",
R"([{"key_name":"color","type":"enum","value_type":"string"},)"
R"({"key_name":"labels","type":"enum","value_type":"string"},)"
R"({"key_name":"price","type":"range","value_type":"float"},)"
R"({"key_name":"scores","type":"range","value_type":"float"},)"
R"({"key_name":"category","type":"enum","value_type":"int32"},)"
R"({"key_name":"category_ids","type":"enum","value_type":"int32"}])"}};
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["row0", [0.0, 0.0, 0.0, 0.0], "red", ["hot"], 5.0, [0.25], 7, [1, 2]],
["row1", null, "red", ["vip"], 8.0, [0.8], 7, [9]],
["row2", [0.0, 1.0, 0.0, 1.0], null, null, null, null, null, null],
["row3", [1.0, 0.0, 1.0, 0.0], " ", [], 20.0, [], 0, []],
["row4", [1.0, 1.0, 1.0, 1.0], "blue", ["vip", null], 30.0, [null, 0.75], 3, [null, 9]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> split,
ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<CommitMessage> index_commit_msg,
GlobalIndexWriteTask::WriteIndex(
table_path, "embedding", "lumina",
std::make_shared<IndexedSplitImpl>(split, std::vector<Range>({Range(0, 4)})),
/*options=*/lumina_options, /*task_id=*/std::nullopt, pool_, fs_));
std::shared_ptr<CommitMessageImpl> index_commit_msg_impl =
std::dynamic_pointer_cast<CommitMessageImpl>(index_commit_msg);
ASSERT_TRUE(index_commit_msg_impl);
const auto& new_index_files = index_commit_msg_impl->GetNewFilesIncrement().NewIndexFiles();
ASSERT_EQ(new_index_files.size(), 1u);
ASSERT_EQ(new_index_files[0]->IndexType(), "lumina");
ASSERT_EQ(new_index_files[0]->RowCount(), 5);
const std::optional<GlobalIndexMeta>& global_index_meta =
new_index_files[0]->GetGlobalIndexMeta();
ASSERT_TRUE(global_index_meta);
std::string expected_index_meta_json =
R"({"distance.metric":"l2","encoding.type":"rawf32","extension.build.tag.tag_schema":"[{\"key_name\":\"color\",\"type\":\"enum\",\"value_type\":\"string\"},{\"key_name\":\"labels\",\"type\":\"enum\",\"value_type\":\"string\"},{\"key_name\":\"price\",\"type\":\"range\",\"value_type\":\"float\"},{\"key_name\":\"scores\",\"type\":\"range\",\"value_type\":\"float\"},{\"key_name\":\"category\",\"type\":\"enum\",\"value_type\":\"int32\"},{\"key_name\":\"category_ids\",\"type\":\"enum\",\"value_type\":\"int32\"}]","index.dimension":"4","index.type":"bruteforce","search.parallel_number":"10"})";
GlobalIndexMeta expected_global_index_meta(
/*row_range_start=*/0, /*row_range_end=*/4, /*index_field_id=*/1,
/*extra_field_ids=*/std::optional<std::vector<int32_t>>({2, 3, 4, 5, 6, 7}),
std::make_shared<Bytes>(expected_index_meta_json, pool_.get()));
ASSERT_EQ(global_index_meta.value(), expected_global_index_meta);
ASSERT_OK(Commit(table_path, {index_commit_msg}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lumina_options, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto lumina_readers,
global_index_scan->CreateReaders("embedding", std::nullopt));
ASSERT_EQ(lumina_readers.size(), 1u);
std::vector<float> query = {1.0f, 1.0f, 1.0f, 1.1f};
auto search_and_check_with_filter = [&](VectorSearch::PreFilter pre_filter,
const std::shared_ptr<Predicate>& predicate,
const std::string& expected) {
std::shared_ptr<VectorSearch> vector_search = std::make_shared<VectorSearch>(
"embedding", /*limit=*/5, query, pre_filter, predicate,
/*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> scored_result,
lumina_readers[0]->VisitVectorSearch(vector_search));
ASSERT_EQ(scored_result->ToString(), expected);
};
auto search_and_check = [&](const std::shared_ptr<Predicate>& predicate,
const std::string& expected) {
search_and_check_with_filter(/*pre_filter=*/nullptr, predicate, expected);
};
search_and_check(/*predicate=*/nullptr, "row ids: {0,2,3,4}, scores: {4.21,2.01,2.21,0.01}");
search_and_check(
PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "red", 3)),
"row ids: {0}, scores: {4.21}");
search_and_check(PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color",
FieldType::STRING, Literal(FieldType::STRING, " ", 1)),
"row ids: {3}, scores: {2.21}");
search_and_check(
PredicateBuilder::In(/*field_index=*/3, /*field_name=*/"labels", FieldType::STRING,
{Literal(FieldType::STRING, "vip", 3)}),
"row ids: {4}, scores: {0.01}");
search_and_check(PredicateBuilder::LessOrEqual(/*field_index=*/4, /*field_name=*/"price",
FieldType::FLOAT, Literal(10.0f)),
"row ids: {0}, scores: {4.21}");
search_and_check(PredicateBuilder::GreaterOrEqual(/*field_index=*/5, /*field_name=*/"scores",
FieldType::FLOAT, Literal(0.5f)),
"row ids: {4}, scores: {0.01}");
search_and_check(PredicateBuilder::Equal(/*field_index=*/6, /*field_name=*/"category",
FieldType::INT, Literal(7)),
"row ids: {0}, scores: {4.21}");
search_and_check(PredicateBuilder::In(/*field_index=*/7, /*field_name=*/"category_ids",
FieldType::INT, {Literal(9)}),
"row ids: {4}, scores: {0.01}");
search_and_check(
PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "green", 5)),
"row ids: {}, scores: {}");
{
std::shared_ptr<Predicate> predicate = PredicateBuilder::Equal(
/*field_index=*/2, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "red", 3));
search_and_check_with_filter([](int64_t id) { return id == 0 || id == 4; }, predicate,
"row ids: {0}, scores: {4.21}");
search_and_check_with_filter([](int64_t id) { return id == 4; }, predicate,
"row ids: {}, scores: {}");
}
{
std::shared_ptr<Predicate> red_predicate = PredicateBuilder::Equal(
/*field_index=*/2, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "red", 3));
std::shared_ptr<Predicate> blue_predicate = PredicateBuilder::Equal(
/*field_index=*/2, /*field_name=*/"color", FieldType::STRING,
Literal(FieldType::STRING, "blue", 4));
std::shared_ptr<Predicate> high_price_predicate = PredicateBuilder::GreaterOrEqual(
/*field_index=*/4, /*field_name=*/"price", FieldType::FLOAT, Literal(30.0f));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> blue_high_price_predicate,
PredicateBuilder::And({blue_predicate, high_price_predicate}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> compound_predicate,
PredicateBuilder::Or({red_predicate, blue_high_price_predicate}));
search_and_check(compound_predicate, "row ids: {0,4}, scores: {4.21,0.01}");
search_and_check_with_filter([](int64_t id) { return id == 4; }, compound_predicate,
"row ids: {4}, scores: {0.01}");
}
}
#endif
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScan) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write and commit data
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Alice", 20, null, 18.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 10, 0, 14.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_all_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 20, null, 18.1],
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 10, 1, 16.1],
[0, "Emily", 10, 0, 13.1],
[0, "Lucy", 20, 1, 15.1],
[0, "Tony", 10, 0, 14.1],
[0, "Tony", 20, 0, 17.1]
])")
.ValueOrDie();
{
// read when no index is built
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, expected_all_array, predicate, plan));
}
// write and commit global index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap", /*options=*/{},
Range(0, 7)));
{
auto cache = std::make_shared<CountingRoutingCache>(CacheKind::MANIFEST, 64 * 1024 * 1024);
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK(ScanGlobalIndexAndData(table_path, predicate, /*options=*/{},
/*index_result=*/nullptr, cache));
ASSERT_GE(cache->GetCount(CacheKind::MANIFEST), 2);
int64_t first_supplier_calls = cache->SupplierCallCount(CacheKind::MANIFEST);
ASSERT_OK(ScanGlobalIndexAndData(table_path, predicate, /*options=*/{},
/*index_result=*/nullptr, cache));
ASSERT_EQ(cache->SupplierCallCount(CacheKind::MANIFEST), first_supplier_calls);
}
// scan and read with global index
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 20, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice2", 6));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, /*expected_array=*/nullptr, predicate, plan));
}
{
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 20, null, 18.1],
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 10, 1, 16.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr));
ASSERT_OK(
ReadData(table_path, write_cols, expected_all_array, /*predicate=*/nullptr, plan));
}
{
// f1 does not have global index
auto predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(19));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, expected_all_array, predicate, plan));
}
{
// disable global index
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(
auto plan,
ScanGlobalIndexAndData(table_path, predicate, {{"global-index.enabled", "false"}}));
ASSERT_OK(ReadData(table_path, write_cols, expected_all_array, predicate, plan));
}
}
TEST_P(GlobalIndexTest, TestDataEvolutionGlobalIndexSnapshotSelection) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 20, 0, 12.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap", /*options=*/{},
Range(0, 1)));
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "missing", 7));
ASSERT_OK_AND_ASSIGN(auto latest_plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_TRUE(latest_plan->Splits().empty());
ASSERT_EQ(latest_plan->SnapshotId(), std::optional<int64_t>(2));
const std::map<std::string, std::string> explicit_latest_options = {
{Options::SCAN_MODE, "latest"},
{Options::SCAN_SNAPSHOT_ID, "999"},
{Options::SCAN_TAG_NAME, "ignored"},
{Options::SCAN_TIMESTAMP_MILLIS, "0"}};
ASSERT_OK_AND_ASSIGN(auto explicit_latest_plan,
ScanGlobalIndexAndData(table_path, predicate, explicit_latest_options));
ASSERT_TRUE(explicit_latest_plan->Splits().empty());
ASSERT_EQ(explicit_latest_plan->SnapshotId(), std::optional<int64_t>(2));
auto empty_index_result = BitmapGlobalIndexResult::FromRanges({});
ASSERT_OK_AND_ASSIGN(auto supplied_latest_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
explicit_latest_options, empty_index_result));
ASSERT_TRUE(supplied_latest_plan->Splits().empty());
ASSERT_FALSE(supplied_latest_plan->SnapshotId());
ASSERT_OK_AND_ASSIGN(auto supplied_explicit_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
{{Options::SCAN_SNAPSHOT_ID, "1"},
{Options::SCAN_TAG_NAME, "ignored"},
{Options::SCAN_TIMESTAMP_MILLIS, "0"}},
empty_index_result));
ASSERT_TRUE(supplied_explicit_plan->Splits().empty());
ASSERT_FALSE(supplied_explicit_plan->SnapshotId());
ASSERT_NOK_WITH_MSG(
ScanGlobalIndexAndData(table_path, predicate, {{Options::SCAN_MODE, "from-snapshot"}}),
"scan.snapshot-id or scan.tag-name must be set when startup mode is FROM_SNAPSHOT");
std::vector<std::map<std::string, std::string>> time_travel_options = {
{{Options::SCAN_TAG_NAME, "tag"}},
{{Options::SCAN_TIMESTAMP_MILLIS, std::to_string(std::numeric_limits<int64_t>::max())}}};
for (const auto& options : time_travel_options) {
Result<std::shared_ptr<Plan>> result =
ScanGlobalIndexAndData(table_path, predicate, options);
ASSERT_TRUE(result.status().IsNotImplemented()) << result.status().ToString();
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Plan> supplied_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr, options, empty_index_result));
ASSERT_TRUE(supplied_plan->Splits().empty());
ASSERT_FALSE(supplied_plan->SnapshotId());
}
auto unindexed_predicate = PredicateBuilder::Equal(/*field_index=*/3, /*field_name=*/"f3",
FieldType::DOUBLE, Literal(99.9));
Result<std::shared_ptr<Plan>> unindexed_time_travel_result =
ScanGlobalIndexAndData(table_path, unindexed_predicate, time_travel_options.back());
ASSERT_TRUE(unindexed_time_travel_result.status().IsNotImplemented())
<< unindexed_time_travel_result.status().ToString();
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Plan> no_index_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr, time_travel_options.back()));
ASSERT_EQ(no_index_plan->SnapshotId(), std::optional<int64_t>(2));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Plan> supplied_nonexistent_snapshot_plan,
ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
{{Options::SCAN_SNAPSHOT_ID, "999"}}, empty_index_result));
ASSERT_TRUE(supplied_nonexistent_snapshot_plan->Splits().empty());
ASSERT_FALSE(supplied_nonexistent_snapshot_plan->SnapshotId());
ASSERT_NOK_WITH_MSG(
ScanGlobalIndexAndData(table_path, predicate, {{Options::SCAN_SNAPSHOT_ID, "999"}}),
"snapshot-999' not exists");
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithOnlyOnePartitionHasIndex) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write and commit data
std::vector<std::string> write_cols = schema->field_names();
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Tony", 10, 0, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 18.1],
["Lucy", 20, 1, 15.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
// write and commit global index for f1 = 10
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "bitmap",
/*options=*/{}, Range(0, 4)));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithTwoIndexInDiffTwoPartition) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write and commit data
std::vector<std::string> write_cols = schema->field_names();
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Tony", 10, 0, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 18.1],
["Tony", 20, 0, 17.1],
["Lucy", 20, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
// write and commit global index for f1 = 10
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "bitmap",
/*options=*/{}, Range(0, 4)));
// write and commit global index for f1 = 20
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}}}, "f2", "bitmap",
/*options=*/{}, Range(5, 7)));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Lucy", 20, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// only f1 = 10 partition has f0 index, query predicate1 results in ["Alice", 10, 1, 11.1]
// only f2 = 20 partition has f2 index, query predicate2 results in ["Lucy", 20, 1, 15.1]
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, /*expected_array=*/nullptr, predicate, plan));
}
{
// predicate2 is partition filter
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithTwoPartitionAllWithIndex) {
CreateTable(/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write and commit data
std::vector<std::string> write_cols = schema->field_names();
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1],
["Emily", 10, 0, 13.1],
["Tony", 10, 0, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 18.1],
["Lucy", 20, 1, 15.1],
["Tony", 20, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
// write and commit global index for f1 = 10
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "bitmap",
/*options=*/{}, Range(0, 4)));
// write and commit global index for f1 = 20
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}}}, "f0", "bitmap",
/*options=*/{}, Range(5, 7)));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 20, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// predicate2 is partition filter
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// predicate2 is partition filter
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// predicate2 is partition filter
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(30));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 20, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithPartitionWithTwoFields) {
CreateTable(/*partition_keys=*/{"f1", "f2"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write and commit data
std::vector<std::string> write_cols = schema->field_names();
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 1, 16.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, {{"f1", "10"}, {"f2", "1"}},
write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Lucy", 20, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, {{"f1", "20"}, {"f2", "1"}},
write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
auto src_array3 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Emily", 10, 0, 13.1],
["Tony", 10, 0, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs3, WriteArray(table_path, {{"f1", "10"}, {"f2", "0"}},
write_cols, src_array3));
ASSERT_OK(Commit(table_path, commit_msgs3));
// build index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}, {"f2", "1"}}}, "f0",
"bitmap",
/*options=*/{}, Range(0, 2)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}, {"f2", "1"}}}, "f0",
"bitmap",
/*options=*/{}, Range(3, 3)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}, {"f2", "0"}}}, "f0",
"bitmap",
/*options=*/{}, Range(4, 5)));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
{
auto predicate1 = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
auto predicate2 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 10, 1, 16.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
auto predicate1 = PredicateBuilder::GreaterThan(
/*field_index=*/1, /*field_name=*/"f1", FieldType::INT, Literal(10));
auto predicate2 = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 10, 1, 16.1],
[0, "Lucy", 20, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
auto predicate1 = PredicateBuilder::GreaterThan(
/*field_index=*/1, /*field_name=*/"f1", FieldType::INT, Literal(10));
auto predicate2 = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(auto or_predicate, PredicateBuilder::Or({predicate1, predicate2}));
auto predicate3 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Emily", 5));
ASSERT_OK_AND_ASSIGN(auto and_predicate, PredicateBuilder::And({predicate3, or_predicate}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, and_predicate));
ASSERT_OK(
ReadData(table_path, write_cols, /*expected_array=*/nullptr, and_predicate, plan));
}
}
#ifdef PAIMON_ENABLE_LUMINA
TEST_P(GlobalIndexTest, TestScanIndexWithTwoIndexes) {
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1],
["Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1],
["Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1],
["Tony", [11.0, 10.0, 11.0, 10.0], 20, 17.1]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write and commit bitmap global index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap",
/*options=*/{}, Range(0, 8)));
// write and commit lumina global index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina",
/*options=*/lumina_options, Range(0, 8)));
ASSERT_OK_AND_ASSIGN(
auto global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt,
/*options=*/lumina_options, fs_, /*executor=*/nullptr, pool_));
// query f0
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1);
ASSERT_OK_AND_ASSIGN(auto index_result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_EQ(index_result->ToString(), "{0,1}");
// query f1
ASSERT_OK_AND_ASSIGN(index_readers, global_index_scan->CreateReaders("f1", std::nullopt));
ASSERT_EQ(index_readers.size(), 1);
std::vector<float> query = {11.0f, 11.0f, 11.0f, 11.0f};
ASSERT_OK_AND_ASSIGN(
auto scored_result,
index_readers[0]->VisitVectorSearch(std::make_shared<VectorSearch>(
"f1", 1, query, /*filter=*/nullptr,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/lumina_options)));
ASSERT_EQ(scored_result->ToString(), "row ids: {1}, scores: {0.00}");
// query f2
ASSERT_OK_AND_ASSIGN(index_readers, global_index_scan->CreateReaders("f2", std::nullopt));
ASSERT_EQ(index_readers.size(), 0);
}
TEST_P(GlobalIndexTest, TestLuminaSearchSkipsRowsDeletedByDeletionVectors) {
arrow::FieldVector fields = {arrow::field("f0", arrow::int32()),
arrow::field("f1", arrow::list(arrow::float32()))};
std::map<std::string, std::string> lumina_options = {
{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"},
{"lumina.extension.build.tag.tag_schema",
R"([{"key_name":"f0","type":"enum","value_type":"int32"}])"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::DELETION_VECTORS_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[0, [0.0, 0.0, 0.0, 0.0]],
[1, [1.0, 1.0, 1.0, 1.0]],
[2, [2.0, 2.0, 2.0, 2.0]],
[3, [3.0, 3.0, 3.0, 3.0]],
[4, [4.0, 4.0, 4.0, 4.0]],
[5, [5.0, 5.0, 5.0, 5.0]]
])")
.ValueOrDie();
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[6, [6.0, 6.0, 6.0, 6.0]],
[7, [7.0, 7.0, 7.0, 7.0]],
[8, [8.0, 8.0, 8.0, 8.0]],
[9, [9.0, 9.0, 9.0, 9.0]]
])")
.ValueOrDie();
auto src_array3 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[10, [10.0, 10.0, 10.0, 10.0]],
[11, [11.0, 11.0, 11.0, 11.0]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, schema->field_names(), src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, schema->field_names(), src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
ASSERT_OK_AND_ASSIGN(auto commit_msgs3,
WriteArray(table_path, schema->field_names(), src_array3));
ASSERT_OK(Commit(table_path, commit_msgs3));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina", lumina_options,
Range(0, 5)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina", lumina_options,
Range(6, 9)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina", lumina_options,
Range(10, 11)));
auto create_scan = [&]() {
return GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lumina_options, fs_,
/*executor=*/nullptr, pool_);
};
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> scan_before_deletion, create_scan());
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
DeletionVectorTestHelper::RetrieveAnchorFileNames(split->DataFiles()));
ASSERT_EQ(anchor_file_names.size(), 3u);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<CommitMessage> dv_commit_msg,
DeletionVectorTestHelper::CreateDeletionVectorCommitMessage(
dir_->GetFileSystem(), table_path, file_format_, commit_msgs1[0],
{{anchor_file_names[0], /*deleted_positions=*/{0, 1, 2}},
{anchor_file_names[1], /*deleted_positions=*/{1}},
{anchor_file_names[2], /*deleted_positions=*/{0, 1}}},
pool_));
ASSERT_OK(Commit(table_path, {dv_commit_msg}));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> scan, create_scan());
std::vector<float> query = {0.0f, 0.0f, 0.0f, 0.0f};
auto search = [&](const std::shared_ptr<GlobalIndexScan>& index_scan,
const std::optional<RowRangeIndex>& row_range_index,
VectorSearch::PreFilter pre_filter,
const std::shared_ptr<Predicate>& predicate)
-> Result<std::shared_ptr<ScoredGlobalIndexResult>> {
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<GlobalIndexReader> reader,
index_scan->CreateReader("f1", "lumina", row_range_index));
if (!reader) {
return Status::Invalid("no lumina index on f1");
}
return reader->VisitVectorSearch(std::make_shared<VectorSearch>(
"f1", /*limit=*/3, query, std::move(pre_filter), predicate,
/*distance_type=*/std::nullopt, lumina_options));
};
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ScoredGlobalIndexResult> index_result,
search(scan_before_deletion, /*row_range_index=*/std::nullopt,
/*pre_filter=*/nullptr, /*predicate=*/nullptr));
ASSERT_EQ(index_result->ToString(),
"row ids: {0,1,2,6,7,8,10,11}, "
"scores: {0.00,4.00,16.00,144.00,196.00,256.00,400.00,484.00}");
ASSERT_OK_AND_ASSIGN(index_result, search(
scan, /*row_range_index=*/std::nullopt,
[](int64_t row_id) -> bool { return row_id != 8; },
/*predicate=*/nullptr));
ASSERT_EQ(index_result->ToString(),
"row ids: {3,4,5,6,9}, scores: {36.00,64.00,100.00,144.00,324.00}");
std::shared_ptr<Predicate> predicate = PredicateBuilder::In(
/*field_index=*/0, /*field_name=*/"f0", FieldType::INT,
{Literal(1), Literal(4), Literal(7), Literal(8), Literal(9)});
ASSERT_OK_AND_ASSIGN(index_result,
search(
scan, /*row_range_index=*/std::nullopt,
[](int64_t row_id) -> bool { return row_id != 9; }, predicate));
ASSERT_EQ(index_result->ToString(), "row ids: {4,8}, scores: {64.00,256.00}");
ASSERT_OK_AND_ASSIGN(index_result, search(scan, /*row_range_index=*/std::nullopt,
/*pre_filter=*/nullptr, /*predicate=*/nullptr));
ASSERT_EQ(index_result->ToString(),
"row ids: {3,4,5,6,8,9}, scores: {36.00,64.00,100.00,144.00,256.00,324.00}");
ASSERT_OK_AND_ASSIGN(RowRangeIndex deleted_shard_range, RowRangeIndex::Create({Range(10, 11)}));
ASSERT_OK_AND_ASSIGN(index_result, search(scan, deleted_shard_range, /*pre_filter=*/nullptr,
/*predicate=*/nullptr));
ASSERT_EQ(index_result->ToString(), "row ids: {}, scores: {}");
ASSERT_OK_AND_ASSIGN(RowRangeIndex first_shard_range, RowRangeIndex::Create({Range(0, 5)}));
ASSERT_OK_AND_ASSIGN(index_result, search(scan, first_shard_range, /*pre_filter=*/nullptr,
/*predicate=*/nullptr));
ASSERT_EQ(index_result->ToString(), "row ids: {3,4,5}, scores: {36.00,64.00,100.00}");
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({SpecialFields::ValueKind().ArrowField(), fields[0]}), R"([
[0, 3],
[0, 4],
[0, 5]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, {"f0"}, expected_array, /*predicate=*/nullptr, plan));
}
TEST_P(GlobalIndexTest, TestLuminaSearchSkipsUnindexedDeletionVectors) {
arrow::FieldVector fields = {arrow::field("f0", arrow::int32()),
arrow::field("f1", arrow::list(arrow::float32()))};
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::DELETION_VECTORS_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto indexed_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[0, [0.0, 0.0, 0.0, 0.0]],
[1, [1.0, 1.0, 1.0, 1.0]]
])")
.ValueOrDie();
auto unindexed_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[2, [2.0, 2.0, 2.0, 2.0]],
[3, [3.0, 3.0, 3.0, 3.0]]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto indexed_commit_msgs,
WriteArray(table_path, schema->field_names(), indexed_array));
ASSERT_OK(Commit(table_path, indexed_commit_msgs));
ASSERT_OK_AND_ASSIGN(auto unindexed_commit_msgs,
WriteArray(table_path, schema->field_names(), unindexed_array));
ASSERT_OK(Commit(table_path, unindexed_commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina", lumina_options,
Range(0, 1)));
ASSERT_OK_AND_ASSIGN(auto split, ScanData(table_path, /*partition_filters=*/{}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
DeletionVectorTestHelper::RetrieveAnchorFileNames(split->DataFiles()));
ASSERT_EQ(anchor_file_names.size(), 2u);
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<CommitMessage> dv_commit_msg,
DeletionVectorTestHelper::CreateDeletionVectorCommitMessage(
dir_->GetFileSystem(), table_path, file_format_, unindexed_commit_msgs[0],
{{anchor_file_names[1], /*deleted_positions=*/{1}}}, pool_));
ASSERT_OK(Commit(table_path, {dv_commit_msg}));
ASSERT_OK_AND_ASSIGN(split, ScanData(table_path, /*partition_filters=*/{}));
std::vector<std::string> deletion_file_paths;
for (const std::optional<DeletionFile>& deletion_file : split->DeletionFiles()) {
if (deletion_file) {
deletion_file_paths.push_back(deletion_file->path);
}
}
ASSERT_EQ(deletion_file_paths.size(), 1u);
ASSERT_OK(dir_->GetFileSystem()->Delete(deletion_file_paths[0]));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lumina_options, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> reader,
scan->CreateReader("f1", "lumina", /*row_range_index=*/std::nullopt));
ASSERT_TRUE(reader);
std::vector<float> query = {0.0f, 0.0f, 0.0f, 0.0f};
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<ScoredGlobalIndexResult> index_result,
reader->VisitVectorSearch(std::make_shared<VectorSearch>(
"f1", /*limit=*/2, query, /*pre_filter=*/nullptr,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, lumina_options)));
ASSERT_EQ(index_result->ToString(), "row ids: {0,1}, scores: {0.00,4.00}");
}
#endif
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithExternalPath) {
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1],
["Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1],
["Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1],
["Tony", [11.0, 10.0, 11.0, 10.0], 20, 17.1]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write and commit bitmap global index
auto external_dir1 = UniqueTestDirectory::Create("local");
std::map<std::string, std::string> index_options = {
{"global-index.external-path", "FILE://" + external_dir1->Str()}};
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap", index_options,
Range(0, 8)));
auto result_fields = fields;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
// test scan and read
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate, index_options));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
[0, "Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
TEST_P(GlobalIndexTest, TestIOException) {
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
auto schema = arrow::schema(fields);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Alice", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1]
])")
.ValueOrDie();
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
std::string table_path;
bool write_run_complete = false;
auto io_hook = IOHook::GetInstance();
for (size_t i = 0; i < 2000; i += paimon::test::RandomNumber(20, 30)) {
ScopeGuard guard([&io_hook]() { io_hook->Clear(); });
dir_ = UniqueTestDirectory::Create("local");
// create table and write data
CreateTable(/*partition_keys=*/{}, schema, options);
table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
io_hook->Reset(i, IOHook::Mode::RETURN_ERROR);
// write bitmap index
auto bitmap_index_write_status =
WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap",
/*options=*/{}, Range(0, 3));
CHECK_HOOK_STATUS(bitmap_index_write_status, i);
write_run_complete = true;
break;
}
ASSERT_TRUE(write_run_complete);
// read for bitmap
bool read_run_complete = false;
for (size_t i = 0; i < 2000; i += paimon::test::RandomNumber(20, 30)) {
ScopeGuard guard([&io_hook]() { io_hook->Clear(); });
io_hook->Reset(i, IOHook::Mode::RETURN_ERROR);
auto result_fields = fields;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
[0, "Alice", [1.0, 0.0, 1.0, 0.0], 10, 13.1]
])")
.ValueOrDie();
auto plan_result = ScanGlobalIndexAndData(table_path, predicate);
CHECK_HOOK_STATUS_WITHOUT_MESSAGE_CHECK(plan_result.status());
auto plan = std::move(plan_result).value();
auto read_status = ReadData(table_path, write_cols, expected_array, predicate, plan);
CHECK_HOOK_STATUS(read_status, i);
read_run_complete = true;
break;
}
ASSERT_TRUE(read_run_complete);
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithRangeBitmap) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Emily", 15, 0, 13.1],
["Tony", 20, 0, 14.1],
["Lucy", 20, 1, 15.1],
["Bob", 25, 1, 16.1],
["Tony", 30, 0, 17.1],
["Alice", 30, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_all_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Bob", 10, 1, 12.1],
[0, "Emily", 15, 0, 13.1],
[0, "Tony", 20, 0, 14.1],
[0, "Lucy", 20, 1, 15.1],
[0, "Bob", 25, 1, 16.1],
[0, "Tony", 30, 0, 17.1],
[0, "Alice", 30, null, 18.1]
])")
.ValueOrDie();
{
// read when no index is built
auto predicate = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, expected_all_array, predicate, plan));
}
// write and commit range-bitmap global index on f1 (int32)
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "range-bitmap", /*options=*/{},
Range(0, 7)));
// scan and read with range-bitmap global index
{
// f1 < 20: rows with f1=10,10,15 -> Alice(10), Bob(10), Emily(15)
auto predicate = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Bob", 10, 1, 12.1],
[0, "Emily", 15, 0, 13.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// f1 >= 25: rows with f1=25,30,30 -> Bob(25), Tony(30), Alice(30)
auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(25));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", 25, 1, 16.1],
[0, "Tony", 30, 0, 17.1],
[0, "Alice", 30, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// f1 > 30: no rows match
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(30));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, /*expected_array=*/nullptr, predicate, plan));
}
{
// f1 <= 10: rows with f1=10,10 -> Alice(10), Bob(10)
auto predicate = PredicateBuilder::LessOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Bob", 10, 1, 12.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// no predicate: return all rows
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr));
ASSERT_OK(
ReadData(table_path, write_cols, expected_all_array, /*predicate=*/nullptr, plan));
}
{
// f0 does not have range-bitmap index, should return all rows
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
ASSERT_OK(ReadData(table_path, write_cols, expected_all_array, predicate, plan));
}
}
TEST_P(GlobalIndexTest, TestDataEvolutionBatchScanWithRangeBitmapAndBitmap) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Alice", 30, null, 18.1],
["Bob", 10, 1, 12.1],
["Bob", 25, 1, 16.1],
["Emily", 15, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 20, 0, 14.1],
["Tony", 30, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_all_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Alice", 30, null, 18.1],
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 25, 1, 16.1],
[0, "Emily", 15, 0, 13.1],
[0, "Lucy", 20, 1, 15.1],
[0, "Tony", 20, 0, 14.1],
[0, "Tony", 30, 0, 17.1]
])")
.ValueOrDie();
// write and commit bitmap global index on f0 (string)
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap", /*options=*/{},
Range(0, 7)));
// write and commit range-bitmap global index on f1 (int32)
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "range-bitmap", /*options=*/{},
Range(0, 7)));
// scan and read with both indexes
{
// bitmap: f0 == "Alice" AND range-bitmap: f1 < 20
// Alice has f1=10 and f1=30, only f1=10 < 20 -> Alice(10)
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// bitmap: f0 == "Bob" AND range-bitmap: f1 >= 20
// Bob has f1=10 and f1=25, only f1=25 >= 20 -> Bob(25)
auto predicate1 =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
auto predicate2 = PredicateBuilder::GreaterOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({predicate1, predicate2}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", 25, 1, 16.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// bitmap: (f0 == "Alice" OR f0 == "Tony") AND range-bitmap: f1 > 20
// Alice: f1=10,30 -> f1=30 > 20; Tony: f1=20,30 -> f1=30 > 20
auto predicate_alice =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
auto predicate_tony =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Tony", 4));
ASSERT_OK_AND_ASSIGN(auto or_predicate,
PredicateBuilder::Or({predicate_alice, predicate_tony}));
auto range_predicate = PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate,
PredicateBuilder::And({or_predicate, range_predicate}));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 30, null, 18.1],
[0, "Tony", 30, 0, 17.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// only bitmap predicate: f0 == "Emily"
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Emily", 5));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Emily", 15, 0, 13.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// only range-bitmap predicate: f1 <= 15
auto predicate = PredicateBuilder::LessOrEqual(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(15));
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, predicate));
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1],
[0, "Bob", 10, 1, 12.1],
[0, "Emily", 15, 0, 13.1]
])")
.ValueOrDie();
ASSERT_OK(ReadData(table_path, write_cols, expected_array, predicate, plan));
}
{
// no predicate: return all rows
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr));
ASSERT_OK(
ReadData(table_path, write_cols, expected_all_array, /*predicate=*/nullptr, plan));
}
}
#ifdef PAIMON_ENABLE_LUCENE
TEST_P(GlobalIndexTest, TestLuceneWriteCommitScanReadIndexWithScore) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
auto tmp_dir = paimon::test::UniqueTestDirectory::Create();
ASSERT_TRUE(tmp_dir);
std::map<std::string, std::string> lucene_options = {
{"lucene-fts.write.omit-term-freq-and-position", "false"},
{"lucene-fts.write.tmp.directory", tmp_dir->Str()}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["This is an test document.", 0],
["This is an new document document document.", 1],
["Document document document document test.", 2],
["unordered user-defined doc id", 3]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write and commit lucene index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "lucene-fts",
/*options=*/lucene_options, Range(0, 3)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// test f0 field
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
auto index_reader = index_readers[0];
auto search = [&](const std::string& query,
const std::optional<RoaringBitmap64>& include_row_ids,
const std::string& expected) {
ASSERT_OK_AND_ASSIGN(auto index_result,
index_reader->VisitFullTextSearch(std::make_shared<FullTextSearch>(
"f0", query, /*limit=*/10, include_row_ids)));
ASSERT_TRUE(index_result->ToString().find("row ids: " + expected + ", scores: ") !=
std::string::npos)
<< query << ": " << index_result->ToString();
};
search(R"({"match":{"query":"document","operator":"And"}})", std::nullopt, "{0,1,2}");
search(R"({"match":{"query":"document","operator":"And"}})", RoaringBitmap64::From({1, 2, 3}),
"{1,2}");
search(R"({"match_phrase":{"query":"test document"}})", std::nullopt, "{0}");
search(R"({"boolean":{"should":[{"match":{"query":"unordered"}},{"match":{"query":"new"}}]}})",
std::nullopt, "{1,3}");
}
TEST_P(GlobalIndexTest, TestWriteCommitScanReadLuceneIndexWithPartition) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
auto tmp_dir = paimon::test::UniqueTestDirectory::Create();
ASSERT_TRUE(tmp_dir);
std::map<std::string, std::string> lucene_options = {
{"lucene-fts.write.omit-term-freq-and-position", "false"},
{"lucene-fts.write.tmp.directory", tmp_dir->Str()}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{"f1"}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto write_data_and_index = [&](const std::shared_ptr<arrow::Array>& src_array,
const std::map<std::string, std::string>& partition,
const Range& expected_range) {
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, partition, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// write lucene index
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{partition}, "f0", "lucene-fts",
/*options=*/lucene_options, expected_range));
};
// write partition f1 = 10
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["This is an test document.", 10],
["This is an new document document document.", 10]
])")
.ValueOrDie();
write_data_and_index(src_array1, {{"f1", "10"}}, Range(0, 1));
// write partition f1 = 20
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Document document document document test.", 20],
["unordered user-defined doc id", 20]
])")
.ValueOrDie();
write_data_and_index(src_array2, {{"f1", "20"}}, Range(2, 3));
auto scan_and_check_result = [&](const std::map<std::string, std::string>& partition,
const std::optional<RowRangeIndex>& row_range_index,
const std::optional<RoaringBitmap64>& include_row_ids,
const std::string& index_expected) {
std::vector<std::map<std::string, std::string>> partitions = {partition};
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, partitions,
lucene_options, fs_, /*executor=*/nullptr, pool_));
// check lucene index
ASSERT_OK_AND_ASSIGN(auto readers, global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto index_result,
readers[0]->VisitFullTextSearch(std::make_shared<FullTextSearch>(
"f0", R"({"match":{"query":"document","operator":"And"}})",
/*limit=*/10, include_row_ids)));
ASSERT_TRUE(index_result->ToString().find(index_expected) != std::string::npos);
};
{
// test scan and read for f1=10
auto filter = RoaringBitmap64::From(std::vector<int64_t>({0l}));
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(0, 1)}));
scan_and_check_result({{"f1", "10"}}, row_range_index, filter, "row ids: {0}");
}
{
// test scan and read for f1=20
auto filter = RoaringBitmap64::From(std::vector<int64_t>({2l}));
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(2, 3)}));
scan_and_check_result({{"f1", "20"}}, row_range_index, filter, "row ids: {2}");
}
{
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, lucene_options, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(
auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers.size(), 1u);
std::string query = R"({"match":{"query":"document","operator":"And"}})";
CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2});
CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2});
}
CheckTableSearchWithPartitionFilter(table_path,
R"({"match":{"query":"document","operator":"And"}})");
}
TEST_P(GlobalIndexTest, TestLuceneSearchOverlappingIndexRanges) {
auto tmp_dir = paimon::test::UniqueTestDirectory::Create();
ASSERT_TRUE(tmp_dir);
CheckSearchOverlappingIndexRanges("lucene-fts",
{{"lucene-fts.write.omit-term-freq-and-position", "false"},
{"lucene-fts.write.tmp.directory", tmp_dir->Str()}},
R"({"match":{"query":"document","operator":"And"}})");
}
TEST_P(GlobalIndexTest, TestLuceneSearchSkipsRowsDeletedByDeletionVectors) {
auto tmp_dir = paimon::test::UniqueTestDirectory::Create();
ASSERT_TRUE(tmp_dir);
CheckSearchSkipsDeletedRows("lucene-fts",
{{"lucene-fts.write.omit-term-freq-and-position", "false"},
{"lucene-fts.write.tmp.directory", tmp_dir->Str()}});
}
TEST_P(GlobalIndexTest, TestLuceneTableFullTextSearch) {
auto tmp_dir = paimon::test::UniqueTestDirectory::Create();
ASSERT_TRUE(tmp_dir);
CheckTableFullTextSearch("lucene-fts",
{{"lucene-fts.write.omit-term-freq-and-position", "false"},
{"lucene-fts.write.tmp.directory", tmp_dir->Str()}});
}
#endif
#ifdef PAIMON_ENABLE_FULL_TEXT
TEST_P(GlobalIndexTest, TestFullTextWriteCommitScanReadIndex) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
std::map<std::string, std::string> full_text_options = {{"full-text.tokenizer", "default"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["This is an test document.", 0],
["This is an new document document document.", 1],
["Document document document document test.", 2],
["unordered user-defined doc id", 3],
[null, 4]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "full-text",
/*options=*/full_text_options, Range(0, 4)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
auto index_reader = index_readers[0];
auto search = [&](const std::string& query,
const std::optional<RoaringBitmap64>& include_row_ids,
const std::string& expected) {
ASSERT_OK_AND_ASSIGN(auto index_result,
index_reader->VisitFullTextSearch(std::make_shared<FullTextSearch>(
"f0", query, /*limit=*/10, include_row_ids)));
ASSERT_TRUE(index_result->ToString().find("row ids: " + expected + ", scores: ") !=
std::string::npos)
<< query << ": " << index_result->ToString();
};
search(R"({"match":{"query":"document"}})", std::nullopt, "{0,1,2}");
search(R"({"match":{"query":"document"}})", RoaringBitmap64::From({1, 2, 3}), "{1,2}");
search(R"({"match":{"query":"new document","operator":"And"}})", std::nullopt, "{1}");
search(R"({"match_phrase":{"query":"test document"}})", std::nullopt, "{0}");
search(R"({"match":{"query":"unordered"}})", std::nullopt, "{3}");
}
TEST_P(GlobalIndexTest, TestWriteCommitScanReadFullTextIndexWithPartition) {
arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()),
arrow::field("f1", arrow::int32())};
std::map<std::string, std::string> full_text_options = {{"full-text.tokenizer", "default"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{"f1"}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
auto write_data_and_index = [&](const std::shared_ptr<arrow::Array>& src_array,
const std::map<std::string, std::string>& partition,
const Range& expected_range) {
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, partition, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{partition}, "f0", "full-text",
/*options=*/full_text_options, expected_range));
};
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["This is an test document.", 10],
["This is an new document document document.", 10]
])")
.ValueOrDie();
write_data_and_index(src_array1, {{"f1", "10"}}, Range(0, 1));
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Document document document document test.", 20],
["unordered user-defined doc id", 20]
])")
.ValueOrDie();
write_data_and_index(src_array2, {{"f1", "20"}}, Range(2, 3));
auto scan_and_check_result = [&](const std::map<std::string, std::string>& partition,
const std::optional<RowRangeIndex>& row_range_index,
const std::optional<RoaringBitmap64>& include_row_ids,
const std::string& index_expected) {
std::vector<std::map<std::string, std::string>> partitions = {partition};
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, partitions,
/*options=*/{}, fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto readers, global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(
auto index_result,
readers[0]->VisitFullTextSearch(std::make_shared<FullTextSearch>(
"f0", R"({"match":{"query":"document"}})", /*limit=*/10, include_row_ids)));
ASSERT_TRUE(index_result->ToString().find("row ids: " + index_expected + ", scores: ") !=
std::string::npos)
<< index_result->ToString();
};
{
auto filter = RoaringBitmap64::From(std::vector<int64_t>({0l}));
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(0, 1)}));
scan_and_check_result({{"f1", "10"}}, row_range_index, filter, "{0}");
}
{
auto filter = RoaringBitmap64::From(std::vector<int64_t>({2l}));
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(2, 3)}));
scan_and_check_result({{"f1", "20"}}, row_range_index, filter, "{2}");
}
{
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(2, 3)}));
scan_and_check_result({{"f1", "20"}}, row_range_index, std::nullopt, "{2}");
}
{
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(
auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(readers.size(), 1u);
std::string query = R"({"match":{"query":"document"}})";
CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2});
CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2});
}
CheckTableSearchWithPartitionFilter(table_path, R"({"match":{"query":"document"}})");
}
TEST_P(GlobalIndexTest, TestFullTextSearchOverlappingIndexRanges) {
CheckSearchOverlappingIndexRanges("full-text", {{"full-text.tokenizer", "default"}},
R"({"match":{"query":"document"}})");
}
TEST_P(GlobalIndexTest, TestFullTextSearchSkipsRowsDeletedByDeletionVectors) {
CheckSearchSkipsDeletedRows("full-text", {{"full-text.tokenizer", "default"}});
}
TEST_P(GlobalIndexTest, TestFullTextTableFullTextSearch) {
CheckTableFullTextSearch("full-text", {{"full-text.tokenizer", "default"}});
}
#endif
TEST_P(GlobalIndexTest, TestBTreeWriteCommitScanReadIndex) {
// BTreeGlobalIndexWriter requires keys to be written in monotonically increasing order.
// Therefore the source data must be pre-sorted by the indexed column (f0, string).
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
// Data sorted by f0 (string, ascending): Alice < Bob < Bob < Emily < Lucy < Tony < Tony
// The last row has f0=null which is treated separately by the null bitmap.
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 20, 0, 16.1],
["Emily", 10, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 10, 0, 14.1],
["Tony", 20, 0, 17.1],
[null, 20, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// Write btree-global index on f0
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "btree",
/*options=*/{}, Range(0, 7)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
auto index_reader = index_readers[0];
{
// VisitEqual: "Alice" -> row 0
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0}");
}
{
// VisitEqual: "Bob" -> rows 1,2
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
{
// VisitEqual: non-existent key -> empty
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitEqual(Literal(FieldType::STRING, "Zara", 4)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{}");
}
{
// VisitNotEqual: "Bob" -> all non-null except Bob rows -> {0,3,4,5,6}
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitNotEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,3,4,5,6}");
}
{
// VisitIsNull -> row 7 (null key)
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitIsNull());
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{7}");
}
{
// VisitIsNotNull -> rows 0-6
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitIsNotNull());
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,1,2,3,4,5,6}");
}
{
// VisitIn: {"Alice", "Lucy"} -> rows {0, 4}
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitIn({Literal(FieldType::STRING, "Alice", 5),
Literal(FieldType::STRING, "Lucy", 4)}));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,4}");
}
{
// VisitNotIn: {"Alice", "Lucy"} -> all non-null except {0,4} -> {1,2,3,5,6}
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitNotIn({Literal(FieldType::STRING, "Alice", 5),
Literal(FieldType::STRING, "Lucy", 4)}));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2,3,5,6}");
}
{
// VisitLessThan: "Emily" -> keys < "Emily" -> Alice(0), Bob(1,2) -> {0,1,2}
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitLessThan(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,1,2}");
}
{
// VisitLessOrEqual: "Emily" -> keys <= "Emily" -> Alice(0), Bob(1,2), Emily(3) ->
// {0,1,2,3}
ASSERT_OK_AND_ASSIGN(
auto result, index_reader->VisitLessOrEqual(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,1,2,3}");
}
{
// VisitGreaterThan: "Emily" -> keys > "Emily" -> Lucy(4), Tony(5,6) -> {4,5,6}
ASSERT_OK_AND_ASSIGN(
auto result, index_reader->VisitGreaterThan(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{4,5,6}");
}
{
// VisitGreaterOrEqual: "Emily" -> keys >= "Emily" -> Emily(3), Lucy(4), Tony(5,6) ->
// {3,4,5,6}
ASSERT_OK_AND_ASSIGN(
auto result, index_reader->VisitGreaterOrEqual(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{3,4,5,6}");
}
auto scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scan_impl);
{
// Equal predicate via evaluator
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Tony", 4));
ASSERT_OK_AND_ASSIGN(auto result, scan_impl->Scan(predicate));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{5,6}");
}
{
// AND predicate: f0 == "Bob" AND f1 == 20
// f0 == "Bob" -> {1,2}, but f1 index does not exist -> AND yields {1,2}
// (fields without index return nullptr, AND with nullptr keeps the other side)
auto f0_predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
auto f1_predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(20));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_predicate, f1_predicate}));
ASSERT_OK_AND_ASSIGN(auto result, scan_impl->Scan(predicate));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
{
// row_range_index filtering: range [0,2] should only load that range
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(0, 2)}));
ASSERT_OK_AND_ASSIGN(auto range_readers,
global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(range_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto result,
range_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0}");
}
{
// Invalid row_range_index: no intersection -> empty readers
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index,
RowRangeIndex::Create({Range(100, 200)}));
ASSERT_OK_AND_ASSIGN(auto range_readers,
global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(range_readers.size(), 0u);
}
// Test full pipeline: scan with predicate -> read data
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
auto scan_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_OK_AND_ASSIGN(auto index_result, scan_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{1,2}");
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 20, 0, 16.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
}
TEST_P(GlobalIndexTest, TestBTreeEmptyStringKeyPredicates) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
// Null keys are tracked separately. The non-null keys remain monotonically increasing.
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[null, 10, 0, 10.0],
["", 20, 1, 20.0],
["abc", 30, 2, 30.0]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "btree",
/*options=*/{}, Range(0, 2)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
auto index_reader = index_readers[0];
Literal empty(FieldType::STRING, "", 0);
Literal abc(FieldType::STRING, "abc", 3);
Literal missing(FieldType::STRING, "missing", 7);
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitEqual(empty));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitEqual(abc));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitEqual(missing));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitNotEqual(empty));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitIsNull());
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitIsNotNull());
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitIn({empty, abc, missing}));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitNotIn({empty, missing}));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitLessThan(abc));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitLessOrEqual(empty));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitGreaterThan(empty));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{2}");
}
{
ASSERT_OK_AND_ASSIGN(auto result, index_reader->VisitGreaterOrEqual(abc));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{2}");
}
}
TEST_P(GlobalIndexTest, TestBTreeWriteCommitScanReadIndexWithPartition) {
// BTree index with partitioned table. Each partition's data is sorted by f0 independently.
auto schema = arrow::schema(fields_);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{"f1"}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
// Write partition f1=10. Data sorted by f0: Alice < Bob < Bob < Emily < Tony
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 0, 13.1],
["Emily", 10, 0, 14.1],
["Tony", 10, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "btree",
/*options=*/{}, Range(0, 4)));
// Write partition f1=20. Data sorted by f0: Alice < Lucy < Tony
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 16.1],
["Lucy", 20, 1, 17.1],
["Tony", 20, 0, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}}}, "f0", "btree",
/*options=*/{}, Range(5, 7)));
// Scan all partitions
{
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", std::nullopt));
// One reader per partition range -> 2 ranges -> UnionGlobalIndexReader wraps them
ASSERT_EQ(index_readers.size(), 1u);
// "Alice" exists in both partitions: local ids {0} in range [0,4] -> global 0,
// and local ids {0} in range [5,7] -> global 5
ASSERT_OK_AND_ASSIGN(auto result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,5}");
// "Bob" only in partition f1=10: local ids {1,2} -> global {1,2}
ASSERT_OK_AND_ASSIGN(auto result2,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result2);
ASSERT_EQ(result2->ToString(), "{1,2}");
// "Lucy" only in partition f1=20: local ids {1} -> global {6}
ASSERT_OK_AND_ASSIGN(auto result3,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Lucy", 4)));
ASSERT_TRUE(result3);
ASSERT_EQ(result3->ToString(), "{6}");
}
// Scan with partition filter: only f1=10
{
std::vector<std::map<std::string, std::string>> partitions = {{{"f1", "10"}}};
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, partitions,
/*options=*/{}, fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
// "Alice" in f1=10 only -> global {0}
ASSERT_OK_AND_ASSIGN(auto result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0}");
// "Lucy" not in f1=10 -> empty
ASSERT_OK_AND_ASSIGN(auto result2,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Lucy", 4)));
ASSERT_TRUE(result2);
ASSERT_EQ(result2->ToString(), "{}");
}
// Scan with row_range_index filtering: only global row 6 in range [5,7]. This verifies
// that the global range is translated to the file-local range [1,1] before creating the
// BTree reader.
{
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(6, 6)}));
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", row_range_index));
ASSERT_EQ(index_readers.size(), 1u);
// "Lucy" is local row 1 in the index range [5,7] -> global row 6.
ASSERT_OK_AND_ASSIGN(auto result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Lucy", 4)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{6}");
// "Tony" is in the selected index file but outside the requested row range.
ASSERT_OK_AND_ASSIGN(result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Tony", 4)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{}");
}
// Full pipeline with evaluator: Scan(predicate) -> read data
{
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{},
fs_, /*executor=*/nullptr, pool_));
auto scanner_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scanner_impl);
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Tony", 4));
ASSERT_OK_AND_ASSIGN(auto index_result, scanner_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{4,7}");
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Tony", 10, 1, 15.1],
[0, "Tony", 20, 0, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
}
TEST_P(GlobalIndexTest, TestBTreeWithPartitionAndCustomExecutor) {
// Test that UnionGlobalIndexReader uses a custom 8-thread executor to read
// btree indexes from two partitions in parallel.
auto schema = arrow::schema(fields_);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{"f1"}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
// Write partition f1=10 (5 rows, sorted by f0)
auto src_array1 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 10, 0, 13.1],
["Emily", 10, 0, 14.1],
["Tony", 10, 1, 15.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "10"}}, write_cols, src_array1));
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "10"}}}, "f0", "btree",
/*options=*/{}, Range(0, 4)));
// Write partition f1=20 (3 rows, sorted by f0)
auto src_array2 = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 20, null, 16.1],
["Lucy", 20, 1, 17.1],
["Tony", 20, 0, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs2,
WriteArray(table_path, {{"f1", "20"}}, write_cols, src_array2));
ASSERT_OK(Commit(table_path, commit_msgs2));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{{{"f1", "20"}}}, "f0", "btree",
/*options=*/{}, Range(5, 7)));
// Create a GlobalIndexScan with an explicit 8-thread executor
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Executor> executor,
CreateDefaultExecutor(/*thread_count=*/8));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_, executor, pool_));
// CreateReaders should return 1 UnionGlobalIndexReader (2 sub-readers for 2 ranges)
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
auto union_reader = std::dynamic_pointer_cast<UnionGlobalIndexReader>(index_readers[0]);
ASSERT_TRUE(union_reader);
ASSERT_EQ(union_reader->executor_, executor);
// "Alice" in both partitions: global ids {0, 5}
ASSERT_OK_AND_ASSIGN(auto result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,5}");
// "Bob" only in f1=10: global ids {1, 2}
ASSERT_OK_AND_ASSIGN(auto result2,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result2);
ASSERT_EQ(result2->ToString(), "{1,2}");
// "Lucy" only in f1=20: global id {6}
ASSERT_OK_AND_ASSIGN(auto result3,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Lucy", 4)));
ASSERT_TRUE(result3);
ASSERT_EQ(result3->ToString(), "{6}");
// Full pipeline: evaluator with the 8-thread executor
auto scanner_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scanner_impl);
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Tony", 4));
ASSERT_OK_AND_ASSIGN(auto index_result, scanner_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{4,7}");
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Tony", 10, 1, 15.1],
[0, "Tony", 20, 0, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
TEST_P(GlobalIndexTest, TestBTreeAndBitmapCoexist) {
// Test btree-global and bitmap index coexisting on the same field (f0).
// The evaluator should AND their results, producing the intersection.
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
std::vector<std::string> write_cols = schema->field_names();
// Data sorted by f0 for btree: Alice < Bob < Bob < Emily < Lucy < Tony < Tony
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
["Alice", 10, 1, 11.1],
["Bob", 10, 1, 12.1],
["Bob", 25, 1, 16.1],
["Emily", 15, 0, 13.1],
["Lucy", 20, 1, 15.1],
["Tony", 20, 0, 14.1],
["Tony", 30, 0, 17.1],
[null, 30, null, 18.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// Build both indexes on f0
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "btree",
/*options=*/{}, Range(0, 7)));
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "bitmap",
/*options=*/{}, Range(0, 7)));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
// Two index types on f0 -> 2 readers
ASSERT_OK_AND_ASSIGN(auto index_readers, global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 2u);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> btree_reader,
global_index_scan->CreateReader("f0", "btree", std::nullopt));
ASSERT_TRUE(btree_reader);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> btree_result,
btree_reader->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(btree_result);
ASSERT_EQ(btree_result->ToString(), "{1,2}");
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> btree_less_than_result,
btree_reader->VisitLessThan(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(btree_less_than_result);
ASSERT_EQ(btree_less_than_result->ToString(), "{0,1,2}");
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> bitmap_reader,
global_index_scan->CreateReader("f0", "bitmap", std::nullopt));
ASSERT_TRUE(bitmap_reader);
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> bitmap_result,
bitmap_reader->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(bitmap_result);
ASSERT_EQ(bitmap_result->ToString(), "{1,2}");
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexResult> bitmap_less_than_result,
bitmap_reader->VisitLessThan(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(bitmap_less_than_result);
ASSERT_EQ(bitmap_less_than_result->ToString(), "{0,1,2}");
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexReader> missing_reader,
global_index_scan->CreateReader("f0", "lucene", std::nullopt));
ASSERT_FALSE(missing_reader);
// Each reader individually should return the same result for Equal("Bob")
for (const auto& index_reader : index_readers) {
ASSERT_OK_AND_ASSIGN(auto result,
index_reader->VisitEqual(Literal(FieldType::STRING, "Bob", 3)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
// Via evaluator: the two indexes' results get AND, still {1,2}
auto scanner_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scanner_impl);
ASSERT_OK_AND_ASSIGN(auto evaluator, scanner_impl->GetOrCreateIndexEvaluator());
{
// Equal predicate
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto result, evaluator->Evaluate(predicate));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{1,2}");
}
{
// NotEqual predicate: both indexes agree on non-null, non-"Bob" rows -> {0,3,4,5,6}
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto result, evaluator->Evaluate(predicate));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,3,4,5,6}");
}
{
// IsNull: both agree on row 7
auto predicate =
PredicateBuilder::IsNull(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING);
ASSERT_OK_AND_ASSIGN(auto result, evaluator->Evaluate(predicate));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{7}");
}
// Full pipeline: f0 == "Alice" -> read data
{
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Alice", 5));
ASSERT_OK_AND_ASSIGN(auto index_result, scanner_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{0}");
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Alice", 10, 1, 11.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
// Full pipeline with AND across btree(f0) and bitmap(f0):
// both indexes support Equal on f0
// So AND(Equal, unsupported field) keeps the result evaluated from f0
{
// f0 == "Bob" AND f1 == 10 (f1 has no index -> nullptr -> keeps btree+bitmap result)
auto f0_pred =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
auto f1_pred = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_pred, f1_pred}));
ASSERT_OK_AND_ASSIGN(auto index_result, scanner_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{1,2}");
auto result_fields = fields_;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", 10, 1, 12.1],
[0, "Bob", 25, 1, 16.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
}
TEST_P(GlobalIndexTest, TestBTreeScanWithPartitionWithMultiMeta) {
if (file_format_ == "avro") {
return;
}
std::string table_path =
paimon::test::GetDataDir() + "/" + file_format_ +
"/append_with_btree_with_partition.db/append_with_btree_with_partition";
ASSERT_OK_AND_ASSIGN(std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt, /*options=*/{}, fs_,
/*executor=*/nullptr, pool_));
auto count_rows = [](const std::shared_ptr<GlobalIndexResult>& result) -> int64_t {
EXPECT_TRUE(result);
EXPECT_OK_AND_ASSIGN(std::vector<Range> ranges, result->ToRanges());
int64_t total = 0;
for (const auto& range : ranges) {
total += range.Count();
}
return total;
};
auto get_reader = [&](const std::string& column) -> std::shared_ptr<GlobalIndexReader> {
EXPECT_OK_AND_ASSIGN(auto readers, global_index_scan->CreateReaders(column, std::nullopt));
EXPECT_EQ(readers.size(), 1u);
return readers[0];
};
// ---- col_boolean ----
{
auto reader = get_reader("col_boolean");
ASSERT_TRUE(reader);
ASSERT_OK_AND_ASSIGN(auto eq_true, reader->VisitEqual(Literal(true)));
ASSERT_EQ(count_rows(eq_true), 20);
ASSERT_OK_AND_ASSIGN(auto eq_false, reader->VisitEqual(Literal(false)));
ASSERT_EQ(count_rows(eq_false), 20);
}
// ---- col_int ----
{
auto reader = get_reader("col_int");
ASSERT_TRUE(reader);
ASSERT_OK_AND_ASSIGN(auto eq_15, reader->VisitEqual(Literal(15)));
ASSERT_EQ(count_rows(eq_15), 2);
ASSERT_OK_AND_ASSIGN(auto eq_missing, reader->VisitEqual(Literal(100)));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan(30): i*3 > 30 -> i in [11, 19], 9 indices per partition -> 18 rows.
ASSERT_OK_AND_ASSIGN(auto gt_30, reader->VisitGreaterThan(Literal(30)));
ASSERT_EQ(count_rows(gt_30), 18);
// GreaterThan(57): nothing greater than the max value.
ASSERT_OK_AND_ASSIGN(auto gt_max, reader->VisitGreaterThan(Literal(57)));
ASSERT_EQ(count_rows(gt_max), 0);
}
// ---- col_date (values are 18000 + i for i in [0,19]) ----
{
auto reader = get_reader("col_date");
ASSERT_TRUE(reader);
// 18005 is present at i=5 in both partitions.
ASSERT_OK_AND_ASSIGN(auto eq_present, reader->VisitEqual(Literal(FieldType::DATE, 18005)));
ASSERT_EQ(count_rows(eq_present), 2);
ASSERT_OK_AND_ASSIGN(auto eq_missing, reader->VisitEqual(Literal(FieldType::DATE, 17999)));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan(18010): i in [11, 19] -> 9 per partition -> 18 rows.
ASSERT_OK_AND_ASSIGN(auto gt_mid,
reader->VisitGreaterThan(Literal(FieldType::DATE, 18010)));
ASSERT_EQ(count_rows(gt_mid), 18);
}
// ---- col_double (values are i * 2.2 for i in [0,19]) ----
{
auto reader = get_reader("col_double");
ASSERT_TRUE(reader);
// i=5 -> 11.0
ASSERT_OK_AND_ASSIGN(auto eq_present, reader->VisitEqual(Literal(11.0)));
ASSERT_EQ(count_rows(eq_present), 2);
ASSERT_OK_AND_ASSIGN(auto eq_missing, reader->VisitEqual(Literal(123.456)));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan(10 * 2.2 = 22.0): i in [11, 19] -> 9 per partition -> 18 rows.
ASSERT_OK_AND_ASSIGN(auto gt_mid, reader->VisitGreaterThan(Literal(10 * 2.2)));
ASSERT_EQ(count_rows(gt_mid), 18);
}
// ---- col_timestamp (Timestamp from epoch millis = 1700000000000 + i*1000) ----
{
auto reader = get_reader("col_timestamp");
ASSERT_TRUE(reader);
// i=5 -> 1700000005000 ms.
ASSERT_OK_AND_ASSIGN(
auto eq_present,
reader->VisitEqual(Literal(Timestamp::FromEpochMillis(1700000000000L + 5 * 1000L))));
ASSERT_EQ(count_rows(eq_present), 2);
ASSERT_OK_AND_ASSIGN(auto eq_missing,
reader->VisitEqual(Literal(Timestamp::FromEpochMillis(1L))));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan(i=10 boundary): i in [11, 19] -> 18 rows globally.
ASSERT_OK_AND_ASSIGN(auto gt_mid,
reader->VisitGreaterThan(
Literal(Timestamp::FromEpochMillis(1700000000000L + 10 * 1000L))));
ASSERT_EQ(count_rows(gt_mid), 18);
}
// ---- col_timestamp_ltz (same physical values as col_timestamp) ----
{
auto reader = get_reader("col_timestamp_ltz");
ASSERT_TRUE(reader);
ASSERT_OK_AND_ASSIGN(
auto eq_present,
reader->VisitEqual(Literal(Timestamp::FromEpochMillis(1700000000000L + 7 * 1000L))));
ASSERT_EQ(count_rows(eq_present), 2);
ASSERT_OK_AND_ASSIGN(auto gt_mid,
reader->VisitGreaterThan(
Literal(Timestamp::FromEpochMillis(1700000000000L + 10 * 1000L))));
ASSERT_EQ(count_rows(gt_mid), 18);
}
// ---- col_decimal (unscaled = i * 123456, precision=18, scale=6) ----
{
auto reader = get_reader("col_decimal");
ASSERT_TRUE(reader);
// i=5 -> unscaled 617280
ASSERT_OK_AND_ASSIGN(
auto eq_present,
reader->VisitEqual(Literal(Decimal::FromUnscaledLong(5 * 123456L, 18, 6))));
ASSERT_EQ(count_rows(eq_present), 2);
ASSERT_OK_AND_ASSIGN(auto eq_missing,
reader->VisitEqual(Literal(Decimal::FromUnscaledLong(1L, 18, 6))));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan(i=10): i in [11, 19] -> 18 rows globally.
ASSERT_OK_AND_ASSIGN(
auto gt_mid,
reader->VisitGreaterThan(Literal(Decimal::FromUnscaledLong(10 * 123456L, 18, 6))));
ASSERT_EQ(count_rows(gt_mid), 18);
// Global index readers do not rescale Decimal literals. A mathematically equivalent
// literal with a different scale produces an incorrect empty result, so callers must use
// the field's scale.
ASSERT_OK_AND_ASSIGN(auto eq_same_value_different_scale,
reader->VisitEqual(Literal(Decimal::FromUnscaledLong(
5 * 1234560L, /*precision=*/18, /*scale=*/7))));
ASSERT_EQ(count_rows(eq_same_value_different_scale), 0);
}
// ---- col_string (values are "str_00000" .. "str_00019") ----
{
auto reader = get_reader("col_string");
ASSERT_TRUE(reader);
std::string present_value = "str_00005";
ASSERT_OK_AND_ASSIGN(auto eq_present,
reader->VisitEqual(Literal(FieldType::STRING, present_value.data(),
present_value.size())));
ASSERT_EQ(count_rows(eq_present), 2);
std::string missing_value = "str_99999";
ASSERT_OK_AND_ASSIGN(auto eq_missing,
reader->VisitEqual(Literal(FieldType::STRING, missing_value.data(),
missing_value.size())));
ASSERT_EQ(count_rows(eq_missing), 0);
// GreaterThan("str_00010"): lexicographically greater values are i in [11, 19].
std::string mid_value = "str_00010";
ASSERT_OK_AND_ASSIGN(auto gt_mid,
reader->VisitGreaterThan(
Literal(FieldType::STRING, mid_value.data(), mid_value.size())));
ASSERT_EQ(count_rows(gt_mid), 18);
}
}
#ifdef PAIMON_ENABLE_LUMINA
TEST_P(GlobalIndexTest, TestBTreeWithLumina) {
// Test btree on f0 (string) and lumina on f1 (vector) coexisting on different fields.
arrow::FieldVector fields = {
arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::list(arrow::float32())),
arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::float64())};
std::map<std::string, std::string> lumina_options = {{"lumina.index.dimension", "4"},
{"lumina.index.type", "bruteforce"},
{"lumina.distance.metric", "l2"},
{"lumina.encoding.type", "rawf32"},
{"lumina.search.parallel_number", "10"}};
auto schema = arrow::schema(fields);
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, file_format_},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
CreateTable(/*partition_keys=*/{}, schema, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
std::vector<std::string> write_cols = schema->field_names();
// Data sorted by f0 for btree: Alice < Alice < Bob < Bob < Emily < Lucy < Paul < Tony
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
["Alice", [0.0, 0.0, 0.0, 0.0], 10, 11.1],
["Alice", [11.0, 11.0, 11.0, 11.0], 20, 18.1],
["Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
["Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1],
["Emily", [1.0, 0.0, 1.0, 0.0], 10, 13.1],
["Lucy", [10.0, 10.0, 10.0, 10.0], 20, 15.1],
["Paul", [10.0, 10.0, 10.0, 10.0], 20, 19.1],
["Tony", [1.0, 1.0, 1.0, 1.0], 10, 14.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array));
ASSERT_OK(Commit(table_path, commit_msgs));
// Build btree index on f0
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "btree",
/*options=*/{}, Range(0, 7)));
// Build lumina index on f1
ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f1", "lumina",
/*options=*/lumina_options, Range(0, 7)));
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<GlobalIndexScan> global_index_scan,
GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt,
/*partitions=*/std::nullopt,
/*options=*/lumina_options, fs_, /*executor=*/nullptr, pool_));
// Query f0 via btree
{
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
ASSERT_OK_AND_ASSIGN(auto result,
index_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(result);
ASSERT_EQ(result->ToString(), "{0,1}");
ASSERT_OK_AND_ASSIGN(
auto result2, index_readers[0]->VisitLessThan(Literal(FieldType::STRING, "Emily", 5)));
ASSERT_TRUE(result2);
ASSERT_EQ(result2->ToString(), "{0,1,2,3}");
}
// Query f1 via lumina (vector search)
{
ASSERT_OK_AND_ASSIGN(auto index_readers,
global_index_scan->CreateReaders("f1", std::nullopt));
ASSERT_EQ(index_readers.size(), 1u);
std::vector<float> query = {11.0f, 11.0f, 11.0f, 11.0f};
auto vector_search = std::make_shared<VectorSearch>(
"f1", /*limit=*/1, query, /*filter=*/nullptr,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(auto scored_result,
index_readers[0]->VisitVectorSearch(vector_search));
ASSERT_TRUE(scored_result);
ASSERT_EQ(scored_result->ToString(), "row ids: {1}, scores: {0.00}");
}
// Evaluator: btree on f0 = "Bob"
{
auto scanner_impl = std::dynamic_pointer_cast<GlobalIndexScanImpl>(global_index_scan);
ASSERT_TRUE(scanner_impl);
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Bob", 3));
ASSERT_OK_AND_ASSIGN(auto index_result, scanner_impl->Scan(predicate));
ASSERT_TRUE(index_result);
ASSERT_EQ(index_result->ToString(), "{2,3}");
// Read data for Bob
auto result_fields = fields;
result_fields.insert(result_fields.begin(), SpecialFields::ValueKind().ArrowField());
auto expected_array =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), R"([
[0, "Bob", [0.0, 1.0, 0.0, 1.0], 10, 12.1],
[0, "Bob", [10.0, 11.0, 10.0, 11.0], 20, 16.1]
])")
.ValueOrDie();
ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr,
/*options=*/{}, index_result));
ASSERT_OK(ReadData(table_path, write_cols, expected_array, /*predicate=*/nullptr, plan));
}
// Combined: btree f0 filter + lumina vector search with pre-filter
// Use btree result as pre_filter for lumina search
{
ASSERT_OK_AND_ASSIGN(auto btree_readers,
global_index_scan->CreateReaders("f0", std::nullopt));
ASSERT_EQ(btree_readers.size(), 1u);
// Get rows where f0 == "Alice" -> {0, 1}
ASSERT_OK_AND_ASSIGN(auto btree_result,
btree_readers[0]->VisitEqual(Literal(FieldType::STRING, "Alice", 5)));
ASSERT_TRUE(btree_result);
ASSERT_EQ(btree_result->ToString(), "{0,1}");
// Now vector search on f1 with pre_filter limiting to Alice's rows {0, 1}
ASSERT_OK_AND_ASSIGN(auto lumina_readers,
global_index_scan->CreateReaders("f1", std::nullopt));
ASSERT_EQ(lumina_readers.size(), 1u);
std::vector<float> query = {11.0f, 11.0f, 11.0f, 11.0f};
auto filter = [](int64_t id) -> bool { return id == 0 || id == 1; };
auto vector_search = std::make_shared<VectorSearch>(
"f1", /*limit=*/1, query, filter,
/*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/lumina_options);
ASSERT_OK_AND_ASSIGN(auto scored_result,
lumina_readers[0]->VisitVectorSearch(vector_search));
ASSERT_EQ(scored_result->ToString(), "row ids: {1}, scores: {0.00}");
}
}
#endif
std::vector<ParamType> GetTestValuesForGlobalIndexTest() {
std::vector<ParamType> values;
values.emplace_back("parquet", false);
values.emplace_back("parquet", true);
#ifdef PAIMON_ENABLE_ORC
values.emplace_back("orc", false);
values.emplace_back("orc", true);
#endif
#ifdef PAIMON_ENABLE_AVRO
values.emplace_back("avro", false);
values.emplace_back("avro", true);
#endif
return values;
}
INSTANTIATE_TEST_SUITE_P(FileFormat, GlobalIndexTest,
::testing::ValuesIn(GetTestValuesForGlobalIndexTest()));
} // namespace paimon::test