Google Git Google Git
Sign in
apache/paimon-cpp/refs/heads/release-0.4/./test/inte/data_evolution_table_test.cpp
blob: 79e16a6f267ac035ff4f77264129b0e72ffbe316 [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 <iterator>
#include <set>
#include <tuple>
#include "arrow/type.h"
#include "fmt/format.h"
#include "gtest/gtest.h"
#include "paimon/common/factories/io_hook.h"
#include "paimon/common/io/cache/lru_cache.h"
#include "paimon/common/table/special_fields.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/checked_cast.h"
#include "paimon/common/utils/date_time_utils.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/common/utils/scope_guard.h"
#include "paimon/common/utils/vector_store_utils.h"
#include "paimon/core/core_options.h"
#include "paimon/core/global_index/global_index_live_row_filter.h"
#include "paimon/core/global_index/indexed_split_impl.h"
#include "paimon/core/operation/file_store_scan.h"
#include "paimon/core/schema/schema_manager.h"
#include "paimon/core/snapshot.h"
#include "paimon/core/table/source/data_split_impl.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/defs.h"
#include "paimon/executor.h"
#include "paimon/format/mosaic/mosaic_format_defs.h"
#include "paimon/fs/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/indexed_split.h"
#include "paimon/predicate/literal.h"
#include "paimon/predicate/predicate_builder.h"
#include "paimon/predicate/predicate_utils.h"
#include "paimon/result.h"
#include "paimon/status.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"
#include "paimon/utils/row_range_index.h"
namespace paimon::test {
using DataEvolutionTableParam = std::tuple<std::string, bool, bool>;
// This is a sdk end-to-end test for data evolution
class DataEvolutionTableTest : public ::testing::Test,
public ::testing::WithParamInterface<DataEvolutionTableParam> {
protected:
void SetUp() override {
dir_ = UniqueTestDirectory::Create("local");
int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs();
std::srand(seed);
}
void TearDown() override {
dir_.reset();
}
void CreateTable(const arrow::FieldVector& fields,
const std::vector<std::string>& partition_keys,
const std::map<std::string, std::string>& options) const {
auto schema = arrow::schema(fields);
::ArrowSchema c_schema;
ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok());
ASSERT_OK_AND_ASSIGN(auto catalog, Catalog::Create(dir_->Str(), {}));
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) const {
CreateTable(fields_, partition_keys, options);
}
void CreateTable(const std::vector<std::string>& partition_keys) const {
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
return CreateTable(partition_keys, options);
}
void CreateTable() const {
return CreateTable(/*partition_keys=*/{});
}
Result<std::vector<std::shared_ptr<CommitMessage>>> WriteArrays(
const std::string& table_path, const std::map<std::string, std::string>& partition,
const std::vector<std::string>& write_cols,
const std::vector<std::shared_ptr<arrow::Array>>& write_arrays,
const std::map<std::string, std::string>& write_options = {}) const {
// write
WriteContextBuilder write_builder(table_path, "commit_user_1");
write_builder.WithWriteSchema(write_cols);
for (const auto& [key, value] : write_options) {
write_builder.AddOption(key, value);
}
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)));
for (const auto& write_array : write_arrays) {
ArrowArray c_array;
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*write_array, &c_array));
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::map<std::string, std::string>& partition,
const std::vector<std::string>& write_cols,
const std::shared_ptr<arrow::Array>& write_array,
const std::map<std::string, std::string>& write_options = {}) const {
return WriteArrays(table_path, partition, write_cols, {write_array}, write_options);
}
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);
}
/// Stamps `reset_first_row_id` on the metas the caller holds. The commit assigns row ids
/// onto its own copies, so a caller that later looks a file up by row id range, to find a
/// row range group's anchor, has to mirror the assignment here.
void SetFirstRowId(int64_t reset_first_row_id,
std::vector<std::shared_ptr<CommitMessage>>& commit_msgs) const {
for (auto& commit_msg : commit_msgs) {
auto commit_msg_impl = std::dynamic_pointer_cast<CommitMessageImpl>(commit_msg);
ASSERT_TRUE(commit_msg_impl);
for (auto& file : commit_msg_impl->data_increment_.new_files_) {
file->AssignFirstRowId(reset_first_row_id);
}
}
}
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");
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);
}
/// Creates the default (f0, f1, f2) data evolution table and returns the options it was
/// created with, which a test that later evolves the schema has to pass along.
/// `extra_options` is merged in for tests that also pin a read batch or split target size.
std::map<std::string, std::string> CreateDataEvolutionTable(
bool deletion_vectors_enabled, const std::map<std::string, std::string>& extra_options = {},
const std::vector<std::string>& partition_keys = {}) const {
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
if (deletion_vectors_enabled) {
options.emplace(Options::DELETION_VECTORS_ENABLED, "true");
options.emplace(Options::DELETION_VECTOR_BITMAP64,
std::get<2>(GetParam()) ? "true" : "false");
}
options.insert(extra_options.begin(), extra_options.end());
CreateTable(partition_keys, options);
return options;
}
/// DeletionVectorTestHelper::CreateDeletionVectorCommitMessage, documented there, plus the
/// commit of the message it returns.
Result<std::shared_ptr<CommitMessage>> CommitDeletionVectors(
const std::string& table_path, const std::shared_ptr<CommitMessage>& base_commit_msg,
const std::map<std::string, std::vector<int64_t>>& deleted_positions_by_anchor,
const std::shared_ptr<CommitMessage>& replaced_commit_msg = nullptr) const {
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<CommitMessage> commit_message,
DeletionVectorTestHelper::CreateDeletionVectorCommitMessage(
dir_->GetFileSystem(), table_path, /*file_format_identifier=*/FileFormat(),
base_commit_msg, deleted_positions_by_anchor, GetDefaultPool(),
replaced_commit_msg));
PAIMON_RETURN_NOT_OK(Commit(table_path, {commit_message}));
return commit_message;
}
/// Writes one row range group holding `f0_values`, starting at `first_row_id`: a full-row
/// write plus a partial f2 write over the same rows, so the group's split merges columns
/// from two files and is not raw convertible. When `partition` is set the files land in that
/// partition and every row carries its value in f1, the partition key.
///
/// The f2 write starts only once the full-row write is committed, which is what makes the
/// f2 file the newer of the pair. Writing both before either commit left the column merge
/// free to serve f2 from the full-row write instead.
Result<std::vector<std::shared_ptr<CommitMessage>>> WriteAndCommitGroup(
const std::string& table_path, int64_t first_row_id, const std::vector<int32_t>& f0_values,
const std::map<std::string, std::string>& partition = {},
const std::map<std::string, std::string>& partial_write_options = {}) const {
auto partition_f1 = partition.find("f1");
std::string base_json = "[";
std::string f2_json = "[";
for (size_t i = 0; i < f0_values.size(); i++) {
if (i > 0) {
base_json += ", ";
f2_json += ", ";
}
std::string f1_value =
partition_f1 == partition.end() ? fmt::format("a{}", i) : partition_f1->second;
base_json += fmt::format(R"([{}, "{}", "x{}"])", f0_values[i], f1_value, i);
f2_json += fmt::format(R"(["y{}"])", i);
}
base_json += "]";
f2_json += "]";
auto base_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), base_json)
.ValueOrDie());
PAIMON_ASSIGN_OR_RAISE(
std::vector<std::shared_ptr<CommitMessage>> base_msgs,
WriteArray(table_path, partition, arrow::schema(fields_)->field_names(), base_array));
// both writes are stamped with the same first row id, so their files cover the same row
// id range and form one row range group
SetFirstRowId(first_row_id, base_msgs);
PAIMON_RETURN_NOT_OK(Commit(table_path, base_msgs));
arrow::FieldVector f2_fields = {fields_[2]};
auto f2_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(f2_fields), f2_json)
.ValueOrDie());
PAIMON_ASSIGN_OR_RAISE(
std::vector<std::shared_ptr<CommitMessage>> f2_msgs,
WriteArray(table_path, partition, {"f2"}, f2_array, partial_write_options));
SetFirstRowId(first_row_id, f2_msgs);
PAIMON_RETURN_NOT_OK(Commit(table_path, f2_msgs));
std::vector<std::shared_ptr<CommitMessage>> group_msgs;
group_msgs.insert(group_msgs.end(), base_msgs.begin(), base_msgs.end());
group_msgs.insert(group_msgs.end(), f2_msgs.begin(), f2_msgs.end());
return group_msgs;
}
/// Flattens the f0 column of a read result, looked up by name so the row kind column the
/// reader prepends does not shift it.
static Result<std::vector<int32_t>> CollectF0Values(
const std::shared_ptr<arrow::ChunkedArray>& rows) {
std::vector<int32_t> values;
if (!rows) {
return values;
}
for (const std::shared_ptr<arrow::Array>& chunk : rows->chunks()) {
auto struct_array = std::dynamic_pointer_cast<arrow::StructArray>(chunk);
if (!struct_array) {
return Status::Invalid("read result chunk is not a struct array");
}
auto f0_array =
std::dynamic_pointer_cast<arrow::Int32Array>(struct_array->GetFieldByName("f0"));
if (!f0_array) {
return Status::Invalid("read result has no int32 f0 column");
}
for (int64_t i = 0; i < f0_array->length(); i++) {
values.push_back(f0_array->Value(i));
}
}
return values;
}
struct LimitScanResult {
/// The splits the limit push down kept in the plan.
std::vector<std::shared_ptr<Split>> splits;
/// The rows reading those splits produced, null when the read returned nothing. The
/// read does not truncate to the limit, so this is everything the kept splits hold, or
/// everything that survives the predicate when `enable_predicate_filter` is set.
std::shared_ptr<arrow::ChunkedArray> rows;
};
/// Scans with a pushed-down row limit and reads the planned splits back. The push down only
/// prunes splits, so a correct plan must still expose at least `limit` rows to the read.
///
/// `enable_predicate_filter` turns the predicate into a row filter while reading. It is off
/// by default, matching the read context default: a predicate then only prunes while
/// planning, and the read returns every row the kept splits hold.
Result<LimitScanResult> ScanAndReadWithLimit(
const std::string& table_path, const std::vector<std::string>& read_schema, int32_t limit,
const std::shared_ptr<Predicate>& predicate = nullptr,
const std::vector<Range>& row_ranges = {}, bool enable_predicate_filter = false) const {
ScanContextBuilder scan_context_builder(table_path);
scan_context_builder.SetLimit(limit);
scan_context_builder.SetPredicate(predicate);
if (!row_ranges.empty()) {
scan_context_builder.SetGlobalIndexResult(
BitmapGlobalIndexResult::FromRanges(row_ranges));
}
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context,
FinishScanContext(scan_context_builder));
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> table_scan,
TableScan::Create(std::move(scan_context)));
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> result_plan, table_scan->CreatePlan());
LimitScanResult result;
result.splits = result_plan->Splits();
ReadContextBuilder read_context_builder(table_path);
read_context_builder.SetReadFieldNames(read_schema)
.SetPredicate(predicate)
.EnablePredicateFilter(enable_predicate_filter);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context,
read_context_builder.Finish());
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> table_read,
TableRead::Create(std::move(read_context)));
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> batch_reader,
table_read->CreateReader(result.splits));
PAIMON_ASSIGN_OR_RAISE(result.rows,
ReadResultCollector::CollectResult(std::move(batch_reader)));
return result;
}
/// Plans the table without any push down and returns the planned splits, so a test can
/// assert what the scan handed the read: which data file each deletion file landed on, and
/// the row count derived from them. Reading the splits back is ScanAndReadWithLimit's job,
/// which verifies the returned rows after their reader has been destroyed.
Result<std::vector<std::shared_ptr<Split>>> PlanSplits(const std::string& table_path) const {
ScanContextBuilder scan_context_builder(table_path);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context,
FinishScanContext(scan_context_builder));
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> table_scan,
TableScan::Create(std::move(scan_context)));
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, table_scan->CreatePlan());
return plan->Splits();
}
/// Anchor file name of every row range group of the table, ordered by ascending first row
/// id, derived from the planned splits.
///
/// Deriving them from the same metas the read groups keeps the two in step. Deriving them
/// from the metas a write hands back uses pre-commit copies, and a vector keyed by a file
/// the read does not consider the anchor is never found, silently leaving rows undeleted.
Result<std::vector<std::string>> PlannedAnchorFileNames(const std::string& table_path) const {
PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::vector<std::shared_ptr<DataFileMeta>> data_files;
for (const std::shared_ptr<Split>& split : splits) {
auto split_impl = std::dynamic_pointer_cast<DataSplitImpl>(split);
if (!split_impl) {
return Status::Invalid("split cannot cast to DataSplitImpl");
}
for (const std::shared_ptr<DataFileMeta>& file : split_impl->DataFiles()) {
data_files.push_back(file);
}
}
return DeletionVectorTestHelper::RetrieveAnchorFileNames(data_files);
}
/// PlannedAnchorFileNames for a table holding a single row range group.
Result<std::string> PlannedAnchorFileName(const std::string& table_path) const {
PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
if (anchor_file_names.size() != 1) {
return Status::Invalid("expected exactly one row range group in the table");
}
return anchor_file_names[0];
}
using DeletionCardinalityMap = std::map<std::string, int64_t>;
/// Maps every data file of `split` that carries a deletion file to that file's cardinality.
static Result<DeletionCardinalityMap> DeletionCardinalityByDataFile(
const std::shared_ptr<Split>& split) {
auto split_impl = std::dynamic_pointer_cast<DataSplitImpl>(split);
if (!split_impl) {
return Status::Invalid("split cannot cast to DataSplitImpl");
}
const std::vector<std::shared_ptr<DataFileMeta>>& data_files = split_impl->DataFiles();
const std::vector<std::optional<DeletionFile>>& deletion_files =
split_impl->DeletionFiles();
if (!deletion_files.empty() && deletion_files.size() != data_files.size()) {
return Status::Invalid("deletion files are not aligned with data files");
}
DeletionCardinalityMap cardinality_by_file;
for (size_t i = 0; i < deletion_files.size(); i++) {
if (deletion_files[i] == std::nullopt) {
continue;
}
if (deletion_files[i].value().cardinality == std::nullopt) {
return Status::Invalid("deletion file is missing its cardinality");
}
cardinality_by_file[data_files[i]->file_name] =
deletion_files[i].value().cardinality.value();
}
return cardinality_by_file;
}
Status CommitWithRowIdCheckFromSnapshot(
const std::string& table_path,
const std::vector<std::shared_ptr<CommitMessage>>& commit_msgs,
std::optional<int64_t> row_id_check_from_snapshot) const {
CommitContextBuilder commit_builder(table_path, "commit_user_1");
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)));
file_store_commit->RowIdCheckConflict(row_id_check_from_snapshot);
return file_store_commit->Commit(commit_msgs);
}
void CheckGlobalIndexScoresWithFiltering(bool merge_files) const {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true,
{{Options::READ_BATCH_SIZE, "4"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
// Row id = 100 + f0; f2 is "x<f0>" before the update and "y<f0>" after it.
// Each candidate's score is row id + 0.5.
constexpr int64_t first_row_id = 100;
auto base_array = PrepareBulkData(
12, [](int32_t i) { return fmt::format(R"({}, "a{}", "x{}")", i, i, i); }, fields_);
ASSERT_OK_AND_ASSIGN(auto base_msgs,
WriteArray(table_path, {"f0", "f1", "f2"}, base_array));
SetFirstRowId(first_row_id, base_msgs);
ASSERT_OK(Commit(table_path, base_msgs));
if (merge_files) {
auto update_array = PrepareBulkData(
12, [](int32_t i) { return fmt::format(R"("y{}")", i); }, {fields_[2]});
ASSERT_OK_AND_ASSIGN(auto update_msgs, WriteArray(table_path, {"f2"}, update_array));
SetFirstRowId(first_row_id, update_msgs);
ASSERT_OK(Commit(table_path, update_msgs));
}
auto not_two = PredicateBuilder::NotEqual(0, "f0", FieldType::INT, Literal(2));
auto below_ten = PredicateBuilder::LessThan(0, "f0", FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> predicate,
PredicateBuilder::And({not_two, below_ten}));
auto reject_all = PredicateBuilder::GreaterOrEqual(0, "f0", FieldType::INT, Literal(100));
const std::vector<int64_t> candidates = {100, 101, 102, 104, 105, 106,
107, 108, 109, 110, 111};
const std::vector<float> scores = {100.5f, 101.5f, 102.5f, 104.5f, 105.5f, 106.5f,
107.5f, 108.5f, 109.5f, 110.5f, 111.5f};
// DV removes global row ids 101, 104, 105, 106, 107 and 111.
const std::vector<int64_t> deleted = {1, 4, 5, 6, 7, 11};
struct FilterCase {
std::string name;
std::shared_ptr<Predicate> predicate;
std::vector<int64_t> expected_ids_without_dv;
std::vector<int64_t> expected_ids_with_dv;
};
const std::vector<FilterCase> filter_cases = {
{"no predicate",
nullptr,
{100, 101, 102, 104, 105, 106, 107, 108, 109, 110, 111},
{100, 102, 108, 109, 110}},
{"f0 != 2 AND f0 < 10",
predicate,
{100, 101, 104, 105, 106, 107, 108, 109},
{100, 108, 109}},
{"f0 >= 100", reject_all, {}, {}},
};
for (bool with_dv : {false, true}) {
if (with_dv) {
ASSERT_OK_AND_ASSIGN(std::string anchor, PlannedAnchorFileName(table_path));
ASSERT_OK(CommitDeletionVectors(table_path, base_msgs[0], {{anchor, deleted}}));
}
ScanContextBuilder scan_builder(table_path);
scan_builder.SetGlobalIndexResult(std::make_shared<BitmapScoredGlobalIndexResult>(
RoaringBitmap64::From(candidates), std::vector<float>(scores)));
ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_builder));
ASSERT_OK_AND_ASSIGN(auto scan, TableScan::Create(std::move(scan_context)));
ASSERT_OK_AND_ASSIGN(auto plan, scan->CreatePlan());
ASSERT_EQ(plan->Splits().size(), 1);
auto scored_split = std::dynamic_pointer_cast<IndexedSplitImpl>(plan->Splits()[0]);
ASSERT_TRUE(scored_split);
ASSERT_EQ(scored_split->RowRanges(),
std::vector<Range>({Range(100, 102), Range(104, 111)}));
ASSERT_EQ(scored_split->Scores().size(), candidates.size());
auto data_split =
std::dynamic_pointer_cast<DataSplitImpl>(scored_split->GetDataSplit());
ASSERT_TRUE(data_split);
ASSERT_EQ(data_split->DataFiles().size(), merge_files ? 2 : 1);
auto unscored_split =
std::make_shared<IndexedSplitImpl>(data_split, scored_split->RowRanges());
for (const auto& filter_case : filter_cases) {
const auto& expected_ids = with_dv ? filter_case.expected_ids_with_dv
: filter_case.expected_ids_without_dv;
for (bool project_row_id : {false, true}) {
std::vector<std::string> read_fields = {"_INDEX_SCORE", "f2", "f0"};
arrow::FieldVector expected_fields = {SpecialFields::ValueKind().ArrowField(),
SpecialFields::IndexScore().ArrowField(),
fields_[2], fields_[0]};
if (project_row_id) {
read_fields.insert(read_fields.begin(), "_ROW_ID");
expected_fields.insert(expected_fields.begin() + 1,
SpecialFields::RowId().ArrowField());
}
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<Predicate> bound_predicate,
PredicateUtils::CreatePickedFieldFilter(filter_case.predicate,
{{"f0", project_row_id ? 3 : 2}}));
ReadContextBuilder read_builder(table_path);
read_builder.SetReadFieldNames(read_fields)
.SetPredicate(bound_predicate)
.EnablePredicateFilter(true);
ASSERT_OK_AND_ASSIGN(auto read_context, read_builder.Finish());
ASSERT_OK_AND_ASSIGN(auto read, TableRead::Create(std::move(read_context)));
// Reuse the same TableRead with and without scores: the internal row-id
// projection must not affect a subsequent split's output schema.
for (bool with_scores : {true, false}) {
SCOPED_TRACE(fmt::format(
"merge={}, dv={}, predicate={}, row_id={}, scores={}", merge_files,
with_dv, filter_case.name, project_row_id, with_scores));
ASSERT_OK_AND_ASSIGN(
auto reader,
read->CreateReader(with_scores ? scored_split : unscored_split));
ASSERT_OK_AND_ASSIGN(auto result,
ReadResultCollector::CollectResult(std::move(reader)));
if (expected_ids.empty()) {
ASSERT_FALSE(result);
continue;
}
ASSERT_TRUE(result);
std::vector<std::string> expected_rows;
for (int64_t row_id : expected_ids) {
int64_t value = row_id - first_row_id;
std::string row_id_json =
project_row_id ? fmt::format("{}, ", row_id) : "";
std::string score_json =
with_scores ? fmt::format("{}", row_id + 0.5f) : "null";
expected_rows.push_back(
fmt::format(R"([0, {}{}, "{}{}", {}])", row_id_json, score_json,
merge_files ? "y" : "x", value, value));
}
std::shared_ptr<arrow::ChunkedArray> expected;
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(
arrow::struct_(expected_fields),
{fmt::format("[{}]", fmt::join(expected_rows, ","))},
&expected)
.ok());
ASSERT_TRUE(expected->Equals(*result))
<< "actual=" << result->ToString()
<< "\nexpected=" << expected->ToString();
}
}
}
}
}
Status ScanAndRead(const std::string& table_path, const std::vector<std::string>& read_schema,
const std::shared_ptr<arrow::StructArray>& expected_array,
const std::shared_ptr<Predicate>& predicate = nullptr,
const std::vector<Range>& row_ranges = {},
bool check_scan_plan_when_empty_result = true,
bool apply_predicate_to_scan = true) const {
// scan
ScanContextBuilder scan_context_builder(table_path);
scan_context_builder.SetPredicate(apply_predicate_to_scan ? predicate : nullptr);
if (!row_ranges.empty()) {
auto global_index_result = BitmapGlobalIndexResult::FromRanges(row_ranges);
scan_context_builder.SetGlobalIndexResult(global_index_result);
}
PAIMON_ASSIGN_OR_RAISE(auto scan_context, FinishScanContext(scan_context_builder));
PAIMON_ASSIGN_OR_RAISE(auto table_scan, TableScan::Create(std::move(scan_context)));
PAIMON_ASSIGN_OR_RAISE(auto result_plan, table_scan->CreatePlan());
if (!expected_array && check_scan_plan_when_empty_result) {
if (!result_plan->Splits().empty()) {
return Status::Invalid("check_scan_plan_when_empty_result but plan is not empty");
}
}
// read
auto splits = result_plan->Splits();
ReadContextBuilder read_context_builder(table_path);
read_context_builder.SetReadFieldNames(read_schema).SetPredicate(predicate);
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();
}
// add row kind array for expected array
auto row_kind_scalar =
std::make_shared<arrow::Int8Scalar>(RowKind::Insert()->ToByteValue());
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
auto row_kind_array,
arrow::MakeArrayFromScalar(*row_kind_scalar, expected_array->length()));
arrow::ArrayVector expected_with_row_kind_fields = expected_array->fields();
std::vector<std::string> expected_with_row_kind_field_names =
arrow::schema(expected_array->type()->fields())->field_names();
expected_with_row_kind_fields.insert(expected_with_row_kind_fields.begin(), row_kind_array);
expected_with_row_kind_field_names.insert(expected_with_row_kind_field_names.begin(),
"_VALUE_KIND");
// check read result
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
auto expected_with_row_kind_array,
arrow::StructArray::Make(expected_with_row_kind_fields,
expected_with_row_kind_field_names));
auto expected_chunk_array =
std::make_shared<arrow::ChunkedArray>(expected_with_row_kind_array);
if (!expected_chunk_array->Equals(read_result)) {
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();
}
void CheckScanResult(const std::string& table_path, const std::shared_ptr<Predicate>& predicate,
const std::vector<Range>& row_ranges,
const std::vector<std::optional<int64_t>>& expected_first_row_ids,
const std::vector<int64_t>& expected_row_counts) {
ASSERT_EQ(expected_first_row_ids.size(), expected_row_counts.size());
ScanContextBuilder scan_context_builder(table_path);
scan_context_builder.SetPredicate(predicate);
if (!row_ranges.empty()) {
auto global_index_result = BitmapGlobalIndexResult::FromRanges(row_ranges);
scan_context_builder.SetGlobalIndexResult(global_index_result);
}
ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder));
ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context)));
ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan());
const auto& result_splits = result_plan->Splits();
if (expected_first_row_ids.empty()) {
ASSERT_EQ(result_splits.size(), 0);
return;
}
ASSERT_EQ(result_splits.size(), 1);
std::shared_ptr<DataSplitImpl> data_split;
if (auto indexed_split = std::dynamic_pointer_cast<IndexedSplit>(result_splits[0])) {
data_split = std::dynamic_pointer_cast<DataSplitImpl>(indexed_split->GetDataSplit());
} else {
data_split = std::dynamic_pointer_cast<DataSplitImpl>(result_splits[0]);
}
ASSERT_TRUE(data_split);
std::vector<std::optional<int64_t>> result_first_row_ids;
std::vector<int64_t> result_row_counts;
for (const auto& meta : data_split->DataFiles()) {
result_first_row_ids.push_back(meta->first_row_id);
result_row_counts.push_back(meta->row_count);
}
ASSERT_EQ(result_first_row_ids, expected_first_row_ids);
ASSERT_EQ(result_row_counts, expected_row_counts);
}
Result<std::unique_ptr<ScanContext>> FinishScanContext(ScanContextBuilder& builder) const {
if (EnableSnapshotLiveManifestCache()) {
if (!snapshot_live_manifest_cache_) {
snapshot_live_manifest_cache_ =
std::make_shared<LruCache>(/*max_weight=*/64 * 1024 * 1024);
}
builder.AddOption(Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS, "3")
.WithCache(snapshot_live_manifest_cache_);
}
return builder.Finish();
}
std::string FileFormat() const {
return std::get<0>(GetParam());
}
bool EnableSnapshotLiveManifestCache() const {
return std::get<1>(GetParam());
}
std::shared_ptr<arrow::StructArray> PrepareBulkData(
int32_t write_batch_size, std::function<std::string(int32_t)> data_generator,
const arrow::FieldVector& fields) const {
std::string data_str = "[";
for (int32_t i = 0; i < write_batch_size; i++) {
data_str.append("[");
auto row_str = data_generator(i);
data_str.append(row_str);
data_str.append("],");
}
data_str.pop_back();
data_str.append("]");
return std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), data_str)
.ValueOrDie());
}
protected:
std::unique_ptr<UniqueTestDirectory> dir_;
mutable std::shared_ptr<Cache> snapshot_live_manifest_cache_;
arrow::FieldVector fields_ = {
arrow::field("f0", arrow::int32()),
arrow::field("f1", arrow::utf8()),
arrow::field("f2", arrow::utf8()),
};
};
class DataEvolutionDvTableTest : public DataEvolutionTableTest {};
TEST_P(DataEvolutionTableTest, TestBasic) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1, f2
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f2
std::vector<std::string> write_cols1 = {"f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(commit_msgs, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], fields_[0], SpecialFields::SequenceNumber().ArrowField(),
SpecialFields::RowId().ArrowField(), fields_[2]}),
R"([
["a", 1, 2, 0, "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1", "f0", "_SEQUENCE_NUMBER", "_ROW_ID", "f2"},
expected_row_tracking_array));
// read score but not indexed split
ASSERT_NOK_WITH_MSG(
ScanAndRead(table_path, {"f0", "f1", "_INDEX_SCORE"}, expected_row_tracking_array,
/*predicate=*/nullptr,
/*row_ranges=*/{}),
"Invalid read schema, read _INDEX_SCORE while split cannot cast to IndexedSplit");
}
TEST_P(DataEvolutionTableTest, TestCommitConflictOnOverlappedRowIdAndWriteColumns) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
// Snapshot 1: initialize row id range [0, 0].
std::vector<std::string> init_write_cols = {"f0", "f1", "f2"};
auto init_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto init_msgs, WriteArray(table_path, init_write_cols, init_array));
ASSERT_OK(Commit(table_path, init_msgs));
// Snapshot 2: update f2 at row id 0.
std::vector<std::string> write_cols = {"f2"};
auto src_array_1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs_1, WriteArray(table_path, write_cols, src_array_1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs_1);
ASSERT_OK(Commit(table_path, commit_msgs_1));
// Snapshot 3 attempt: update f2 at row id 0 again, and check history from snapshot 1.
// This should conflict with snapshot 2 because row-id range and write columns overlap.
auto src_array_2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs_2, WriteArray(table_path, write_cols, src_array_2));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs_2);
ASSERT_NOK_WITH_MSG(
CommitWithRowIdCheckFromSnapshot(table_path, commit_msgs_2,
/*row_id_check_from_snapshot=*/1),
"multiple 'MERGE INTO' operations have encountered conflicts, updating the same file");
}
TEST_P(DataEvolutionTableTest, TestMultipleAppends) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1, f2
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f0, f1
std::vector<std::string> write_cols1 = {"f0", "f1"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[1, "a"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(10, commit_msgs1);
// write field: f2
std::vector<std::string> write_cols2 = {"f2"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(10, commit_msgs2);
std::vector<std::shared_ptr<CommitMessage>> total_msgs;
total_msgs.insert(total_msgs.end(), commit_msgs1.begin(), commit_msgs1.end());
total_msgs.insert(total_msgs.end(), commit_msgs2.begin(), commit_msgs2.end());
ASSERT_OK(Commit(table_path, total_msgs));
// write field: f0, f1
std::vector<std::string> write_cols3 = {"f0", "f1"};
auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[2, "c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs3, WriteArray(table_path, write_cols3, src_array3));
SetFirstRowId(11, commit_msgs3);
ASSERT_OK(Commit(table_path, commit_msgs3));
// write field: f2
std::vector<std::string> write_cols4 = {"f2"};
auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs4, WriteArray(table_path, write_cols4, src_array4));
SetFirstRowId(11, commit_msgs4);
ASSERT_OK(Commit(table_path, commit_msgs4));
{
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[1, "a", "b"],
[2, "c", "d"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
}
{
std::vector<Range> row_ranges = {Range(0l, 0l), Range(11l, 11l)};
// test with row ids
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 11, 11}, /*expected_row_counts=*/{10, 1, 1});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"],
[2, "c", "d"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({
fields_[0],
fields_[1],
fields_[2],
SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField(),
}),
R"([
[1, "a", "b", 0, 1],
[1, "a", "b", 1, 1],
[1, "a", "b", 2, 1],
[1, "a", "b", 3, 1],
[1, "a", "b", 4, 1],
[1, "a", "b", 5, 1],
[1, "a", "b", 6, 1],
[1, "a", "b", 7, 1],
[1, "a", "b", 8, 1],
[1, "a", "b", 9, 1],
[1, "a", "b", 10, 2],
[2, "c", "d", 11, 4]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestOnlySomeColumns) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0
std::vector<std::string> write_cols0 = {"f0"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0]}), R"([
[1]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f1
std::vector<std::string> write_cols1 = {"f1"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[1]}), R"([
["a"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(0, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
// write field: f2
std::vector<std::string> write_cols2 = {"f2"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(0, commit_msgs2);
ASSERT_OK(Commit(table_path, commit_msgs2));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({
fields_[0],
fields_[1],
fields_[2],
SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField(),
}),
R"([
[1, "a", "b", 0, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestMultipleSharedShreddingMapsPartialOverwrite) {
if (FileFormat() == "mosaic" || FileFormat() == "lance") {
return;
}
if (FileFormat() == "avro") {
return;
}
auto map_type = arrow::map(arrow::utf8(), arrow::int64());
arrow::FieldVector fields = {
arrow::field("id", arrow::int32()),
arrow::field("map1", map_type),
arrow::field("map2", map_type),
};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{"fields.map1.map.storage-layout", "shared-shredding"},
{"fields.map1.map.shared-shredding.max-columns", "1"},
{"fields.map2.map.storage-layout", "shared-shredding"},
{"fields.map2.map.shared-shredding.max-columns", "1"},
};
CreateTable(fields, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields);
std::vector<std::string> write_cols0 = {"id", "map1"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[0], fields[1]}), R"([
[1, [["a", 10], ["b", 20]]],
[11, [["a", 11], ["b", 21]]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs0));
std::vector<std::string> write_cols1 = {"id", "map2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[0], fields[2]}), R"([
[2, [["c", 30], ["d", 40]]],
[12, [["c", 31], ["d", 41]]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
std::vector<std::string> write_cols2 = {"map1"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[[["b", 200], ["a", 100]]],
[[["b", 201], ["a", 101]]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs2);
ASSERT_OK(Commit(table_path, commit_msgs2));
// Read all columns and merge values from all partial files.
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[2, [["a", 100], ["b", 200]], [["c", 30], ["d", 40]]],
[12, [["a", 101], ["b", 201]], [["c", 31], ["d", 41]]]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// Read a subset of columns and recall only the requested shared-shredding MAP column.
auto expected_column_pruned_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[0], fields[2]}), R"([
[2, [["c", 30], ["d", 40]]],
[12, [["c", 31], ["d", 41]]]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "map2"}, expected_column_pruned_array));
// Read selected keys from both shared-shredding MAP columns after partial overwrite merge.
{
auto map1_selected_keys =
arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"b"});
auto map2_selected_keys =
arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"d"});
auto read_schema = arrow::schema({
fields[0],
fields[1]->WithMetadata(map1_selected_keys),
fields[2]->WithMetadata(map2_selected_keys),
});
auto c_schema = std::make_unique<ArrowSchema>();
ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok());
ScanContextBuilder scan_context_builder(table_path);
ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder));
ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context)));
ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan());
ReadContextBuilder read_context_builder(table_path);
read_context_builder.SetReadSchema(std::move(c_schema));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context,
read_context_builder.Finish());
ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context)));
ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits()));
ASSERT_OK_AND_ASSIGN(auto actual,
ReadResultCollector::CollectResult(std::move(batch_reader)));
auto expected_type = arrow::struct_({
SpecialFields::ValueKind().ArrowField(),
fields[0],
fields[1],
fields[2],
});
auto expected = arrow::ipc::internal::json::ArrayFromJSON(expected_type, R"([
[0, 2, [["b", 200]], [["d", 40]]],
[0, 12, [["b", 201]], [["d", 41]]]
])")
.ValueOrDie();
auto expected_chunked = std::make_shared<arrow::ChunkedArray>(expected);
ASSERT_TRUE(expected_chunked->Equals(actual))
<< "actual=" << actual->ToString() << "\nexpected=" << expected_chunked->ToString();
}
// Read a subset of rows after merging values from all partial files.
auto expected_partial_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[12, [["a", 101], ["b", 201]], [["c", 31], ["d", 41]]]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_partial_array,
/*predicate=*/nullptr,
/*row_ranges=*/{Range(1l, 1l)}));
// Read row tracking fields and verify the latest partial overwrite sequence number.
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields[0], fields[1], fields[2], SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
[2, [["a", 100], ["b", 200]], [["c", 30], ["d", 40]], 0, 3],
[12, [["a", 101], ["b", 201]], [["c", 31], ["d", 41]], 1, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "map1", "map2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestNullValues) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1
std::vector<std::string> write_cols1 = {"f0", "f1"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[1, null]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(0, commit_msgs1);
// write field: f2
std::vector<std::string> write_cols2 = {"f2"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(0, commit_msgs2);
std::vector<std::shared_ptr<CommitMessage>> total_msgs;
total_msgs.insert(total_msgs.end(), commit_msgs1.begin(), commit_msgs1.end());
total_msgs.insert(total_msgs.end(), commit_msgs2.begin(), commit_msgs2.end());
ASSERT_OK(Commit(table_path, total_msgs));
// Commit 2: Overwrite with non-null
// write field: f2
std::vector<std::string> write_cols3 = {"f2"};
auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs3, WriteArray(table_path, write_cols3, src_array3));
SetFirstRowId(0, commit_msgs3);
ASSERT_OK(Commit(table_path, commit_msgs3));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, null, "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({
fields_[0],
fields_[1],
fields_[2],
SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField(),
}),
R"([
[1, null, "c", 0, 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestMultipleAppendsDifferentFirstRowIds) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// First commit, firstRowId = 0
// write field: f0, f1
std::vector<std::string> write_cols1 = {"f0", "f1"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[1, "a"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(0, commit_msgs1);
// write field: f2
std::vector<std::string> write_cols2 = {"f2"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(0, commit_msgs2);
std::vector<std::shared_ptr<CommitMessage>> total_msgs;
total_msgs.insert(total_msgs.end(), commit_msgs1.begin(), commit_msgs1.end());
total_msgs.insert(total_msgs.end(), commit_msgs2.begin(), commit_msgs2.end());
ASSERT_OK(Commit(table_path, total_msgs));
// Second commit, firstRowId = 1
// write field: f0, f1
std::vector<std::string> write_cols3 = {"f0", "f1"};
auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[2, "c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs3, WriteArray(table_path, write_cols3, src_array3));
SetFirstRowId(1, commit_msgs3);
ASSERT_OK(Commit(table_path, commit_msgs3));
// Third commit
// write field: f2
std::vector<std::string> write_cols4 = {"f2"};
auto src_array4 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs4, WriteArray(table_path, write_cols4, src_array4));
SetFirstRowId(1, commit_msgs4);
ASSERT_OK(Commit(table_path, commit_msgs4));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"],
[2, "c", "d"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({
fields_[0],
fields_[1],
fields_[2],
SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField(),
}),
R"([
[1, "a", "b", 0, 1],
[2, "c", "d", 1, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestMoreData) {
CreateTable();
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1
std::vector<std::string> write_cols1 = {"f0", "f1"};
// row0: 0, a0; row1: 1, a1 ...
auto src_array1 = PrepareBulkData(
10000, [](int32_t i) { return std::to_string(i) + ", \"a" + std::to_string(i) + "\""; },
{fields_[0], fields_[1]});
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(0, commit_msgs1);
// write field: f2
std::vector<std::string> write_cols2 = {"f2"};
// row0: b0; row1: b1 ...
auto src_array2 = PrepareBulkData(
10000, [](int32_t i) { return "\"b" + std::to_string(i) + "\""; }, {fields_[2]});
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols2, src_array2));
SetFirstRowId(0, commit_msgs2);
std::vector<std::shared_ptr<CommitMessage>> total_msgs;
total_msgs.insert(total_msgs.end(), commit_msgs1.begin(), commit_msgs1.end());
total_msgs.insert(total_msgs.end(), commit_msgs2.begin(), commit_msgs2.end());
ASSERT_OK(Commit(table_path, total_msgs));
// write field: f2
std::vector<std::string> write_cols3 = {"f2"};
// row0: c0; row1: c1 ...
auto src_array3 = PrepareBulkData(
10000, [](int32_t i) { return "\"c" + std::to_string(i) + "\""; }, {fields_[2]});
ASSERT_OK_AND_ASSIGN(auto commit_msgs3, WriteArray(table_path, write_cols3, src_array3));
ASSERT_OK(Commit(table_path, commit_msgs3));
// row0: 0, a0, c0; row1: 1, a1, c1 ...
auto expected_array = PrepareBulkData(10000,
[](int32_t i) {
return std::to_string(i) + ", \"a" +
std::to_string(i) + "\", \"c" +
std::to_string(i) + "\"";
},
{fields_[0], fields_[1], fields_[2]});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
}
TEST_P(DataEvolutionTableTest, TestOnlyRowTrackingEnabled) {
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "false"},
};
CreateTable(/*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1, f2
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"],
[2, "c", "d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], fields_[0], SpecialFields::SequenceNumber().ArrowField(),
SpecialFields::RowId().ArrowField(), fields_[2]}),
R"([
["a", 1, 1, 0, "b"],
["c", 2, 1, 1, "d"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1", "f0", "_SEQUENCE_NUMBER", "_ROW_ID", "f2"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestExternalPath) {
// create external path dir
auto external_dir = UniqueTestDirectory::Create("local");
ASSERT_TRUE(external_dir);
std::string external_test_dir = external_dir->Str();
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::DATA_FILE_EXTERNAL_PATHS, "FILE://" + external_test_dir},
};
CreateTable(/*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1
std::vector<std::string> write_cols0 = {"f0", "f1"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[1, "a"],
[2, "c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs0));
// write field: f0, f2
std::vector<std::string> write_cols1 = {"f0", "f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[2]}), R"([
[10, "b"],
[20, "d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[10, "a", "b"],
[20, "c", "d"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], fields_[0], fields_[2], SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
["a", 10, "b", 0, 2],
["c", 20, "d", 1, 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1", "f0", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestWithPartitionSimple) {
std::vector<std::string> partition_keys = {"f1"};
CreateTable(partition_keys);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1, f2 for f1=2024
std::map<std::string, std::string> partition0 = {{"f1", "2024"}};
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "2024", "b1"],
[2, "2024", "b2"],
[3, "2024", "b3"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, partition0, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f2 for f1=2024
std::vector<std::string> write_cols1 = {"f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c1"],
["c2"],
["c3"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(commit_msgs, WriteArray(table_path, partition0, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
// write fields: f1, f2 for f1=2025
std::map<std::string, std::string> partition2 = {{"f1", "2025"}};
std::vector<std::string> write_cols2 = {"f1", "f2"};
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[1], fields_[2]}), R"([
["2025", "d1"],
["2025", "d2"],
["2025", "d3"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(commit_msgs, WriteArray(table_path, partition2, write_cols2, src_array2));
SetFirstRowId(/*reset_first_row_id=*/3, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
// test read all fields
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "2024", "c1"],
[2, "2024", "c2"],
[3, "2024", "c3"],
[null, "2025", "d1"],
[null, "2025", "d2"],
[null, "2025", "d3"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// test only read partition fields
auto expected_array_only_partition = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[1]}), R"([
["2024"],
["2024"],
["2024"],
["2025"],
["2025"],
["2025"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1"}, expected_array_only_partition));
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[0], fields_[1], fields_[2], SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
[1, "2024", "c1", 0, 2],
[2, "2024", "c2", 1, 2],
[3, "2024", "c3", 2, 2],
[null, "2025", "d1", 3, 3],
[null, "2025", "d2", 4, 3],
[null, "2025", "d3", 5, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
// read only read partition fields and row tracking
auto expected_partition_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
["2024", 0, 2],
["2024", 1, 2],
["2024", 2, 2],
["2025", 3, 3],
["2025", 4, 3],
["2025", 5, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_partition_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestWithPartitionWithoutPartitionFieldsInFile) {
std::vector<std::string> partition_keys = {"f1"};
CreateTable(partition_keys);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f2 for f1=2024
std::map<std::string, std::string> partition0 = {{"f1", "2024"}};
std::vector<std::string> write_cols0 = {"f0", "f2"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[2]}), R"([
[1, "b1"],
[2, "b2"],
[3, "b3"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, partition0, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f2 for f1=2024
std::vector<std::string> write_cols1 = {"f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c1"],
["c2"],
["c3"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(commit_msgs, WriteArray(table_path, partition0, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
{
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "2024", "c1"],
[2, "2024", "c2"],
[3, "2024", "c3"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
}
{
// test with row ids
std::vector<Range> row_ranges = {Range(1l, 1l)};
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{3, 3});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[2, "2024", "c2"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
// read with row tracking
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[0], fields_[1], fields_[2], SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
[1, "2024", "c1", 0, 2],
[2, "2024", "c2", 1, 2],
[3, "2024", "c3", 2, 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array));
}
TEST_P(DataEvolutionTableTest, TestPartitionWithPredicate) {
auto file_format = FileFormat();
if (file_format == "avro" || file_format == "lance") {
return;
}
std::vector<std::string> partition_keys = {"f1"};
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{"parquet.write.max-row-group-length", "1"}};
if (file_format == "mosaic") {
options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0");
}
CreateTable(partition_keys, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1 for partition f1 = "2024"
std::vector<std::string> write_cols0 = {"f0", "f1"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[11, "2024"],
[12, "2024"],
[13, "2024"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0,
WriteArray(table_path, {{"f1", "2024"}}, write_cols0, src_array0));
// write field: f2 for partition f1 = "2024"
std::vector<std::string> write_cols1 = {"f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["a"],
["b"],
["c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, {{"f1", "2024"}}, write_cols1, src_array1));
std::vector<std::shared_ptr<CommitMessage>> total_msgs;
total_msgs.insert(total_msgs.end(), commit_msgs0.begin(), commit_msgs0.end());
total_msgs.insert(total_msgs.end(), commit_msgs1.begin(), commit_msgs1.end());
SetFirstRowId(0, total_msgs);
ASSERT_OK(Commit(table_path, total_msgs));
// write field: f0, f1 for partition f1 = "2025"
std::vector<std::string> write_cols3 = {"f0", "f1"};
auto src_array3 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[21, "2025"],
[22, "2025"],
[23, "2025"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs3,
WriteArray(table_path, {{"f1", "2025"}}, write_cols3, src_array3));
SetFirstRowId(3, commit_msgs3);
ASSERT_OK(Commit(table_path, commit_msgs3));
{
// only set data field predicate, predicate only takes effective in file skip
// read will not push down
auto equal = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT,
Literal(11));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[11, "2024", "a"],
[12, "2024", "b"],
[13, "2024", "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array, equal));
}
{
// only set partition predicate
auto equal =
PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "2024", 4));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[11, "2024", "a"],
[12, "2024", "b"],
[13, "2024", "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array, equal));
}
{
// set partition predicate and data field predicate
auto equal =
PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "2024", 4));
auto greater_than = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(100));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({equal, greater_than}));
ASSERT_OK(
ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr, predicate));
}
{
// read with row tracking
auto equal =
PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "2024", 4));
auto expected_row_tracking_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[0], fields_[1], fields_[2],
SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()}),
R"([
[11, "2024", "a", 0, 1],
[12, "2024", "b", 1, 1],
[13, "2024", "c", 2, 1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID", "_SEQUENCE_NUMBER"},
expected_row_tracking_array, equal));
}
{
// test with row ids
std::vector<Range> row_ranges = {Range(1l, 1l)};
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{3, 3});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[12, "2024", "b"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
{
std::vector<Range> row_ranges = {Range(1l, 1l)};
// test with row ids and partition predicate
auto equal =
PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "2025", 4));
CheckScanResult(table_path, /*predicate=*/equal, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{}, /*expected_row_counts=*/{});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr,
/*predicate=*/equal,
/*row_ranges=*/row_ranges));
}
{
std::vector<Range> row_ranges = {Range(1l, 1l)};
// test with row ids and non-partition predicate
auto equal = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT,
Literal(11));
CheckScanResult(table_path, /*predicate=*/equal, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{3, 3});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[12, "2024", "b"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/equal,
/*row_ranges=*/row_ranges));
}
{
std::vector<Range> row_ranges = {Range(1l, 1l)};
// test with row ids and data predicate
auto equal = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT,
Literal(50));
CheckScanResult(table_path, /*predicate=*/equal, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{}, /*expected_row_counts=*/{});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr,
/*predicate=*/equal,
/*row_ranges=*/row_ranges));
}
}
#ifdef PAIMON_ENABLE_VORTEX
TEST_P(DataEvolutionTableTest, TestReadJavaDedicatedVectorStore) {
if (FileFormat() != "parquet") {
GTEST_SKIP() << "The Java fixture uses parquet main files";
}
const std::string fixture_path = PathUtil::JoinPath(
GetDataDir(), "vortex/dedicated_vector_java_compat.db/dedicated_vector_java_compat");
const std::string table_path = PathUtil::JoinPath(dir_->Str(), "java_dedicated_vector");
// Preserve the Java-produced files: this test only scans and reads the copied table.
ASSERT_TRUE(TestUtil::CopyDirectory(fixture_path, table_path));
// Java DataTypes.FLOAT() serializes the VECTOR element as nullable in schema-0.
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/true), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())};
auto expected = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, [1, 2, 3], "a"], [2, null, "b"], [3, [-1, 0, 1], "c"],
[4, null, "d"], [5, [7, 8, 9], "e"], [6, [10, 11, 12], "f"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, expected));
auto normal_projection =
arrow::StructArray::Make({expected->field(0), expected->field(2)}, {fields[0], fields[2]})
.ValueOrDie();
ASSERT_OK(ScanAndRead(table_path, {"id", "tag"}, normal_projection));
auto vector_projection =
arrow::StructArray::Make({expected->field(1)}, {fields[1]}).ValueOrDie();
ASSERT_OK(ScanAndRead(table_path, {"embedding"}, vector_projection));
// Row ids 1-4 cross the two Java commits and include a null vector from each commit.
auto selected_rows = checked_pointer_cast<arrow::StructArray>(expected->Slice(1, 4));
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, selected_rows,
/*predicate=*/nullptr, /*row_ranges=*/{Range(1, 4)}));
auto selected_vectors =
checked_pointer_cast<arrow::StructArray>(vector_projection->Slice(1, 4));
ASSERT_OK(ScanAndRead(table_path, {"embedding"}, selected_vectors,
/*predicate=*/nullptr, /*row_ranges=*/{Range(1, 4)}));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::map<int64_t, std::vector<std::shared_ptr<DataFileMeta>>> row_groups;
for (const auto& split : splits) {
auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
ASSERT_NE(data_split, nullptr);
for (const auto& file : data_split->DataFiles()) {
ASSERT_TRUE(file->first_row_id.has_value());
row_groups[file->first_row_id.value()].push_back(file);
}
}
ASSERT_EQ(row_groups.size(), 2);
ASSERT_EQ(row_groups.count(0), 1);
ASSERT_EQ(row_groups.count(3), 1);
for (const auto& [first_row_id, files] : row_groups) {
SCOPED_TRACE(first_row_id);
ASSERT_EQ(files.size(), 2);
int32_t vector_file_count = 0;
for (const auto& file : files) {
ASSERT_EQ(file->row_count, 3);
if (VectorStoreUtils::IsVectorStoreFile(file->file_name)) {
++vector_file_count;
ASSERT_TRUE(StringUtils::EndsWith(file->file_name, ".vector.vortex"));
ASSERT_EQ(file->write_cols, std::vector<std::string>({"embedding"}));
} else {
ASSERT_TRUE(StringUtils::EndsWith(file->file_name, ".parquet"));
ASSERT_EQ(file->write_cols, std::vector<std::string>({"id", "tag"}));
}
}
ASSERT_EQ(vector_file_count, 1);
}
}
TEST_P(DataEvolutionTableTest, TestReadJavaIndependentlyRolledVectorStore) {
if (FileFormat() != "parquet") {
GTEST_SKIP() << "The Java fixture uses parquet main files";
}
const std::string fixture_path = PathUtil::JoinPath(
GetDataDir(), "vortex/dedicated_vector_java_rolling.db/dedicated_vector_java_rolling");
const std::string table_path = PathUtil::JoinPath(dir_->Str(), "java_rolled_vector");
ASSERT_TRUE(TestUtil::CopyDirectory(fixture_path, table_path));
constexpr int32_t kRowCount = 3001;
auto vector_type = arrow::fixed_size_list(arrow::float32(), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())};
const std::vector<int32_t> null_vector_ids = {0, 999, 1000, 1999, 2000, 2999, 3000};
auto expected = PrepareBulkData(
kRowCount,
[&null_vector_ids](int32_t id) {
std::string vector_json = std::find(null_vector_ids.begin(), null_vector_ids.end(),
id) != null_vector_ids.end()
? "null"
: fmt::format("[{}, {}, {}]", id, id + 1, id + 2);
return fmt::format(R"({}, {}, "row-{}")", id, vector_json, id);
},
fields);
// The Java fixture sets read.batch-size=127 so batch and vector file boundaries differ.
auto check_projections = [&](const std::shared_ptr<arrow::StructArray>& rows,
const std::vector<Range>& row_ranges) -> Status {
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
std::shared_ptr<arrow::StructArray> normal_projection,
arrow::StructArray::Make({rows->field(0), rows->field(2)}, {fields[0], fields[2]}));
PAIMON_RETURN_NOT_OK(ScanAndRead(table_path, {"id", "tag"}, normal_projection,
/*predicate=*/nullptr, row_ranges));
PAIMON_RETURN_NOT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, rows,
/*predicate=*/nullptr, row_ranges));
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> vector_projection,
arrow::StructArray::Make({rows->field(1)}, {fields[1]}));
return ScanAndRead(table_path, {"embedding"}, vector_projection,
/*predicate=*/nullptr, row_ranges);
};
ASSERT_OK(check_projections(expected, {}));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::shared_ptr<DataFileMeta> main_file;
std::map<int64_t, std::shared_ptr<DataFileMeta>> vector_files;
for (const auto& split : splits) {
auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
ASSERT_NE(data_split, nullptr);
for (const auto& file : data_split->DataFiles()) {
ASSERT_TRUE(file->first_row_id.has_value());
if (VectorStoreUtils::IsVectorStoreFile(file->file_name)) {
ASSERT_TRUE(StringUtils::EndsWith(file->file_name, ".vector.vortex"));
ASSERT_EQ(file->write_cols, std::vector<std::string>({"embedding"}));
ASSERT_TRUE(vector_files.emplace(file->first_row_id.value(), file).second);
} else {
ASSERT_EQ(main_file, nullptr);
ASSERT_TRUE(StringUtils::EndsWith(file->file_name, ".parquet"));
ASSERT_EQ(file->write_cols, std::vector<std::string>({"id", "tag"}));
main_file = file;
}
}
}
ASSERT_NE(main_file, nullptr);
ASSERT_EQ(main_file->first_row_id, 0);
ASSERT_EQ(main_file->row_count, kRowCount);
ASSERT_GT(vector_files.size(), 1);
int64_t next_row_id = 0;
std::vector<Range> selected_ranges = {Range(0, 0)};
for (const auto& [first_row_id, file] : vector_files) {
ASSERT_EQ(first_row_id, next_row_id);
ASSERT_GT(file->row_count, 0);
next_row_id += file->row_count;
if (first_row_id == 0) {
continue;
}
SCOPED_TRACE(first_row_id);
const int64_t from = std::max<int64_t>(0, first_row_id - 2);
const int64_t to = std::min<int64_t>(kRowCount - 1, first_row_id + 2);
auto boundary_rows =
checked_pointer_cast<arrow::StructArray>(expected->Slice(from, to - from + 1));
ASSERT_OK(check_projections(boundary_rows, {Range(from, to)}));
selected_ranges.emplace_back(std::max<int64_t>(0, first_row_id - 1),
std::min<int64_t>(kRowCount - 1, first_row_id + 1));
}
ASSERT_EQ(next_row_id, kRowCount);
// Disjoint selections retain row order while skipping most rows in each physical file.
selected_ranges = Range::SortAndMergeOverlap(selected_ranges, /*adjacent=*/true);
ASSERT_GT(selected_ranges.size(), 1);
arrow::ArrayVector selected_chunks;
for (const auto& range : selected_ranges) {
selected_chunks.push_back(expected->Slice(range.from, range.Count()));
}
auto selected_rows =
checked_pointer_cast<arrow::StructArray>(arrow::Concatenate(selected_chunks).ValueOrDie());
ASSERT_OK(check_projections(selected_rows, selected_ranges));
// Omit the outer VECTOR files in an in-memory split to exercise NULL fill for uncovered
// rows. The main file has no physical embedding column, and the Java fixture stays intact.
ASSERT_EQ(splits.size(), 1);
auto original_split = std::dynamic_pointer_cast<DataSplitImpl>(splits[0]);
ASSERT_NE(original_split, nullptr);
auto middle_file = std::next(vector_files.begin())->second;
const int64_t middle_from = middle_file->first_row_id.value();
const int64_t middle_to = middle_from + middle_file->row_count - 1;
DataSplitImpl::Builder gap_builder(
original_split->Partition(), original_split->Bucket(), original_split->BucketPath(),
std::vector<std::shared_ptr<DataFileMeta>>{main_file, middle_file});
ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> gap_split,
gap_builder.WithTotalBuckets(original_split->TotalBuckets())
.WithSnapshot(original_split->SnapshotId())
.RawConvertible(false)
.Build());
auto gap_expected = PrepareBulkData(
kRowCount,
[&](int32_t id) {
const bool missing_vector = id < middle_from || id > middle_to;
const bool null_vector = std::find(null_vector_ids.begin(), null_vector_ids.end(),
id) != null_vector_ids.end();
std::string vector_json = missing_vector || null_vector
? "null"
: fmt::format("[{}, {}, {}]", id, id + 1, id + 2);
return fmt::format("0, {}", vector_json);
},
{SpecialFields::ValueKind().ArrowField(), fields[1]});
arrow::ArrayVector selected_gap_chunks;
for (const auto& range : selected_ranges) {
selected_gap_chunks.push_back(gap_expected->Slice(range.from, range.Count()));
}
auto selected_gap_expected = checked_pointer_cast<arrow::StructArray>(
arrow::Concatenate(selected_gap_chunks).ValueOrDie());
ReadContextBuilder read_builder(table_path);
read_builder.SetReadFieldNames({"embedding"});
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, read_builder.Finish());
ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read,
TableRead::Create(std::move(read_context)));
const std::vector<std::pair<std::shared_ptr<Split>, std::shared_ptr<arrow::StructArray>>>
gap_cases = {{gap_split, gap_expected},
{std::make_shared<IndexedSplitImpl>(gap_split, selected_ranges),
selected_gap_expected}};
for (const auto& [split, expected_rows] : gap_cases) {
ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader, read->CreateReader(split));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> result,
ReadResultCollector::CollectResult(std::move(reader)));
ASSERT_NE(result, nullptr);
ASSERT_TRUE(result->Equals(arrow::ChunkedArray(expected_rows)))
<< "actual=" << result->ToString() << "\nexpected=" << expected_rows->ToString();
}
// Like Java, the normal file supplies sequence numbers even when VECTOR files are newer.
auto sequence_main = std::make_shared<DataFileMeta>(*main_file);
sequence_main->AssignSequenceNumber(3, 3);
std::vector<std::shared_ptr<DataFileMeta>> sequence_files = {sequence_main};
int64_t sequence = 2;
for (const auto& [first, file] : vector_files) {
auto sequence_vector = std::make_shared<DataFileMeta>(*file);
sequence_vector->AssignSequenceNumber(sequence, sequence);
sequence_files.push_back(sequence_vector);
++sequence;
}
auto sequence_tail = sequence_files.back();
DataSplitImpl::Builder sequence_builder(original_split->Partition(), original_split->Bucket(),
original_split->BucketPath(),
std::move(sequence_files));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> sequence_split,
sequence_builder.RawConvertible(false).Build());
auto tracking = PrepareBulkData(
kRowCount, [](int32_t id) { return fmt::format("0, {}, 3", id); },
{SpecialFields::ValueKind().ArrowField(), SpecialFields::RowId().ArrowField(),
SpecialFields::SequenceNumber().ArrowField()});
for (bool sequential : {true, false}) {
SCOPED_TRACE(sequential);
std::shared_ptr<Split> split = sequence_split;
std::shared_ptr<arrow::Array> vectors = expected->field(1);
if (!sequential) {
// A newer VECTOR covering only the tail forces the range-reader fallback.
const auto& tail = sequence_tail;
ASSERT_GT(tail->max_sequence_number, sequence_main->max_sequence_number);
DataSplitImpl::Builder tail_builder(
original_split->Partition(), original_split->Bucket(), original_split->BucketPath(),
std::vector<std::shared_ptr<DataFileMeta>>{sequence_main, tail});
ASSERT_OK_AND_ASSIGN(split, tail_builder.RawConvertible(false).Build());
const int64_t first = tail->first_row_id.value();
auto missing = arrow::MakeArrayOfNull(vector_type, first).ValueOrDie();
vectors = arrow::Concatenate({missing, vectors->Slice(first)}).ValueOrDie();
}
auto sequence_expected =
arrow::StructArray::Make(
{tracking->field(0), expected->field(0), vectors, tracking->field(1),
tracking->field(2)},
{"_VALUE_KIND", "id", "embedding", "_ROW_ID", "_SEQUENCE_NUMBER"})
.ValueOrDie();
for (const std::vector<std::string>& projection :
{std::vector<std::string>{"id", "embedding", "_ROW_ID", "_SEQUENCE_NUMBER"},
std::vector<std::string>{"id", "_ROW_ID", "_SEQUENCE_NUMBER"},
std::vector<std::string>{"_ROW_ID", "_SEQUENCE_NUMBER"}}) {
ReadContextBuilder sequence_read_builder(table_path);
sequence_read_builder.SetReadFieldNames(projection);
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> sequence_context,
sequence_read_builder.Finish());
ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> sequence_read,
TableRead::Create(std::move(sequence_context)));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> sequence_reader,
sequence_read->CreateReader(split));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> sequence_result,
ReadResultCollector::CollectResult(std::move(sequence_reader)));
ASSERT_NE(sequence_result, nullptr);
arrow::ArrayVector projected_arrays = {tracking->field(0)};
std::vector<std::string> projected_names = {"_VALUE_KIND"};
for (const auto& name : projection) {
projected_arrays.push_back(sequence_expected->GetFieldByName(name));
projected_names.push_back(name);
}
auto projected_expected =
arrow::StructArray::Make(projected_arrays, projected_names).ValueOrDie();
ASSERT_TRUE(sequence_result->Equals(arrow::ChunkedArray(projected_expected)))
<< sequence_result->ToString();
}
}
}
#endif
class DedicatedVectorTableTest : public DataEvolutionTableTest {};
TEST_P(DedicatedVectorTableTest, TestDedicatedVectorStore) {
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, "parquet"}, {Options::VECTOR_FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::TARGET_FILE_ROW_NUM, "2"}, {Options::READ_BATCH_SIZE, "2"}};
CreateTable(fields, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto input = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[0, [0, 0, 0], "skip"],
[1, [1, 2, 3], "a"], [2, null, "b"],
[3, null, "c"], [4, null, "d"],
[5, [7, 8, 9], "e"]
])")
.ValueOrDie());
// Sliced input also exercises a pair containing only null vectors.
std::vector<std::shared_ptr<arrow::Array>> batches = {input->Slice(1, 2), input->Slice(3, 2),
input->Slice(5, 1)};
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> messages,
WriteArrays(table_path, {}, {}, batches));
ASSERT_EQ(messages.size(), 1);
auto message = std::dynamic_pointer_cast<CommitMessageImpl>(messages[0]);
ASSERT_NE(message, nullptr);
const auto& files = message->data_increment_.new_files_;
ASSERT_EQ(files.size(), 6);
for (size_t i = 0; i < files.size(); i += 2) {
ASSERT_FALSE(VectorStoreUtils::IsVectorStoreFile(files[i]->file_name));
ASSERT_TRUE(StringUtils::EndsWith(files[i + 1]->file_name,
fmt::format(".vector.{}", FileFormat())));
ASSERT_EQ(files[i]->row_count, files[i + 1]->row_count);
ASSERT_EQ(files[i]->write_cols, std::vector<std::string>({"id", "tag"}));
ASSERT_EQ(files[i + 1]->write_cols, std::vector<std::string>({"embedding"}));
}
ASSERT_OK(Commit(table_path, messages));
auto expected = checked_pointer_cast<arrow::StructArray>(input->Slice(1));
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, expected));
auto normal_projection =
arrow::StructArray::Make({expected->field(0), expected->field(2)}, {fields[0], fields[2]})
.ValueOrDie();
ASSERT_OK(ScanAndRead(table_path, {"id", "tag"}, normal_projection));
auto vector_projection = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[null], [null], [null], [[7, 8, 9]]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"embedding"}, vector_projection, /*predicate=*/nullptr,
/*row_ranges=*/{Range(1, 4)}));
ASSERT_OK_AND_ASSIGN(messages, WriteArray(table_path, {}, input->Slice(1, 2)));
ASSERT_OK(Commit(table_path, messages));
auto appended_expected = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, [1, 2, 3], "a"], [2, null, "b"],
[3, null, "c"], [4, null, "d"], [5, [7, 8, 9], "e"],
[1, [1, 2, 3], "a"], [2, null, "b"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, appended_expected));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::map<int64_t, std::vector<std::shared_ptr<DataFileMeta>>> row_groups;
for (const auto& split : splits) {
auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
ASSERT_NE(data_split, nullptr);
for (const auto& file : data_split->DataFiles()) {
ASSERT_TRUE(file->first_row_id.has_value());
row_groups[file->first_row_id.value()].push_back(file);
}
}
ASSERT_EQ(row_groups.size(), 4);
ASSERT_EQ(row_groups.count(0), 1);
ASSERT_EQ(row_groups.count(2), 1);
ASSERT_EQ(row_groups.count(4), 1);
ASSERT_EQ(row_groups.count(5), 1);
for (const auto& [first_row_id, group] : row_groups) {
ASSERT_EQ(group.size(), 2);
ASSERT_EQ(group[0]->row_count, group[1]->row_count);
}
WriteContextBuilder partial_builder(table_path, "partial");
partial_builder.WithWriteSchema({"embedding"});
ASSERT_OK_AND_ASSIGN(std::unique_ptr<WriteContext> partial_context, partial_builder.Finish());
ASSERT_NOK_WITH_MSG(FileStoreWrite::Create(std::move(partial_context)).status(),
"partial writes");
WriteContextBuilder invalid_builder(table_path, "invalid");
invalid_builder.AddOption(Options::DATA_EVOLUTION_ENABLED, "false");
ASSERT_OK_AND_ASSIGN(std::unique_ptr<WriteContext> invalid_context, invalid_builder.Finish());
ASSERT_NOK_WITH_MSG(FileStoreWrite::Create(std::move(invalid_context)).status(),
"requires row-tracking.enabled and data-evolution.enabled");
}
TEST_P(DedicatedVectorTableTest, TestIndependentVectorRolling) {
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())};
const std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, "parquet"},
{Options::VECTOR_FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::WRITE_ONLY, "true"},
{Options::TARGET_FILE_SIZE, "128 mb"},
{Options::TARGET_FILE_ROW_NUM, "2500"},
{Options::VECTOR_TARGET_FILE_SIZE, "1 b"},
{Options::WRITE_BATCH_SIZE, "100"},
{Options::READ_BATCH_SIZE, "127"},
{Options::DELETION_VECTORS_ENABLED, "true"}};
CreateTable(fields, /*partition_keys=*/{}, options);
const std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto input = PrepareBulkData(
6001,
[](int32_t id) {
const std::string vector = id % 1000 == 0 || id % 1000 == 999
? "null"
: fmt::format("[{}, {}, {}]", id, id + 1, id + 2);
return fmt::format(R"({}, {}, "row-{}")", id, vector, id);
},
fields);
std::vector<std::shared_ptr<arrow::Array>> batches;
for (int64_t offset = 0; offset < input->length(); offset += 1000) {
batches.push_back(input->Slice(offset, std::min<int64_t>(1000, input->length() - offset)));
}
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> messages,
WriteArrays(table_path, {}, {}, batches));
ASSERT_EQ(messages.size(), 1);
auto message = std::dynamic_pointer_cast<CommitMessageImpl>(messages[0]);
ASSERT_NE(message, nullptr);
int32_t main_count = 0;
int32_t vector_count = 0;
int64_t main_rows = 0;
int64_t vector_rows = 0;
for (const auto& file : message->data_increment_.new_files_) {
if (VectorStoreUtils::IsVectorStoreFile(file->file_name)) {
ASSERT_GT(main_count, 0);
ASSERT_GT(file->row_count, 0);
vector_rows += file->row_count;
ASSERT_LE(vector_rows, main_rows);
++vector_count;
} else {
ASSERT_EQ(main_rows, vector_rows);
main_rows += file->row_count;
++main_count;
}
}
ASSERT_EQ(main_count, 3);
ASSERT_EQ(main_rows, input->length());
ASSERT_EQ(vector_rows, main_rows);
// Lance checks encoded bytes only; a small file may stay buffered until the main file closes.
if (FileFormat() != "lance") {
ASSERT_GT(vector_count, main_count);
}
ASSERT_OK(Commit(table_path, messages));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::map<int64_t, std::shared_ptr<DataFileMeta>> main_files;
std::map<int64_t, std::shared_ptr<DataFileMeta>> vector_files;
for (const auto& split : splits) {
auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
ASSERT_NE(data_split, nullptr);
for (const auto& file : data_split->DataFiles()) {
ASSERT_OK_AND_ASSIGN(int64_t first_row_id, file->NonNullFirstRowId());
ASSERT_EQ(file->min_sequence_number, 1);
ASSERT_EQ(file->max_sequence_number, 1);
auto& files =
VectorStoreUtils::IsVectorStoreFile(file->file_name) ? vector_files : main_files;
ASSERT_TRUE(files.emplace(first_row_id, file).second);
}
}
ASSERT_EQ(main_files.size(), 3);
auto vector_file = vector_files.begin();
int64_t next_row_id = 0;
for (const auto& [first, main] : main_files) {
ASSERT_EQ(first, next_row_id);
const int64_t end = first + main->row_count;
while (next_row_id < end) {
ASSERT_NE(vector_file, vector_files.end());
ASSERT_EQ(vector_file->first, next_row_id);
next_row_id += vector_file->second->row_count;
ASSERT_LE(next_row_id, end);
++vector_file;
}
}
ASSERT_EQ(next_row_id, input->length());
ASSERT_EQ(vector_file, vector_files.end());
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, input));
auto normal_projection =
arrow::StructArray::Make({input->field(0), input->field(2)}, {fields[0], fields[2]})
.ValueOrDie();
ASSERT_OK(ScanAndRead(table_path, {"id", "tag"}, normal_projection));
auto vector_projection = arrow::StructArray::Make({input->field(1)}, {fields[1]}).ValueOrDie();
ASSERT_OK(ScanAndRead(table_path, {"embedding"}, vector_projection));
for (const Range& range : {Range(999, 1001), Range(2999, 3001), Range(5999, 6000)}) {
auto expected = checked_pointer_cast<arrow::StructArray>(
input->Slice(range.from, range.to - range.from + 1));
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, expected,
/*predicate=*/nullptr, /*row_ranges=*/{range}));
}
// Delete across main/VECTOR boundaries, including a whole VECTOR file and the last row.
std::set<int64_t> deleted_ids = {0, 999, 1000, 2499, 2500, 2999, 3000, 4999, 6000};
if (FileFormat() != "lance") {
const auto& deleted_file = std::next(vector_files.begin())->second;
for (int64_t i = 0; i < deleted_file->row_count; ++i) {
deleted_ids.insert(deleted_file->first_row_id.value() + i);
}
}
std::map<std::string, std::vector<int64_t>> deleted_positions;
for (const auto& [first, file] : main_files) {
for (int64_t id : deleted_ids) {
if (id >= first && id < first + file->row_count) {
deleted_positions[file->file_name].push_back(id - first);
}
}
}
ASSERT_OK(CommitDeletionVectors(table_path, messages[0], deleted_positions));
for (const std::vector<Range>& selection :
{std::vector<Range>{Range(0, 6000)},
std::vector<Range>{Range(998, 1002), Range(2498, 2502), Range(5998, 6000)}}) {
arrow::ArrayVector surviving_rows;
for (const auto& range : selection) {
int64_t next = range.from;
for (int64_t id : deleted_ids) {
if (id < next || id > range.to) {
continue;
}
if (next < id) {
surviving_rows.push_back(input->Slice(next, id - next));
}
next = id + 1;
}
if (next <= range.to) {
surviving_rows.push_back(input->Slice(next, range.to - next + 1));
}
}
auto expected = checked_pointer_cast<arrow::StructArray>(
arrow::Concatenate(surviving_rows).ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, expected,
/*predicate=*/nullptr, selection));
}
ASSERT_OK_AND_ASSIGN(messages, WriteArray(table_path, {}, input->Slice(0, 2)));
ASSERT_OK(Commit(table_path, messages));
auto appended = checked_pointer_cast<arrow::StructArray>(input->Slice(0, 2));
ASSERT_OK(ScanAndRead(table_path, {"id", "embedding", "tag"}, appended,
/*predicate=*/nullptr, /*row_ranges=*/{Range(6001, 6002)}));
}
TEST_P(DataEvolutionTableTest, TestVectorReadWrite) {
if (FileFormat() != "parquet" && FileFormat() != "lance") {
return;
}
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::READ_BATCH_SIZE, "2"},
};
CreateTable(fields, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto field_names = arrow::schema(fields)->field_names();
auto initial_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, [1, 2, 3], "a"],
[2, null, "b"],
[3, [7, 8, 9], "c"],
[4, [10, 11, 12], "d"],
[5, null, "e"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> initial_msgs,
WriteArray(table_path, field_names, initial_array));
ASSERT_OK(Commit(table_path, initial_msgs));
ASSERT_OK(ScanAndRead(table_path, field_names, initial_array));
auto appended_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[6, [16, 17, 18], "f"], [7, null, "g"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> appended_msgs,
WriteArray(table_path, field_names, appended_array));
ASSERT_OK(Commit(table_path, appended_msgs));
// Replace only the first row range, including transitions to and from NULL VECTOR values.
auto updated_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[null], [[4, 5, 6]], [[70, 80, 90]], [null], [[13, 14, 15]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> update_msgs,
WriteArray(table_path, {"embedding"}, updated_array));
SetFirstRowId(/*reset_first_row_id=*/0, update_msgs);
ASSERT_OK(Commit(table_path, update_msgs));
auto expected_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, null, "a"],
[2, [4, 5, 6], "b"],
[3, [70, 80, 90], "c"],
[4, null, "d"],
[5, [13, 14, 15], "e"],
[6, [16, 17, 18], "f"],
[7, null, "g"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, field_names, expected_array));
auto projected_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[[4, 5, 6]], [[70, 80, 90]], [null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"embedding"}, projected_array, /*predicate=*/nullptr,
/*row_ranges=*/{Range(1, 3)}));
auto null_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[null], [null], [null], [null], [null]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> null_msgs,
WriteArray(table_path, {"embedding"}, null_array));
SetFirstRowId(/*reset_first_row_id=*/0, null_msgs);
ASSERT_OK(Commit(table_path, null_msgs));
auto expected_null_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, null, "a"], [2, null, "b"], [3, null, "c"], [4, null, "d"],
[5, null, "e"], [6, [16, 17, 18], "f"], [7, null, "g"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, field_names, expected_null_array));
}
TEST_P(DataEvolutionTableTest, TestNestedVectorReadWrite) {
if (FileFormat() != "parquet") {
return;
}
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
arrow::FieldVector fields = {
arrow::field("id", arrow::int32()),
arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_type),
arrow::field("tag", arrow::utf8())})),
};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, "parquet"},
{Options::FILE_SYSTEM, "local"},
{Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::READ_BATCH_SIZE, "2"},
};
CreateTable(fields, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto initial_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, [[1, 2, 3], "a"]], [2, [null, "b"]], [3, null]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> initial_msgs,
WriteArray(table_path, {"id", "payload"}, initial_array));
ASSERT_OK(Commit(table_path, initial_msgs));
ASSERT_OK(ScanAndRead(table_path, {"id", "payload"}, initial_array));
auto updated_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), R"([
[null], [[[4, 5, 6], "updated"]], [[null, "c"]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> update_msgs,
WriteArray(table_path, {"payload"}, updated_array));
SetFirstRowId(/*reset_first_row_id=*/0, update_msgs);
ASSERT_OK(Commit(table_path, update_msgs));
auto expected_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, null], [2, [[4, 5, 6], "updated"]], [3, [null, "c"]]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"id", "payload"}, expected_array));
}
TEST_P(DataEvolutionTableTest, TestVectorSchemaEvolution) {
if (FileFormat() != "parquet" && FileFormat() != "lance") {
return;
}
auto retained_vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
auto dropped_vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::int64(), /*nullable=*/false), 2);
auto added_vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float64(), /*nullable=*/false), 2);
arrow::FieldVector fields_v0 = {
arrow::field("id", arrow::int32()),
arrow::field("retained_embedding", retained_vector_type),
arrow::field("dropped_embedding", dropped_vector_type),
};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{Options::READ_BATCH_SIZE, "1"},
};
CreateTable(fields_v0, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto initial_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_v0), R"([
[1, [1, 2, 3], [10, 11]], [2, null, [20, 21]]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> initial_msgs,
WriteArray(table_path, arrow::schema(fields_v0)->field_names(), initial_array));
ASSERT_OK(Commit(table_path, initial_msgs));
arrow::FieldVector fields_v1 = {fields_v0[0], fields_v0[1],
arrow::field("added_embedding", added_vector_type)};
ASSERT_OK(TestHelper::WriteNextSchema(
dir_->GetFileSystem(), table_path,
{DataField(0, fields_v1[0]), DataField(1, fields_v1[1]), DataField(3, fields_v1[2])},
/*highest_field_id=*/3, options));
auto expected_after_evolution = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_v1), R"([
[1, [1, 2, 3], null], [2, null, null]
])")
.ValueOrDie());
ASSERT_OK(
ScanAndRead(table_path, arrow::schema(fields_v1)->field_names(), expected_after_evolution));
auto added_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_v1[2]}), R"([
[[100, 101]], [null]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> added_msgs,
WriteArray(table_path, {"added_embedding"}, added_array));
SetFirstRowId(/*reset_first_row_id=*/0, added_msgs);
ASSERT_OK(Commit(table_path, added_msgs));
auto expected_after_partial_write = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_v1), R"([
[1, [1, 2, 3], [100, 101]], [2, null, null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(fields_v1)->field_names(),
expected_after_partial_write));
}
TEST_P(DataEvolutionTableTest, TestVectorSchemaEvolutionRejectsTypeChange) {
if (FileFormat() != "parquet" && FileFormat() != "lance") {
return;
}
auto vector_type =
arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3);
arrow::FieldVector fields = {arrow::field("id", arrow::int32()),
arrow::field("embedding", vector_type)};
std::map<std::string, std::string> options = {
{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::BUCKET, "-1"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
};
CreateTable(fields, /*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto initial_array = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
[1, [1, 2, 3]], [2, null]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> initial_msgs,
WriteArray(table_path, {"id", "embedding"}, initial_array));
ASSERT_OK(Commit(table_path, initial_msgs));
for (const auto& incompatible_type : {arrow::fixed_size_list(arrow::float32(), 5),
arrow::fixed_size_list(arrow::float64(), 3)}) {
SCOPED_TRACE(incompatible_type->ToString());
ASSERT_OK(TestHelper::WriteNextSchema(
dir_->GetFileSystem(), table_path,
{DataField(0, fields[0]), DataField(1, arrow::field("embedding", incompatible_type))},
/*highest_field_id=*/1, options));
ASSERT_NOK_WITH_MSG(ScanAndRead(table_path, {"id", "embedding"}, initial_array),
"VECTOR type mismatch during schema evolution");
}
}
TEST_P(DataEvolutionTableTest, TestAlterTable) {
auto file_format = FileFormat();
if (file_format == "mosaic" || file_format == "lance") {
return;
}
if (file_format == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + file_format +
"/append_table_alter_table_with_cast_with_data_evolution.db/"
"append_table_alter_table_with_cast_with_data_evolution";
std::vector<DataField> read_fields = {
DataField(6, arrow::field("f4", arrow::timestamp(arrow::TimeUnit::NANO))),
DataField(0, arrow::field("key0", arrow::int32())),
DataField(1, arrow::field("key1", arrow::int32())),
DataField(2, arrow::field("f3", arrow::int32())),
DataField(3, arrow::field("f1", arrow::utf8())),
DataField(4, arrow::field("f2", arrow::decimal128(6, 3))),
DataField(5, arrow::field("f0", arrow::boolean())),
DataField(8, arrow::field("f6", arrow::int32())),
SpecialFields::RowId(),
SpecialFields::SequenceNumber()};
std::shared_ptr<arrow::DataType> arrow_data_type =
DataField::ConvertDataFieldsToArrowStructType(read_fields);
{
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["1970-01-05T00:00", 0, 1, 100, "2024-11-26 06:38:56.001000001", "0.020", true, null, 0, 1],
["1969-11-18T00:00", 0, 1, 110, "2024-11-26 06:38:56.011000011", "11.120", true, null, 1, 1],
["1971-03-21T00:00", 0, 1, 120, "2024-11-26 06:38:56.021000021", "22.220", false, null, 2, 1],
["1957-11-01T00:00", 0, 1, 130, "2024-11-26 06:38:56.031000031", "333.320", false, null, 3, 1],
["2091-09-07T00:00", 0, 1, 140, "2024-11-26 06:38:56.041000041", "444.420", true, null, 4, 1],
["2024-11-26T06:38:56.054000154", 0, 1, 150, "2024-11-26 15:28:31", "55.002", true, 56, 5, 3],
["2024-11-26T06:38:56.064000164", 0, 1, 160, "2024-11-26 15:28:41", "666.012", false, 66, 6, 3],
["2024-11-26T06:38:56.074000174", 0, 1, 170, "2024-11-26 15:28:51", "-77.022", true, 76, 7, 3],
["2024-11-26T06:38:56.084000184", 0, 1, 180, "2024-11-26 15:29:01", "8.032", true, -86, 8, 3],
["2024-11-26T06:38:56.094000194", 0, 1, 190, "I'm strange", "-999.420", false, 96, 9, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array));
}
{
// only files with schema-1 will be skipped, while files with schema-0 will be reserved
// as type change
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"f3",
FieldType::INT, Literal(200));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["1970-01-05T00:00", 0, 1, 100, "2024-11-26 06:38:56.001000001", "0.020", true, null, 0, 1],
["1969-11-18T00:00", 0, 1, 110, "2024-11-26 06:38:56.011000011", "11.120", true, null, 1, 1],
["1971-03-21T00:00", 0, 1, 120, "2024-11-26 06:38:56.021000021", "22.220", false, null, 2, 1],
["1957-11-01T00:00", 0, 1, 130, "2024-11-26 06:38:56.031000031", "333.320", false, null, 3, 1],
["2091-09-07T00:00", 0, 1, 140, "2024-11-26 06:38:56.041000041", "444.420", true, null, 4, 1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate));
}
{
// files with schema-0 will be reserved as f6 does not exist in schema-0 (null count =
// null), files with schema-1 will also be reserved
auto predicate =
PredicateBuilder::IsNotNull(/*field_index=*/7, /*field_name=*/"f6", FieldType::INT);
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["1970-01-05T00:00", 0, 1, 100, "2024-11-26 06:38:56.001000001", "0.020", true, null, 0, 1],
["1969-11-18T00:00", 0, 1, 110, "2024-11-26 06:38:56.011000011", "11.120", true, null, 1, 1],
["1971-03-21T00:00", 0, 1, 120, "2024-11-26 06:38:56.021000021", "22.220", false, null, 2, 1],
["1957-11-01T00:00", 0, 1, 130, "2024-11-26 06:38:56.031000031", "333.320", false, null, 3, 1],
["2091-09-07T00:00", 0, 1, 140, "2024-11-26 06:38:56.041000041", "444.420", true, null, 4, 1],
["2024-11-26T06:38:56.054000154", 0, 1, 150, "2024-11-26 15:28:31", "55.002", true, 56, 5, 3],
["2024-11-26T06:38:56.064000164", 0, 1, 160, "2024-11-26 15:28:41", "666.012", false, 66, 6, 3],
["2024-11-26T06:38:56.074000174", 0, 1, 170, "2024-11-26 15:28:51", "-77.022", true, 76, 7, 3],
["2024-11-26T06:38:56.084000184", 0, 1, 180, "2024-11-26 15:29:01", "8.032", true, -86, 8, 3],
["2024-11-26T06:38:56.094000194", 0, 1, 190, "I'm strange", "-999.420", false, 96, 9, 3]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate));
}
{
std::vector<Range> row_ranges = {Range(1l, 1l)};
// test with row ids
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{5, 5});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["1969-11-18T00:00", 0, 1, 110, "2024-11-26 06:38:56.011000011", "11.120", true, null, 1, 1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
}
TEST_P(DataEvolutionTableTest, TestReadCompactFiles) {
if (FileFormat() == "mosaic" || FileFormat() == "lance") {
return;
}
auto file_format = FileFormat();
if (file_format == "avro") {
return;
}
std::string table_path =
paimon::test::GetDataDir() + file_format +
"/append_table_row_tracking_with_compact.db/append_table_row_tracking_with_compact";
std::vector<DataField> read_fields = {DataField(0, arrow::field("f0", arrow::utf8())),
DataField(1, arrow::field("f1", arrow::int32())),
DataField(2, arrow::field("f2", arrow::int32())),
DataField(3, arrow::field("f3", arrow::float64())),
SpecialFields::RowId(),
SpecialFields::SequenceNumber()};
std::shared_ptr<arrow::DataType> arrow_data_type =
DataField::ConvertDataFieldsToArrowStructType(read_fields);
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Lily", 2, 12, 2.1, 0, 1],
["Alice", 3, 13, 3.1, 1, 1],
["Bob", 4, 14, 4.1, 2, 2],
["David", 5, 15, 5.1, 3, 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array));
}
TEST_P(DataEvolutionTableTest, TestReadTableWithDenseStats) {
if (FileFormat() == "mosaic" || FileFormat() == "lance") {
return;
}
auto file_format = FileFormat();
if (file_format == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + file_format +
"/data_evolution_with_dense_stats.db/data_evolution_with_dense_stats";
std::vector<DataField> read_fields = {DataField(0, arrow::field("f0", arrow::utf8())),
DataField(1, arrow::field("f1", arrow::int32())),
DataField(2, arrow::field("f2", arrow::int32())),
DataField(3, arrow::field("f3", arrow::float64())),
SpecialFields::RowId(),
SpecialFields::SequenceNumber()};
std::shared_ptr<arrow::DataType> arrow_data_type =
DataField::ConvertDataFieldsToArrowStructType(read_fields);
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Lily", 2, 102, 2.1, 0, 2],
["Alice", 4, 104, 3.1, 1, 2]
])")
.ValueOrDie());
{
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array));
}
{
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(102));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate));
}
{
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(12));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
/*expected_array=*/nullptr, predicate));
}
{
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(6));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
/*expected_array=*/nullptr, predicate));
}
{
// f3 does not have stats, therefore data will not be filtered
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"f3",
FieldType::DOUBLE, Literal(5.1));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate));
}
{
// test row id with predicate
std::vector<Range> row_ranges = {Range(0l, 0l)};
auto predicate = PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1",
FieldType::INT, Literal(5));
CheckScanResult(table_path, /*predicate=*/predicate, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{}, /*expected_row_counts=*/{});
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
/*expected_array=*/nullptr, predicate,
/*row_ranges=*/row_ranges));
}
{
// test row id with predicate
std::vector<Range> row_ranges = {Range(0l, 0l)};
auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"f3",
FieldType::DOUBLE, Literal(5.1));
CheckScanResult(table_path, /*predicate=*/predicate, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{2, 2});
auto expected_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Lily", 2, 102, 2.1, 0, 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array0, predicate,
/*row_ranges=*/row_ranges));
}
}
TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) {
if (FileFormat() == "mosaic" || FileFormat() == "lance") {
return;
}
auto file_format = FileFormat();
if (file_format == "avro") {
return;
}
// only f2 has index
std::string table_path = paimon::test::GetDataDir() + file_format +
"/data_evolution_with_index.db/data_evolution_with_index";
std::vector<DataField> read_fields = {DataField(0, arrow::field("f0", arrow::utf8())),
DataField(1, arrow::field("f1", arrow::int32())),
DataField(2, arrow::field("f2", arrow::int32())),
DataField(3, arrow::field("f3", arrow::float64()))};
std::shared_ptr<arrow::DataType> arrow_data_type =
DataField::ConvertDataFieldsToArrowStructType(read_fields);
{
// first 4 records are two file with the same first row id with data evolution
// last 2 rows only in one file
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Lily", 2, 102, 2.1],
["Alice", 4, 104, 3.1],
["Bob", 6, 106, 4.1],
["David", 8, 108, 5.1],
[null, null, 202, 6.1],
[null, null, 204, 7.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array));
}
{
// The newer file owns f2; its bitmap selects the same row in every merged reader.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(102));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Lily", 2, 102, 2.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
// The bitmap proves that neither row range group contains f2 = 103.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(103));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
/*expected_array=*/nullptr, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
{
// Scan planning keeps the split, but reader-side indexes skip both row range groups.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(203));
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
/*expected_array=*/nullptr, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
{
// A single-file group applies the exact bitmap row selection.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(202));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
[null, null, 202, 6.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
auto predicate =
PredicateBuilder::IsNull(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING);
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
[null, null, 202, 6.1],
[null, null, 204, 7.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate));
}
{
// test row id with predicate
std::vector<Range> row_ranges = {Range(0l, 2l)};
// Intersect the group's shared bitmap with row ranges while preserving column alignment.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(106));
CheckScanResult(table_path, /*predicate=*/predicate, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 0}, /*expected_row_counts=*/{4, 4});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
["Bob", 6, 106, 4.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate,
/*row_ranges=*/row_ranges,
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
// test row id with predicate
std::vector<Range> row_ranges = {Range(4l, 5l)};
// The single-file bitmap selection is intersected with the row-id selection.
auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2",
FieldType::INT, Literal(204));
CheckScanResult(table_path, /*predicate=*/predicate, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{4}, /*expected_row_counts=*/{2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([
[null, null, 204, 7.1]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(),
expected_array, predicate,
/*row_ranges=*/row_ranges,
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
}
TEST_P(DataEvolutionTableTest, TestDataEvolutionPredicatePushDownBoundaries) {
if (FileFormat() == "mosaic" || FileFormat() == "lance") {
return;
}
auto file_format = FileFormat();
if (file_format == "avro") {
return;
}
std::string table_path = paimon::test::GetDataDir() + file_format +
"/data_evolution_with_index.db/data_evolution_with_index";
{
// A file without f0 must not interpret the predicate as f0 = null.
auto predicate = PredicateBuilder::Equal(
/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING,
Literal(FieldType::STRING, "Lily", 4));
auto read_type =
arrow::struct_({arrow::field("f0", arrow::utf8()), arrow::field("f2", arrow::int32())});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(read_type, R"([
["Lily", 102],
["Alice", 104],
["Bob", 106],
["David", 108],
[null, 202],
[null, 204]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f2"}, expected_array, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
// System fields are completed after reading and cannot be pushed into data files.
auto predicate = PredicateBuilder::Equal(
/*field_index=*/1, /*field_name=*/"_ROW_ID", FieldType::BIGINT, Literal(int64_t{99}));
auto read_type = arrow::struct_(
{arrow::field("f2", arrow::int32()), SpecialFields::RowId().ArrowField()});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(read_type, R"([
[102, 0],
[104, 1],
[106, 2],
[108, 3],
[202, 4],
[204, 5]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f2", "_ROW_ID"}, expected_array, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
// Dropping a system-field conjunct must not drop a pushable data conjunct.
auto data_predicate = PredicateBuilder::Equal(
/*field_index=*/0, /*field_name=*/"f2", FieldType::INT, Literal(103));
auto system_predicate = PredicateBuilder::Equal(
/*field_index=*/1, /*field_name=*/"_ROW_ID", FieldType::BIGINT, Literal(int64_t{0}));
ASSERT_OK_AND_ASSIGN(auto predicate,
PredicateBuilder::And({data_predicate, system_predicate}));
ASSERT_OK(ScanAndRead(table_path, {"f2", "_ROW_ID"}, /*expected_array=*/nullptr, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
{
// Bitmap positions compose with row ranges without changing the physical row id.
auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f2",
FieldType::INT, Literal(204));
auto read_type = arrow::struct_(
{arrow::field("f2", arrow::int32()), SpecialFields::RowId().ArrowField()});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(read_type, R"([
[204, 5]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f2", "_ROW_ID"}, expected_array, predicate,
/*row_ranges=*/{Range(4l, 5l)},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
{
// The bitmap selects row id 5 while the global-index selection keeps only row id 4.
auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f2",
FieldType::INT, Literal(204));
ASSERT_OK(ScanAndRead(table_path, {"f2", "_ROW_ID"}, /*expected_array=*/nullptr, predicate,
/*row_ranges=*/{Range(4l, 4l)},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
{
// The predicate keeps row ids {4, 5}; the global-index selection keeps {0, 1, 2, 3, 4}.
auto equal_202 = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f2",
FieldType::INT, Literal(202));
auto equal_204 = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f2",
FieldType::INT, Literal(204));
ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({equal_202, equal_204}));
auto read_type = arrow::struct_(
{arrow::field("f2", arrow::int32()), SpecialFields::RowId().ArrowField()});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(read_type, R"([
[202, 4]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f2", "_ROW_ID"}, expected_array, predicate,
/*row_ranges=*/{Range(0l, 4l)},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
}
TEST_P(DataEvolutionTableTest, TestFormatPredicatePushDownWithoutFileIndex) {
if (FileFormat() == "avro" || FileFormat() == "lance") {
return;
}
std::map<std::string, std::string> options = {{Options::FILE_INDEX_READ_ENABLED, "false"},
{Options::WRITE_BATCH_SIZE, "1"},
{"parquet.page.size", "1"},
{"parquet.enable.dictionary", "false"},
{"parquet.write.enable-page-index", "true"},
{"parquet.read.enable-page-index-filter", "true"},
{"orc.stripe.size", "1"},
{"orc.row.index.stride", "1"}};
if (FileFormat() == "mosaic") {
options.emplace(Options::FILE_BLOCK_SIZE, "1 B");
options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0");
}
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/false, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto input = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[2, "b", "y"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
std::vector<std::shared_ptr<arrow::Array>> write_arrays;
for (int64_t i = 0; i < input->length(); i++) {
arrow::ArrayVector children;
for (const auto& child : input->fields()) {
children.push_back(child->Slice(i, 1));
}
write_arrays.push_back(arrow::StructArray::Make(children, fields_).ValueOrDie());
}
ASSERT_OK_AND_ASSIGN(auto commit_messages, WriteArrays(table_path, /*partition=*/{},
{"f0", "f1", "f2"}, write_arrays));
ASSERT_OK(Commit(table_path, commit_messages));
auto predicate =
PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, Literal(3));
auto expected = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[3, "c", "z"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, expected, predicate,
/*row_ranges=*/{},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/true));
}
TEST_P(DataEvolutionTableTest, TestPredicate) {
if (FileFormat() == "avro" || FileFormat() == "lance") {
// Avro and Lance do not have stats.
return;
}
if (FileFormat() == "mosaic") {
CreateTable(/*partition_keys=*/{}, {{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"},
{mosaic::MOSAIC_STATS_COLUMNS, "f0,f1,f2"}});
} else {
CreateTable();
}
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// write field: f0, f1, f2
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols0, src_array0));
ASSERT_OK(Commit(table_path, commit_msgs));
// write field: f2
std::vector<std::string> write_cols1 = {"f2"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2]}), R"([
["c"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(commit_msgs, WriteArray(table_path, write_cols1, src_array1));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
{
// test no predicate
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
}
{
// test predicate with f2
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/2, /*field_name=*/"f2", FieldType::STRING,
Literal(FieldType::STRING, "b", 1));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "c"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array, predicate));
}
{
// test predicate with f1
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "a", 1));
ASSERT_OK(
ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr, predicate));
}
{
// test predicate with f2
auto predicate =
PredicateBuilder::NotEqual(/*field_index=*/2, /*field_name=*/"f2", FieldType::STRING,
Literal(FieldType::STRING, "c", 1));
ASSERT_OK(
ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr, predicate));
}
}
TEST_P(DataEvolutionTableTest, TestIOException) {
std::string table_path;
// write and commit with I/O exception
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");
CreateTable();
table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
io_hook->Reset(i, IOHook::Mode::RETURN_ERROR);
// write field: f0, f1, f2
std::vector<std::string> write_cols0 = schema->field_names();
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[10, "a", "b"],
[20, "aa", "bb"],
[23, "aaa", "bbb"]
])")
.ValueOrDie());
auto commit_msgs0_result = WriteArray(table_path, write_cols0, src_array0);
CHECK_HOOK_STATUS(commit_msgs0_result.status(), i);
CHECK_HOOK_STATUS(Commit(table_path, commit_msgs0_result.value()), i);
// write field: f2, f0
std::vector<std::string> write_cols1 = {"f2", "f0"};
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[2], fields_[0]}), R"([
["c", 100],
["cc", 200],
["ccc", 300]
])")
.ValueOrDie());
auto commit_msgs1_result = WriteArray(table_path, write_cols1, src_array1);
CHECK_HOOK_STATUS(commit_msgs1_result.status(), i);
SetFirstRowId(/*reset_first_row_id=*/0,
const_cast<std::vector<std::shared_ptr<paimon::CommitMessage>>&>(
commit_msgs1_result.value()));
CHECK_HOOK_STATUS(Commit(table_path, commit_msgs1_result.value()), i);
write_run_complete = true;
break;
}
ASSERT_TRUE(write_run_complete);
// scan and read with I/O exception
bool read_run_complete = false;
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], fields_[0], SpecialFields::SequenceNumber().ArrowField(),
SpecialFields::RowId().ArrowField(), fields_[2]}),
R"([
["a", 100, 2, 0, "c"],
["aa", 200, 2, 1, "cc"],
["aaa", 300, 2, 2, "ccc"]
])")
.ValueOrDie());
for (size_t i = 0; i < 2000; i++) {
ScopeGuard guard([&io_hook]() { io_hook->Clear(); });
io_hook->Reset(i, IOHook::Mode::RETURN_ERROR);
CHECK_HOOK_STATUS(ScanAndRead(table_path, {"f1", "f0", "_SEQUENCE_NUMBER", "_ROW_ID", "f2"},
expected_array),
i);
read_run_complete = true;
break;
}
ASSERT_TRUE(read_run_complete);
}
TEST_P(DataEvolutionTableTest, TestWithRowIds) {
std::map<std::string, std::string> options = {{Options::FILE_FORMAT, FileFormat()},
{Options::FILE_SYSTEM, "local"},
{Options::ROW_TRACKING_ENABLED, "true"},
{Options::DATA_EVOLUTION_ENABLED, "true"}};
if (FileFormat() == "mosaic") {
options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0,f1");
}
CreateTable(/*partition_keys=*/{}, options);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// turn 0: write field: f0, f1
std::vector<std::string> write_cols = {"f0", "f1"};
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[0, "a"],
[1, "b"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0, WriteArray(table_path, write_cols, src_array0));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs0);
ASSERT_OK(Commit(table_path, commit_msgs0));
// turn 1: write field: f0, f1
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[2, "c"],
[3, "d"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, write_cols, src_array1));
SetFirstRowId(/*reset_first_row_id=*/2, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
// turn 2: write field: f0, f1
auto src_array2 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields_[0], fields_[1]}), R"([
[4, "e"],
[5, "f"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs2, WriteArray(table_path, write_cols, src_array2));
SetFirstRowId(/*reset_first_row_id=*/4, commit_msgs2);
ASSERT_OK(Commit(table_path, commit_msgs2));
{
// test without row ids
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/{},
/*expected_first_row_ids=*/{0, 2, 4}, /*expected_row_counts=*/{2, 2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[0, "a", null],
[1, "b", null],
[2, "c", null],
[3, "d", null],
[4, "e", null],
[5, "f", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
}
{
// test row ids in first file
std::vector<Range> row_ranges = {Range(0l, 1l)};
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0}, /*expected_row_counts=*/{2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[0, "a", null],
[1, "b", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr, /*row_ranges=*/row_ranges));
}
{
// test row ids in last file
std::vector<Range> row_ranges = {Range(4l, 4l)};
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{4}, /*expected_row_counts=*/{2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[4, "e", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr, /*row_ranges=*/row_ranges));
}
{
// test row ids in multiple files
std::vector<Range> row_ranges = {Range(1l, 1l), Range(4l, 4l)};
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 4}, /*expected_row_counts=*/{2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "b", null],
[4, "e", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
{
// test all row ids
std::vector<Range> row_ranges = {Range(0l, 5l)};
CheckScanResult(table_path, /*predicate=*/nullptr, /*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 2, 4}, /*expected_row_counts=*/{2, 2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[0, "a", null],
[1, "b", null],
[2, "c", null],
[3, "d", null],
[4, "e", null],
[5, "f", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr, /*row_ranges=*/row_ranges));
}
{
// test unordered row ids
std::vector<Range> row_ranges = {Range(5l, 5l), Range(3l, 3l)};
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{2, 4}, /*expected_row_counts=*/{2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[3, "d", null],
[5, "f", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
{
// test row ids which partially exist
std::vector<Range> row_ranges = {Range(100l, 100l), Range(5l, 5l), Range(3l, 3l)};
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{2, 4}, /*expected_row_counts=*/{2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[3, "d", null],
[5, "f", null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
{
// test row ids which do not exist
std::vector<Range> row_ranges = {Range(100l, 100l), Range(200l, 200l)};
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{}, /*expected_row_counts=*/{});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr,
/*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
if (FileFormat() == "avro" || FileFormat() == "lance") {
// Avro and Lance do not support stats.
return;
}
{
// test row id with predicate
// row ids {0, 1, 5} and predicate f0 = 2
// scan will filter out all files
std::vector<Range> row_ranges = {Range(0l, 0l), Range(1l, 1l), Range(5l, 5l)};
auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(2));
CheckScanResult(table_path, /*predicate=*/predicate,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{}, /*expected_row_counts=*/{});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), /*expected_array=*/nullptr,
predicate,
/*row_ranges=*/row_ranges));
}
{
// test row id with predicate
// row ids {0, 1, 5} and predicate f0 = 5
// scan will filter out file0 and file1
std::vector<Range> row_ranges = {Range(0l, 0l), Range(1l, 1l), Range(5l, 5l)};
auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(5));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[5, "f", null]
])")
.ValueOrDie());
CheckScanResult(table_path, /*predicate=*/predicate,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{4}, /*expected_row_counts=*/{2});
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array, predicate,
/*row_ranges=*/row_ranges));
}
{
// test with row tracking fields
std::vector<Range> row_ranges = {Range(0l, 0l), Range(1l, 1l), Range(5l, 5l)};
CheckScanResult(table_path, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges,
/*expected_first_row_ids=*/{0, 4}, /*expected_row_counts=*/{2, 2});
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_({fields_[1], fields_[0],
SpecialFields::SequenceNumber().ArrowField(),
SpecialFields::RowId().ArrowField(), fields_[2]}),
R"([
["a", 0, 1, 0, null],
["b", 1, 1, 1, null],
["f", 5, 3, 5, null]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f1", "f0", "_SEQUENCE_NUMBER", "_ROW_ID", "f2"},
expected_array, /*predicate=*/nullptr,
/*row_ranges=*/row_ranges));
}
}
TEST_P(DataEvolutionDvTableTest, TestGlobalIndexScoresWithSingleFileFiltering) {
CheckGlobalIndexScoresWithFiltering(/*merge_files=*/false);
}
TEST_P(DataEvolutionDvTableTest, TestGlobalIndexScoresWithMergedFileFiltering) {
CheckGlobalIndexScoresWithFiltering(/*merge_files=*/true);
}
TEST_P(DataEvolutionDvTableTest, TestGlobalIndexLiveRows) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true, /*extra_options=*/{},
/*partition_keys=*/{"f1"});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> group_msgs0,
WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3},
/*partition=*/{{"f1", "p0"}}));
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> group_msgs1,
WriteAndCommitGroup(table_path, /*first_row_id=*/4, /*f0_values=*/{4, 5, 6, 7},
/*partition=*/{{"f1", "p1"}}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchor_file_names.size(), 2);
std::shared_ptr<FileSystem> fs = dir_->GetFileSystem();
SchemaManager schema_manager(fs, table_path);
ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema,
schema_manager.Latest());
ASSERT_TRUE(table_schema);
ASSERT_OK_AND_ASSIGN(CoreOptions options,
CoreOptions::FromMap(table_schema.value()->Options(), fs));
SnapshotManager snapshot_manager(fs, table_path);
ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot_before_deletion,
snapshot_manager.LatestSnapshot());
ASSERT_TRUE(snapshot_before_deletion);
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs0[0],
{{anchor_file_names[0], /*deleted_positions=*/{1, 3}}}));
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs1[0],
{{anchor_file_names[1], /*deleted_positions=*/{0}}}));
ASSERT_OK_AND_ASSIGN(std::optional<Snapshot> snapshot, snapshot_manager.LatestSnapshot());
ASSERT_TRUE(snapshot);
std::shared_ptr<Executor> executor = CreateDefaultExecutor();
auto read_live_rows = [&](const Snapshot& read_snapshot,
const std::shared_ptr<PredicateFilter>& partition_filter,
const std::optional<RowRangeIndex>& row_range_index) {
return GlobalIndexLiveRowFilter::LiveRows(table_path, table_schema.value(), read_snapshot,
partition_filter, row_range_index, options,
executor, GetDefaultPool());
};
ASSERT_OK_AND_ASSIGN(std::optional<RoaringBitmap64> live_rows,
read_live_rows(snapshot.value(), /*partition_filter=*/nullptr,
/*row_range_index=*/std::nullopt));
ASSERT_TRUE(live_rows);
ASSERT_EQ(live_rows.value(), RoaringBitmap64::From({0, 2, 5, 6, 7}));
ASSERT_OK_AND_ASSIGN(RowRangeIndex p1_row_ranges, RowRangeIndex::Create({Range(4, 5)}));
ASSERT_OK_AND_ASSIGN(
live_rows, read_live_rows(snapshot.value(), /*partition_filter=*/nullptr, p1_row_ranges));
ASSERT_TRUE(live_rows);
ASSERT_EQ(live_rows.value(), RoaringBitmap64::From({5, 6, 7}));
auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema.value()->Fields());
ASSERT_OK_AND_ASSIGN(std::shared_ptr<PredicateFilter> p0_filter,
FileStoreScan::CreatePartitionPredicate(
table_schema.value()->PartitionKeys(),
options.GetPartitionDefaultName(), arrow_schema, {{{"f1", "p0"}}}));
ASSERT_OK_AND_ASSIGN(live_rows, read_live_rows(snapshot.value(), p0_filter,
/*row_range_index=*/std::nullopt));
ASSERT_TRUE(live_rows);
ASSERT_EQ(live_rows.value(), RoaringBitmap64::From({0, 2}));
ASSERT_OK_AND_ASSIGN(live_rows, read_live_rows(snapshot_before_deletion.value(),
/*partition_filter=*/nullptr,
/*row_range_index=*/std::nullopt));
ASSERT_TRUE(live_rows);
ASSERT_EQ(live_rows.value(), RoaringBitmap64::From({0, 1, 2, 3, 4, 5, 6, 7}));
ASSERT_OK(fs->Delete(PathUtil::JoinPath(table_path, "index")));
ASSERT_NOK_WITH_MSG(read_live_rows(snapshot.value(), /*partition_filter=*/nullptr,
/*row_range_index=*/std::nullopt),
"failed to read deletion vector of data file");
}
TEST_P(DataEvolutionDvTableTest, TestReadWithDeletionVectors) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// full-row write assigns row ids 0-3, producing the anchor file of the row range group
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[2, "b", "y"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0,
WriteArray(table_path, schema->field_names(), src_array));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs0);
ASSERT_OK(Commit(table_path, commit_msgs0));
// partial write of f2 over the same row range: the group merges columns from two files
arrow::FieldVector f2_fields = {fields_[2]};
auto update_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(f2_fields), R"([
["x2"],
["y2"],
["z2"],
["w2"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1, WriteArray(table_path, {"f2"}, update_array));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
auto expected_all = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x2"],
[2, "b", "y2"],
[3, "c", "z2"],
[4, "d", "w2"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_all));
ASSERT_OK_AND_ASSIGN(std::string anchor_file_name, PlannedAnchorFileName(table_path));
ASSERT_OK(CommitDeletionVectors(table_path, commit_msgs0[0],
{{anchor_file_name, /*deleted_positions=*/{1, 3}}}));
// both files of the group must drop the same rows to keep the column merge aligned
auto expected_deleted = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x2"],
[3, "c", "z2"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_deleted));
auto expected_with_row_id = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(
arrow::struct_(
{fields_[0], fields_[1], fields_[2], SpecialFields::RowId().ArrowField()}),
R"([
[1, "a", "x2", 0],
[3, "c", "z2", 2]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID"}, expected_with_row_id));
// a row-range selection composes with the deletion vector: rows {1, 2} minus deleted {1}
auto expected_selected = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[3, "c", "z2"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_selected,
/*predicate=*/nullptr, /*row_ranges=*/{Range(1, 2)}));
}
TEST_P(DataEvolutionDvTableTest, TestIndexSkipsBeforeReadingDeletionVector) {
CreateDataEvolutionTable(
/*deletion_vectors_enabled=*/true,
{{"file-index.bitmap.columns", "f2"}, {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto data = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_),
R"([[0,"a0","x0"],[1,"a1","x1"],
[2,"a2","x2"],[3,"a3","x3"]])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> messages,
WriteArray(table_path, {}, {"f0", "f1", "f2"}, data));
SetFirstRowId(100, messages);
ASSERT_OK(Commit(table_path, messages));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> merged_messages,
WriteAndCommitGroup(table_path, 200, {0, 1, 2, 3}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchors, PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchors.size(), 2);
ASSERT_OK(CommitDeletionVectors(table_path, messages[0], {{anchors[0], {2}}}));
ASSERT_OK(CommitDeletionVectors(table_path, merged_messages[0], {{anchors[1], {2}}}));
{
// Indexes keep positions 1 and 2 in both groups; their DVs remove position 2.
auto predicate = PredicateBuilder::In(
2, "f2", FieldType::STRING,
{Literal(FieldType::STRING, "x1", 2), Literal(FieldType::STRING, "x2", 2),
Literal(FieldType::STRING, "y1", 2), Literal(FieldType::STRING, "y2", 2)});
auto expected = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_),
R"([[1,"a1","x1"],[1,"a1","y1"]])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, expected, predicate, {},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
}
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
std::set<std::string> dv_paths;
for (const auto& split : splits) {
auto data_split = checked_pointer_cast<DataSplitImpl>(split);
for (const auto& deletion_file : data_split->DeletionFiles()) {
if (deletion_file) {
dv_paths.insert(deletion_file->path);
}
}
}
ASSERT_FALSE(dv_paths.empty());
for (const auto& path : dv_paths) {
ASSERT_OK(dir_->GetFileSystem()->Delete(path, /*recursive=*/false));
}
auto missing = PredicateBuilder::Equal(2, "f2", FieldType::STRING,
Literal(FieldType::STRING, "missing", 7));
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, nullptr, missing, {},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
auto single_match =
PredicateBuilder::Equal(2, "f2", FieldType::STRING, Literal(FieldType::STRING, "x1", 2));
auto merged_match =
PredicateBuilder::Equal(2, "f2", FieldType::STRING, Literal(FieldType::STRING, "y1", 2));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> match,
PredicateBuilder::Or({single_match, merged_match}));
// Both single-file and merged groups are rejected by rowRanges intersected with the bitmap.
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, nullptr, match,
{Range(100, 100), Range(200, 200)},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
// A selected matching row still needs its DV; do not hide an actual missing sidecar.
ASSERT_NOK(ScanAndRead(table_path, {"f0", "f1", "f2"}, nullptr, match, {Range(101, 101)},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
TEST_P(DataEvolutionDvTableTest, TestSharedBitmapSelectionComposesWithDeletionVectors) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true,
{{"file-index.bitmap.columns", "f0,f2"},
{Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"},
{Options::READ_BATCH_SIZE, "2"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> first,
WriteAndCommitGroup(table_path, 100, {0, 1, 2, 3}, {},
{{"file-index.bitmap.columns", "f2"}}));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> second,
WriteAndCommitGroup(table_path, 200, {0, 1, 2, 3}, {},
{{"file-index.bitmap.columns", "f2"}}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchors, PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchors.size(), 2);
ASSERT_OK(CommitDeletionVectors(table_path, first[0], {{anchors[0], {2}}}));
ASSERT_OK(CommitDeletionVectors(table_path, second[0], {{anchors[1], {1}}}));
auto one =
PredicateBuilder::Equal(2, "f2", FieldType::STRING, Literal(FieldType::STRING, "y1", 2));
auto two =
PredicateBuilder::Equal(2, "f2", FieldType::STRING, Literal(FieldType::STRING, "y2", 2));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> predicate, PredicateBuilder::Or({one, two}));
arrow::FieldVector fields = fields_;
fields.push_back(SpecialFields::RowId().ArrowField());
auto expected = checked_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields),
R"([[1,"a1","y1",101],[2,"a2","y2",202]])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID"}, expected, predicate, {},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
// The newer f2 index must win over its stale copy in the anchor file.
auto newer =
PredicateBuilder::Equal(2, "f2", FieldType::STRING, Literal(FieldType::STRING, "y1", 2));
auto expected_one = checked_pointer_cast<arrow::StructArray>(expected->Slice(0, 1));
auto small = PredicateBuilder::Equal(0, "f0", FieldType::INT, Literal(1));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> combined,
PredicateBuilder::And({small, predicate}));
// f0 and f2 are won by different files. Intersect both indexes before applying the DV.
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID"}, expected_one, combined, {},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID"}, expected_one, newer, {},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> splits, PlanSplits(table_path));
for (const auto& split : splits) {
auto data_split = checked_pointer_cast<DataSplitImpl>(split);
for (const auto& file : data_split->DataFiles()) {
if (file->first_row_id == 200) {
ASSERT_OK(dir_->GetFileSystem()->Delete(
PathUtil::JoinPath(data_split->BucketPath(), file->file_name), false));
}
}
}
// The second group's only bitmap hit is deleted: skip it before opening any data file.
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "_ROW_ID"}, expected_one, one, {},
/*check_scan_plan_when_empty_result=*/true,
/*apply_predicate_to_scan=*/false));
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, nullptr, one, {Range(200, 203)},
/*check_scan_plan_when_empty_result=*/false,
/*apply_predicate_to_scan=*/false));
}
TEST_P(DataEvolutionDvTableTest, TestReadWithDeletionVectorsAcrossReadBatches) {
// the 12 rows below span several read batches, so the deletion vector empties a whole
// batch of every file of the group: each file reader then skips that batch entirely and
// the column merge has to stay aligned on the surviving row count alone
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true, {{Options::READ_BATCH_SIZE, "4"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> group_msgs,
WriteAndCommitGroup(table_path, /*first_row_id=*/0,
/*f0_values=*/{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}));
// positions 4-7 cover a whole read batch, position 9 only part of the next one
ASSERT_OK_AND_ASSIGN(std::string anchor_file_name, PlannedAnchorFileName(table_path));
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs[0],
{{anchor_file_name, /*deleted_positions=*/{4, 5, 6, 7, 9}}}));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[0, "a0", "y0"],
[1, "a1", "y1"],
[2, "a2", "y2"],
[3, "a3", "y3"],
[8, "a8", "y8"],
[10, "a10", "y10"],
[11, "a11", "y11"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
arrow::FieldVector row_id_fields = {SpecialFields::RowId().ArrowField()};
auto expected_row_ids = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(row_id_fields), R"([
[0], [1], [2], [3], [8], [10], [11]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"_ROW_ID"}, expected_row_ids));
// a row-range selection that spans the fully deleted batch composes with it
auto expected_selected = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[3, "a3", "y3"],
[8, "a8", "y8"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_selected,
/*predicate=*/nullptr, /*row_ranges=*/{Range(3, 8)}));
}
TEST_P(DataEvolutionDvTableTest, TestReadWithDeletionVectorsOnPartOfRowRangeGroups) {
// one split per row range group, so a group's deletion file must not reach the other
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true,
{{Options::SOURCE_SPLIT_TARGET_SIZE, "1"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
auto src_array0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[2, "b", "y"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs0,
WriteArray(table_path, schema->field_names(), src_array0));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs0);
ASSERT_OK(Commit(table_path, commit_msgs0));
auto src_array1 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[5, "e", "v"],
[6, "f", "u"],
[7, "g", "t"],
[8, "h", "s"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs1,
WriteArray(table_path, schema->field_names(), src_array1));
SetFirstRowId(/*reset_first_row_id=*/4, commit_msgs1);
ASSERT_OK(Commit(table_path, commit_msgs1));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchor_file_names.size(), 2);
ASSERT_OK(CommitDeletionVectors(table_path, commit_msgs0[0],
{{anchor_file_names[0], /*deleted_positions=*/{1, 3}}}));
// the deletion vector applies to its own group only, the other group keeps every row
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[3, "c", "z"],
[5, "e", "v"],
[6, "f", "u"],
[7, "g", "t"],
[8, "h", "s"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
// a projection-only read drops the same rows
arrow::FieldVector f0_fields = {fields_[0]};
auto expected_f0 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(f0_fields), R"([
[1], [3], [5], [6], [7], [8]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0"}, expected_f0));
arrow::FieldVector row_id_fields = {SpecialFields::RowId().ArrowField()};
auto expected_row_ids = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(row_id_fields), R"([
[0], [2], [4], [5], [6], [7]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"_ROW_ID"}, expected_row_ids));
}
TEST_P(DataEvolutionDvTableTest, TestReadWithDeletionVectorsOnEveryRowRangeGroup) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// one split (the default target size keeps both groups together), so a single split
// deletion vector factory serves two groups anchored at different row ids
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> group_msgs0,
WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4, /*f0_values=*/{4, 5, 6, 7}));
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchor_file_names.size(), 2);
// Positions are anchor-relative, so the groups deliberately delete different ones: group 0
// drops {1, 3} of row ids 0-3, group 1 drops {0, 2} of row ids 4-7. Reading a group with the
// other group's vector, or with its anchor range as the shift base, cannot match below.
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs0[0],
{{anchor_file_names[0], /*deleted_positions=*/{1, 3}},
{anchor_file_names[1], /*deleted_positions=*/{0, 2}}}));
// The read looks a group's deletion vector up by its anchor file name, so the scan has to
// hand it exactly that. Asserting it separates a scan-side mix-up from a read-side one.
ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<Split>> planned_splits,
PlanSplits(table_path));
ASSERT_EQ(planned_splits.size(), 1);
auto planned_split_impl = std::dynamic_pointer_cast<DataSplitImpl>(planned_splits[0]);
ASSERT_TRUE(planned_split_impl);
ASSERT_OK_AND_ASSIGN(DeletionCardinalityMap cardinality_by_file,
DeletionCardinalityByDataFile(planned_splits[0]));
DeletionCardinalityMap expected_cardinalities = {{anchor_file_names[0], 2},
{anchor_file_names[1], 2}};
ASSERT_EQ(cardinality_by_file, expected_cardinalities);
// the same split reports the surviving row count the limit push down prunes on: the two
// groups hold 4 rows each and each deletion vector drops 2 of them
ASSERT_OK_AND_ASSIGN(std::optional<int64_t> merged_row_count,
planned_split_impl->MergedRowCount());
ASSERT_EQ(std::optional<int64_t>(4), merged_row_count);
// a count query answers from that metadata alone, never reading a row, so the deletion
// vectors have to reach it too: without them it reports the 8 rows the files hold
ReadContextBuilder count_context_builder(table_path);
count_context_builder.SetReadFieldNames(schema->field_names());
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> count_context,
count_context_builder.Finish());
ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> count_table_read,
TableRead::Create(std::move(count_context)));
ASSERT_OK_AND_ASSIGN(std::unique_ptr<CountReader> count_reader,
count_table_read->CreateCountReader(planned_splits));
ASSERT_OK_AND_ASSIGN(int64_t counted_rows, count_reader->CountRows());
ASSERT_EQ(counted_rows, 4);
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[0, "a0", "y0"],
[2, "a2", "y2"],
[5, "a1", "y1"],
[7, "a3", "y3"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
arrow::FieldVector row_id_fields = {SpecialFields::RowId().ArrowField()};
auto expected_row_ids = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(row_id_fields), R"([
[0], [2], [5], [7]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"_ROW_ID"}, expected_row_ids));
// a row-range selection straddling the group boundary composes with both deletion vectors:
// row ids {2, 3, 4, 5} minus the deleted {3, 4}
auto expected_selected = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[2, "a2", "y2"],
[5, "a1", "y1"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_selected,
/*predicate=*/nullptr, /*row_ranges=*/{Range(2, 5)}));
}
TEST_P(DataEvolutionDvTableTest, TestReadWithFullyDeletedRowRangeGroup) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
// one split (the default target size keeps both groups together), so the emptied group's
// readers are concatenated with the surviving group's
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> group_msgs0,
WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4, /*f0_values=*/{4, 5, 6, 7}));
// both file readers of the first group then yield nothing, and its column merge must
// produce no rows at all instead of misaligning
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchor_file_names.size(), 2);
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs0[0],
{{anchor_file_names[0], /*deleted_positions=*/{0, 1, 2, 3}}}));
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[4, "a0", "y0"],
[5, "a1", "y1"],
[6, "a2", "y2"],
[7, "a3", "y3"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_array));
arrow::FieldVector row_id_fields = {SpecialFields::RowId().ArrowField()};
auto expected_row_ids = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(row_id_fields), R"([
[4], [5], [6], [7]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"_ROW_ID"}, expected_row_ids));
// a row-range selection spanning both groups keeps only what survives in the second one
auto expected_selected = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[4, "a0", "y0"],
[5, "a1", "y1"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_selected,
/*predicate=*/nullptr, /*row_ranges=*/{Range(0, 5)}));
// a selection covering only deleted rows returns nothing. The plan is asserted non-empty
// too: the scan cannot prune the split on row ids alone, so the emptiness comes from the
// deletion vector rather than from a plan with nothing to read.
ASSERT_OK_AND_ASSIGN(LimitScanResult only_deleted,
ScanAndReadWithLimit(table_path, schema->field_names(), /*limit=*/100,
/*predicate=*/nullptr,
/*row_ranges=*/{Range(0, 3)}));
ASSERT_FALSE(only_deleted.splits.empty());
ASSERT_FALSE(only_deleted.rows);
}
TEST_P(DataEvolutionDvTableTest, TestReadAfterUpdatingDeletionVectors) {
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto schema = arrow::schema(fields_);
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[2, "b", "y"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, schema->field_names(), src_array));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
ASSERT_OK_AND_ASSIGN(std::string anchor_file_name, PlannedAnchorFileName(table_path));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<CommitMessage> first_dv_msg,
CommitDeletionVectors(table_path, commit_msgs[0],
{{anchor_file_name, /*deleted_positions=*/{1}}}));
auto expected_after_first = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_after_first));
// a second deletion vector replaces the first one instead of both staying live
ASSERT_OK(CommitDeletionVectors(table_path, commit_msgs[0],
{{anchor_file_name, /*deleted_positions=*/{1, 3}}},
/*replaced_commit_msg=*/first_dv_msg));
auto expected_after_update = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[3, "c", "z"]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, schema->field_names(), expected_after_update));
}
TEST_P(DataEvolutionDvTableTest, TestReadWithDeletionVectorsAfterAddingColumn) {
if (FileFormat() == "avro") {
GTEST_SKIP() << "Avro has no stats, which the added column's scan pruning relies on";
}
std::map<std::string, std::string> options =
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([
[1, "a", "x"],
[2, "b", "y"],
[3, "c", "z"],
[4, "d", "w"]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto commit_msgs,
WriteArray(table_path, arrow::schema(fields_)->field_names(), src_array));
SetFirstRowId(/*reset_first_row_id=*/0, commit_msgs);
ASSERT_OK(Commit(table_path, commit_msgs));
// add column f3, then fill it for the same row range: the group merges columns from two
// files written under different schema ids
auto f3 = arrow::field("f3", arrow::int64());
ASSERT_OK(TestHelper::WriteNextSchema(dir_->GetFileSystem(), table_path,
{DataField(0, fields_[0]), DataField(1, fields_[1]),
DataField(2, fields_[2]), DataField(3, f3)},
/*highest_field_id=*/3, options));
auto f3_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f3}), R"([
[10], [20], [30], [40]
])")
.ValueOrDie());
ASSERT_OK_AND_ASSIGN(auto f3_commit_msgs, WriteArray(table_path, {"f3"}, f3_array));
SetFirstRowId(/*reset_first_row_id=*/0, f3_commit_msgs);
ASSERT_OK(Commit(table_path, f3_commit_msgs));
// the deletion vector is still anchored on the oldest normal file, written before the
// column was added
ASSERT_OK_AND_ASSIGN(std::string anchor_file_name, PlannedAnchorFileName(table_path));
ASSERT_OK(CommitDeletionVectors(table_path, commit_msgs[0],
{{anchor_file_name, /*deleted_positions=*/{1, 3}}}));
arrow::FieldVector evolved_fields = {fields_[0], fields_[1], fields_[2], f3};
auto expected_array = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(evolved_fields), R"([
[1, "a", "x", 10],
[3, "c", "z", 30]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2", "f3"}, expected_array));
// projecting only the added column keeps the same surviving rows
auto expected_f3 = std::dynamic_pointer_cast<arrow::StructArray>(
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f3}), R"([
[10], [30]
])")
.ValueOrDie());
ASSERT_OK(ScanAndRead(table_path, {"f3"}, expected_f3));
}
TEST_P(DataEvolutionDvTableTest, TestLimitPushDownWithHeavilyDeletedFirstRowRangeGroup) {
// one split per row range group, so the limit has to span both to be satisfied
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/true,
{{Options::SOURCE_SPLIT_TARGET_SIZE, "1"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK_AND_ASSIGN(
std::vector<std::shared_ptr<CommitMessage>> group_msgs0,
WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4, /*f0_values=*/{4, 5, 6, 7}));
// the first group keeps a single surviving row
ASSERT_OK_AND_ASSIGN(std::vector<std::string> anchor_file_names,
PlannedAnchorFileNames(table_path));
ASSERT_EQ(anchor_file_names.size(), 2);
ASSERT_OK(CommitDeletionVectors(table_path, group_msgs0[0],
{{anchor_file_names[0], /*deleted_positions=*/{0, 1, 2}}}));
// the first split alone satisfies a limit of 1: it still holds the one row that survived
// the deletion vector
ASSERT_OK_AND_ASSIGN(LimitScanResult limit_1,
ScanAndReadWithLimit(table_path, {"f0"}, /*limit=*/1));
ASSERT_EQ(limit_1.splits.size(), 1);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> limit_1_values, CollectF0Values(limit_1.rows));
ASSERT_EQ(limit_1_values, (std::vector<int32_t>{3}));
// the first split contributes only one surviving row, so a limit of 3 needs the second one
ASSERT_OK_AND_ASSIGN(LimitScanResult limit_3,
ScanAndReadWithLimit(table_path, {"f0"}, /*limit=*/3));
ASSERT_EQ(limit_3.splits.size(), 2);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> limit_3_values, CollectF0Values(limit_3.rows));
ASSERT_EQ(limit_3_values, (std::vector<int32_t>{3, 4, 5, 6, 7}));
}
TEST_P(DataEvolutionTableTest, TestLimitPushDownDisabledByNonPartitionFilter) {
// one split per row range group, so the plan can drop the group holding the matches
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/false,
{{Options::SOURCE_SPLIT_TARGET_SIZE, "1"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
// First row range group: no f0 value lies in [100, 200], but the 300 keeps the group's stats
// range straddling the filter so the scan cannot prune it. It therefore reaches the plan
// reporting four rows, and contributes none of them to the result.
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 300}));
// second row range group: every f0 value matches
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4,
/*f0_values=*/{100, 101, 102, 103}));
// the metadata row count of the first split alone satisfies the limit, but the filter runs
// while reading and drops every one of its rows. Pruning the plan on the metadata count
// would return nothing, so the push down has to be skipped and both splits kept.
auto at_least_100 = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(100));
auto at_most_200 = PredicateBuilder::LessOrEqual(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(200));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> predicate,
PredicateBuilder::And({at_least_100, at_most_200}));
// filtering while reading is what makes the metadata count an upper bound, so the read has
// to opt into it for the rows below to show what the push down would have thrown away
ASSERT_OK_AND_ASSIGN(LimitScanResult limited,
ScanAndReadWithLimit(table_path, {"f0"}, /*limit=*/2, predicate,
/*row_ranges=*/{},
/*enable_predicate_filter=*/true));
ASSERT_EQ(limited.splits.size(), 2);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> values, CollectF0Values(limited.rows));
ASSERT_EQ(values, (std::vector<int32_t>{100, 101, 102, 103}));
}
TEST_P(DataEvolutionTableTest, TestLimitPushDownKeptByPartitionFilter) {
std::vector<std::string> partition_keys = {"f1"};
CreateTable(partition_keys);
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
// one row range group per partition, each holding four rows with its own f0 range. A split
// never spans partitions, so the plan holds one split per partition for the push down to
// prune.
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3},
/*partition=*/{{"f1", "p0"}}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4,
/*f0_values=*/{10, 11, 12, 13}, /*partition=*/{{"f1", "p1"}}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/8,
/*f0_values=*/{20, 21, 22, 23}, /*partition=*/{{"f1", "p2"}}));
// a predicate on the partition key alone is evaluated while planning, never while reading,
// so every row a surviving split reports is actually returned and the push down stays on
auto not_p0 =
PredicateBuilder::NotEqual(/*field_index=*/1, /*field_name=*/"f1", FieldType::STRING,
Literal(FieldType::STRING, "p0", 2));
// baseline: a limit no split combination can reach prunes nothing, so this only shows which
// splits the partition filter itself leaves behind
ASSERT_OK_AND_ASSIGN(LimitScanResult unpruned,
ScanAndReadWithLimit(table_path, {"f0", "f1"}, /*limit=*/100, not_p0));
ASSERT_EQ(unpruned.splits.size(), 2);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> unpruned_values, CollectF0Values(unpruned.rows));
std::sort(unpruned_values.begin(), unpruned_values.end());
ASSERT_EQ(unpruned_values, (std::vector<int32_t>{10, 11, 12, 13, 20, 21, 22, 23}));
// the first matching split already holds four rows, so a limit of 2 drops the second one.
// Treating the partition filter as a read-time filter would skip the push down and keep both.
ASSERT_OK_AND_ASSIGN(LimitScanResult limited,
ScanAndReadWithLimit(table_path, {"f0", "f1"}, /*limit=*/2, not_p0));
ASSERT_EQ(limited.splits.size(), 1);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> limited_values, CollectF0Values(limited.rows));
std::sort(limited_values.begin(), limited_values.end());
ASSERT_TRUE(limited_values == (std::vector<int32_t>{10, 11, 12, 13}) ||
limited_values == (std::vector<int32_t>{20, 21, 22, 23}))
<< "unexpected kept rows for limit 2";
// one non-partition conjunct added to the same partition filter: the predicate is no longer
// settled while planning, so the push down is skipped and the split the limit would have
// dropped stays
auto at_least_10 = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"f0",
FieldType::INT, Literal(10));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> mixed_predicate,
PredicateBuilder::And({not_p0, at_least_10}));
ASSERT_OK_AND_ASSIGN(LimitScanResult mixed, ScanAndReadWithLimit(table_path, {"f0", "f1"},
/*limit=*/2, mixed_predicate));
ASSERT_EQ(mixed.splits.size(), 2);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> mixed_values, CollectF0Values(mixed.rows));
std::sort(mixed_values.begin(), mixed_values.end());
ASSERT_EQ(mixed_values, (std::vector<int32_t>{10, 11, 12, 13, 20, 21, 22, 23}));
}
TEST_P(DataEvolutionTableTest, TestLimitPushDownDisabledByRowRangeIndex) {
// one split per row range group, so the plan can drop the group holding the selection
CreateDataEvolutionTable(/*deletion_vectors_enabled=*/false,
{{Options::SOURCE_SPLIT_TARGET_SIZE, "1"}});
std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/0, /*f0_values=*/{0, 1, 2, 3}));
ASSERT_OK(WriteAndCommitGroup(table_path, /*first_row_id=*/4,
/*f0_values=*/{100, 101, 102, 103}));
// the index selects one row of the first group and two of the second, but the first split
// still reports four rows. Pruning on that count would drop the second group and return a
// single row for a limit of 2, so the push down has to be skipped.
ASSERT_OK_AND_ASSIGN(LimitScanResult limited,
ScanAndReadWithLimit(table_path, {"f0"}, /*limit=*/2,
/*predicate=*/nullptr,
/*row_ranges=*/{Range(3, 3), Range(5, 6)}));
ASSERT_EQ(limited.splits.size(), 2);
ASSERT_OK_AND_ASSIGN(std::vector<int32_t> values, CollectF0Values(limited.rows));
ASSERT_EQ(values, (std::vector<int32_t>{3, 101, 102}));
}
std::vector<DataEvolutionTableParam> GetTestValuesForDataEvolutionTableTest() {
std::vector<DataEvolutionTableParam> values;
for (bool enable_snapshot_live_manifest_cache : {false, true}) {
values.emplace_back("parquet", enable_snapshot_live_manifest_cache, /*bitmap64=*/false);
#ifdef PAIMON_ENABLE_MOSAIC
values.emplace_back("mosaic", enable_snapshot_live_manifest_cache, /*bitmap64=*/false);
#endif
#ifdef PAIMON_ENABLE_LANCE
values.emplace_back("lance", enable_snapshot_live_manifest_cache, /*bitmap64=*/false);
#endif
#ifdef PAIMON_ENABLE_ORC
values.emplace_back("orc", enable_snapshot_live_manifest_cache, /*bitmap64=*/false);
#endif
#ifdef PAIMON_ENABLE_AVRO
values.emplace_back("avro", enable_snapshot_live_manifest_cache, /*bitmap64=*/false);
#endif
}
return values;
}
std::vector<DataEvolutionTableParam> GetTestValuesForDataEvolutionDvTableTest() {
std::vector<DataEvolutionTableParam> values;
for (const auto& param : GetTestValuesForDataEvolutionTableTest()) {
for (const auto& bitmap64 : {false, true}) {
values.emplace_back(std::get<0>(param), std::get<1>(param), bitmap64);
}
}
return values;
}
INSTANTIATE_TEST_SUITE_P(FileFormat, DataEvolutionDvTableTest,
::testing::ValuesIn(GetTestValuesForDataEvolutionDvTableTest()));
std::vector<DataEvolutionTableParam> GetTestValuesForDedicatedVectorTableTest() {
std::vector<std::string> formats = {"parquet"};
#ifdef PAIMON_ENABLE_LANCE
formats.push_back("lance");
#endif
#ifdef PAIMON_ENABLE_VORTEX
formats.push_back("vortex");
#endif
std::vector<DataEvolutionTableParam> values;
for (const auto& format : formats) {
for (bool enable_cache : {false, true}) {
values.emplace_back(format, enable_cache, /*bitmap64=*/false);
}
}
return values;
}
INSTANTIATE_TEST_SUITE_P(VectorFormat, DedicatedVectorTableTest,
::testing::ValuesIn(GetTestValuesForDedicatedVectorTableTest()));
INSTANTIATE_TEST_SUITE_P(FileFormat, DataEvolutionTableTest,
::testing::ValuesIn(GetTestValuesForDataEvolutionTableTest()));
} // namespace paimon::test