diff --git a/.asf.yaml b/.asf.yaml index 0cc39ed4..e4846c16 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -54,39 +54,7 @@ github: dismiss_stale_reviews: true require_last_push_approval: true required_approving_review_count: 1 - required_status_checks: - - name: "pre-commit" - app_slug: -1 - - name: "rat-license-check" - app_slug: -1 - - name: "script-tests" - app_slug: -1 - - name: "asan-ubsan-x86_64" - app_slug: -1 - - name: "tsan-x86_64" - app_slug: -1 - - name: "clang-debug-x86_64" - app_slug: -1 - - name: "clang-release-x86_64" - app_slug: -1 - - name: "gcc-debug-x86_64" - app_slug: -1 - - name: "gcc-release-x86_64" - app_slug: -1 - - name: "gcc-debug-aarch64" - app_slug: -1 - - name: "gcc-release-aarch64" - app_slug: -1 - - name: "clang-debug-aarch64" - app_slug: -1 - - name: "asan-ubsan-aarch64" - app_slug: -1 - - name: "clang-release-aarch64" - app_slug: -1 - - name: "tsan-aarch64" - app_slug: -1 - - name: "gcc8-test" - app_slug: -1 + pull_requests: allow_auto_merge: false allow_update_branch: true diff --git a/cmake_modules/arrow.diff b/cmake_modules/arrow.diff index f86f36e8..b8b83517 100644 --- a/cmake_modules/arrow.diff +++ b/cmake_modules/arrow.diff @@ -1,22 +1,338 @@ -diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc -index ec3890a41f..943f69bb6c 100644 ---- a/cpp/src/parquet/arrow/schema.cc -+++ b/cpp/src/parquet/arrow/schema.cc -@@ -178,7 +178,7 @@ static Status GetTimestampMetadata(const ::arrow::TimestampType& type, +diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.cmake +index e7523add27..e079a1ad41 100644 +--- a/cpp/cmake_modules/BuildUtils.cmake ++++ b/cpp/cmake_modules/BuildUtils.cmake +@@ -112,7 +112,7 @@ function(arrow_create_merged_static_lib output_target) + execute_process(COMMAND ${LIBTOOL_MACOS} -V + OUTPUT_VARIABLE LIBTOOL_V_OUTPUT + OUTPUT_STRIP_TRAILING_WHITESPACE) +- if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools-([0-9.]+).*") ++ if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools(_ld)?-([0-9.]+).*") + message(FATAL_ERROR "libtool found appears to be the incompatible GNU libtool: ${LIBTOOL_MACOS}" + ) + endif() +diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake +index 8cb3ec83f5..0765df8fa8 100644 +--- a/cpp/cmake_modules/ThirdpartyToolchain.cmake ++++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake +@@ -983,6 +983,11 @@ if(CMAKE_TOOLCHAIN_FILE) + list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}) + endif() - // The user is explicitly asking for Impala int96 encoding, there is no - // logical type. -- if (arrow_properties.support_deprecated_int96_timestamps()) { -+ if (arrow_properties.support_deprecated_int96_timestamps() && target_unit == ::arrow::TimeUnit::NANO) { - *physical_type = ParquetType::INT96; - return Status::OK(); - } ++# Compatibility with bundled dependencies that require old CMake versions. ++if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.30") ++ list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_POLICY_VERSION_MINIMUM=3.5) ++endif() ++ + # and crosscompiling emulator (for try_run() ) + if(CMAKE_CROSSCOMPILING_EMULATOR) + string(REPLACE ";" ${EP_LIST_SEPARATOR} EP_CMAKE_CROSSCOMPILING_EMULATOR +@@ -1716,6 +1721,7 @@ macro(build_thrift) + -DWITH_JAVASCRIPT=OFF + -DWITH_LIBEVENT=OFF + -DWITH_NODEJS=OFF ++ -DWITH_OPENSSL=OFF + -DWITH_PYTHON=OFF + -DWITH_QT5=OFF + -DWITH_ZLIB=OFF) +diff --git a/cpp/src/arrow/io/interfaces.h b/cpp/src/arrow/io/interfaces.h +index b36c38c6d4..f974a33073 100644 +--- a/cpp/src/arrow/io/interfaces.h ++++ b/cpp/src/arrow/io/interfaces.h +@@ -210,7 +210,7 @@ class ARROW_EXPORT InputStream : virtual public FileInterface, virtual public Re + /// \brief Advance or skip stream indicated number of bytes + /// \param[in] nbytes the number to move forward + /// \return Status +- Status Advance(int64_t nbytes); ++ virtual Status Advance(int64_t nbytes); + /// \brief Return zero-copy string_view to upcoming bytes. + /// +diff --git a/cpp/src/arrow/util/bit_run_reader.h b/cpp/src/arrow/util/bit_run_reader.h +index a436a503a0..27d483978c 100644 +--- a/cpp/src/arrow/util/bit_run_reader.h ++++ b/cpp/src/arrow/util/bit_run_reader.h +@@ -168,6 +168,26 @@ class ARROW_EXPORT BitRunReader { + using BitRunReader = BitRunReaderLinear; + #endif + ++template ++inline Status VisitBitRuns(const uint8_t* bitmap, int64_t offset, int64_t length, ++ Visit&& visit) { ++ if (bitmap == NULLPTR) { ++ // Assuming all set (as in a null bitmap) ++ return visit(static_cast(0), length, true); ++ } ++ BitRunReader reader(bitmap, offset, length); ++ int64_t position = 0; ++ while (true) { ++ const auto run = reader.NextRun(); ++ if (run.length == 0) { ++ break; ++ } ++ ARROW_RETURN_NOT_OK(visit(position, run.length, run.set)); ++ position += run.length; ++ } ++ return Status::OK(); ++} ++ + struct SetBitRun { + int64_t position; + int64_t length; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc -index 285e2a5973..aa6f92f077 100644 +index 285e2a5973..52f42cf5b3 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc -@@ -1013,25 +1013,32 @@ Status FileReaderImpl::GetRecordBatchReader(const std::vector& row_groups, +@@ -19,12 +19,14 @@ + + #include + #include ++#include + #include + #include + #include + #include + + #include "arrow/array.h" ++#include "arrow/array/concatenate.h" + #include "arrow/buffer.h" + #include "arrow/extension_type.h" + #include "arrow/io/memory.h" +@@ -32,12 +34,14 @@ + #include "arrow/table.h" + #include "arrow/type.h" + #include "arrow/util/async_generator.h" ++#include "arrow/util/bit_run_reader.h" + #include "arrow/util/bit_util.h" + #include "arrow/util/future.h" + #include "arrow/util/iterator.h" + #include "arrow/util/logging.h" + #include "arrow/util/parallel.h" + #include "arrow/util/range.h" ++#include "arrow/util/span.h" + #include "arrow/util/tracing_internal.h" + #include "parquet/arrow/reader_internal.h" + #include "parquet/column_reader.h" +@@ -254,6 +254,11 @@ class FileReaderImpl : public FileReader { + return GetColumn(i, AllRowGroupsFactory(), out); + } + ++ ::arrow::Status GetColumn( ++ int i, const std::shared_ptr>& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* out) override; ++ + Status GetSchema(std::shared_ptr<::arrow::Schema>* out) override { + return FromParquetSchema(reader_->metadata()->schema(), reader_properties_, + reader_->metadata()->key_value_metadata(), out); +@@ -493,10 +498,40 @@ class LeafReader : public ColumnReaderImpl { + + ::arrow::Status BuildArray(int64_t length_upper_bound, + std::shared_ptr<::arrow::ChunkedArray>* out) final { ++ if (!out_) { ++ BEGIN_PARQUET_CATCH_EXCEPTIONS ++ RETURN_NOT_OK( ++ TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_)); ++ END_PARQUET_CATCH_EXCEPTIONS ++ } + *out = out_; + return Status::OK(); + } + ++ std::vector LeafColumnIndices() const final { ++ return {input_->column_index()}; ++ } ++ ++ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { ++ if (col_idx != input_->column_index()) return Status::OK(); ++ BEGIN_PARQUET_CATCH_EXCEPTIONS ++ out_ = nullptr; ++ record_reader_->Reset(); ++ record_reader_->Reserve(reserve); ++ return Status::OK(); ++ END_PARQUET_CATCH_EXCEPTIONS ++ } ++ ++ int64_t SkipRecords(int col_idx, int64_t num_records) final { ++ if (col_idx != input_->column_index() || num_records <= 0) return 0; ++ return record_reader_->SkipRecords(num_records); ++ } ++ ++ int64_t ReadRecords(int col_idx, int64_t num_records) final { ++ if (col_idx != input_->column_index() || num_records <= 0) return 0; ++ return record_reader_->ReadRecords(num_records); ++ } ++ + const std::shared_ptr field() override { return field_; } + + private: +@@ -532,6 +567,22 @@ class ExtensionReader : public ColumnReaderImpl { + return storage_reader_->LoadBatch(number_of_records); + } + ++ std::vector LeafColumnIndices() const final { ++ return storage_reader_->LeafColumnIndices(); ++ } ++ ++ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { ++ return storage_reader_->ResetLeaf(col_idx, reserve); ++ } ++ ++ int64_t SkipRecords(int col_idx, int64_t num_records) final { ++ return storage_reader_->SkipRecords(col_idx, num_records); ++ } ++ ++ int64_t ReadRecords(int col_idx, int64_t num_records) final { ++ return storage_reader_->ReadRecords(col_idx, num_records); ++ } ++ + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override { + std::shared_ptr storage; +@@ -576,6 +627,22 @@ class ListReader : public ColumnReaderImpl { + return item_reader_->LoadBatch(number_of_records); + } + ++ std::vector LeafColumnIndices() const final { ++ return item_reader_->LeafColumnIndices(); ++ } ++ ++ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { ++ return item_reader_->ResetLeaf(col_idx, reserve); ++ } ++ ++ int64_t SkipRecords(int col_idx, int64_t num_records) final { ++ return item_reader_->SkipRecords(col_idx, num_records); ++ } ++ ++ int64_t ReadRecords(int col_idx, int64_t num_records) final { ++ return item_reader_->ReadRecords(col_idx, num_records); ++ } ++ + virtual ::arrow::Result> AssembleArray( + std::shared_ptr data) { + if (field_->type()->id() == ::arrow::Type::MAP) { +@@ -642,8 +713,10 @@ class ListReader : public ColumnReaderImpl { + + const std::shared_ptr field() override { return field_; } + +- private: ++ protected: + std::shared_ptr ctx_; ++ ++ private: + std::shared_ptr field_; + ::parquet::internal::LevelInfo level_info_; + std::unique_ptr item_reader_; +@@ -662,12 +735,62 @@ class PARQUET_NO_EXPORT FixedSizeListReader : public ListReader { + DCHECK_EQ(field()->type()->id(), ::arrow::Type::FIXED_SIZE_LIST); + const auto& type = checked_cast<::arrow::FixedSizeListType&>(*field()->type()); + const int32_t* offsets = reinterpret_cast(data->buffers[1]->data()); +- for (int x = 1; x <= data->length; x++) { +- int32_t size = offsets[x] - offsets[x - 1]; +- if (size != type.list_size()) { +- return Status::Invalid("Expected all lists to be of size=", type.list_size(), +- " but index ", x, " had size=", size); ++ const int32_t list_size = type.list_size(); ++ auto validate_offsets = [&](int64_t start, int64_t length, ++ bool has_elements) -> Status { ++ const int32_t expected_size = has_elements ? list_size : 0; ++ ::arrow::util::span run_offsets( ++ offsets + start, static_cast(length + 1)); ++ const auto first_invalid_offset = std::adjacent_find( ++ run_offsets.begin(), run_offsets.end(), ++ [&](int32_t left, int32_t right) { return right - left != expected_size; }); ++ if (first_invalid_offset != run_offsets.end()) { ++ const int64_t x = ++ start + std::distance(run_offsets.begin(), first_invalid_offset); ++ const int32_t size = offsets[x + 1] - offsets[x]; ++ if (has_elements) { ++ return Status::Invalid("Expected all lists to be of size=", list_size, ++ " but index ", x + 1, " had size=", size); ++ } ++ return Status::Invalid("Expected null fixed-size list at index ", x + 1, ++ " to have no child values but had size=", size); + } ++ return Status::OK(); ++ }; ++ if (data->GetNullCount() != 0) { ++ // Rebuild the child array run-by-run so null fixed-size list slots still ++ // contribute list_size child values in the final layout. ++ ::arrow::ArrayVector child_arrays; ++ ++ auto visit_run = [&](int64_t start, int64_t length, bool has_elements) -> Status { ++ RETURN_NOT_OK(validate_offsets(start, length, has_elements)); ++ ++ const int64_t child_length = length * list_size; ++ // Valid runs reuse the decoded child slice; null runs materialize null ++ // children to preserve the fixed-size list shape. ++ if (!has_elements) { ++ ARROW_ASSIGN_OR_RAISE( ++ auto null_array, ++ ::arrow::MakeArrayOfNull(type.value_type(), child_length, ctx_->pool)); ++ child_arrays.push_back(std::move(null_array)); ++ return Status::OK(); ++ } ++ child_arrays.push_back( ++ ::arrow::MakeArray(data->child_data[0]->Slice(offsets[start], child_length))); ++ return Status::OK(); ++ }; ++ ++ DCHECK_NE(data->buffers[0], nullptr); ++ RETURN_NOT_OK(::arrow::internal::VisitBitRuns( ++ data->buffers[0]->data(), data->offset, data->length, visit_run)); ++ ++ // TODO(GH-50271): Build one padded child array directly instead of creating ++ // one temporary Array/ArrayData per validity run and concatenating them. ++ ARROW_ASSIGN_OR_RAISE(auto child_array_with_padding, ++ ::arrow::Concatenate(child_arrays, ctx_->pool)); ++ data->child_data[0] = child_array_with_padding->data(); ++ } else { ++ RETURN_NOT_OK(validate_offsets(/*start=*/0, data->length, /*valid=*/true)); + } + data->buffers.resize(1); + std::shared_ptr result = ::arrow::MakeArray(data); +@@ -709,6 +832,39 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { + } + return Status::OK(); + } ++ ++ std::vector LeafColumnIndices() const override { ++ std::vector indices; ++ for (const std::unique_ptr& reader : children_) { ++ std::vector child_indices = reader->LeafColumnIndices(); ++ indices.insert(indices.end(), child_indices.begin(), child_indices.end()); ++ } ++ return indices; ++ } ++ ++ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) override { ++ for (const std::unique_ptr& reader : children_) { ++ RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve)); ++ } ++ return Status::OK(); ++ } ++ ++ int64_t SkipRecords(int col_idx, int64_t num_records) override { ++ int64_t skipped = 0; ++ for (const std::unique_ptr& reader : children_) { ++ skipped += reader->SkipRecords(col_idx, num_records); ++ } ++ return skipped; ++ } ++ ++ int64_t ReadRecords(int col_idx, int64_t num_records) override { ++ int64_t read = 0; ++ for (const std::unique_ptr& reader : children_) { ++ read += reader->ReadRecords(col_idx, num_records); ++ } ++ return read; ++ } ++ + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override; + Status GetDefLevels(const int16_t** data, int64_t* length) override; +@@ -1013,25 +1169,32 @@ Status FileReaderImpl::GetRecordBatchReader(const std::vector& row_groups, return Status::OK(); } @@ -55,30 +371,238 @@ index 285e2a5973..aa6f92f077 100644 RETURN_NOT_OK(::arrow::internal::OptionalParallelFor( reader_properties_.use_threads(), static_cast(readers.size()), -diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc -index 4fd7ef1b47..87326a54f1 100644 ---- a/cpp/src/parquet/arrow/writer.cc -+++ b/cpp/src/parquet/arrow/writer.cc -@@ -314,6 +314,14 @@ class FileWriterImpl : public FileWriter { - return Status::OK(); - } +@@ -1224,6 +1387,23 @@ Status FileReaderImpl::GetColumn(int i, FileColumnIteratorFactory iterator_facto + return Status::OK(); + } -+ int64_t GetBufferedSize() override { -+ if (row_group_writer_ == nullptr) { -+ return 0; -+ } -+ return row_group_writer_->total_compressed_bytes() + -+ row_group_writer_->total_compressed_bytes_written(); -+ } ++::arrow::Status FileReaderImpl::GetColumn( ++ int i, const std::shared_ptr>& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* out) { ++ RETURN_NOT_OK(BoundsCheckColumn(i)); ++ auto ctx = std::make_shared(); ++ ctx->reader = reader_.get(); ++ ctx->pool = pool_; ++ ctx->iterator_factory = iterator_factory; ++ ctx->filter_leaves = true; ++ ctx->included_leaves = column_indices; ++ std::unique_ptr result; ++ RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); ++ *out = std::move(result); ++ return Status::OK(); ++} + - Status Close() override { - if (!closed_) { - // Make idempotent -@@ -418,10 +426,13 @@ class FileWriterImpl : public FileWriter { + Status FileReaderImpl::ReadRowGroups(const std::vector& row_groups, + const std::vector& column_indices, + std::shared_ptr* out) { +diff --git a/cpp/src/parquet/arrow/reader.h b/cpp/src/parquet/arrow/reader.h +index 6e46ca43f7..e86ff0ef52 100644 +--- a/cpp/src/parquet/arrow/reader.h ++++ b/cpp/src/parquet/arrow/reader.h +@@ -21,6 +21,7 @@ + // N.B. we don't include async_generator.h as it's relatively heavy + #include + #include ++#include + #include - // Max number of rows allowed in a row group. - const int64_t max_row_group_length = this->properties().max_row_group_length(); -+ const int64_t max_row_group_size = this->properties().max_row_group_size(); + #include "parquet/file_reader.h" +@@ -48,9 +49,13 @@ namespace arrow { + + class ColumnChunkReader; + class ColumnReader; ++class FileColumnIterator; + struct SchemaManifest; + class RowGroupReader; + ++using FileColumnIteratorFactory = ++ std::function; ++ + /// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches. + /// + /// This interfaces caters for different use cases and thus provides different +@@ -136,6 +141,27 @@ class PARQUET_EXPORT FileReader { + // The indicated column index is relative to the schema + virtual ::arrow::Status GetColumn(int i, std::unique_ptr* out) = 0; + ++ /// \brief Return a ColumnReader with a custom FileColumnIteratorFactory ++ /// and leaf column filtering. ++ /// ++ /// This allows callers to customize page reading behavior (e.g., setting ++ /// data_page_filter for page-level skipping) and to select only specific ++ /// leaf columns within a nested field. The factory is called once per leaf ++ /// column included in column_indices. ++ /// ++ /// \param i top-level field index (same as GetColumn(int i, ...)) ++ /// \param column_indices leaf column indices to include (enables sub-column ++ /// projection within nested types) ++ /// \param iterator_factory factory to create FileColumnIterator per leaf ++ /// \param[out] out the ColumnReader (may be nullptr if all leaves are pruned) ++ virtual ::arrow::Status GetColumn( ++ int i, const std::shared_ptr>& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* out) { ++ return ::arrow::Status::NotImplemented( ++ "GetColumn with factory not implemented"); ++ } ++ + /// \brief Return arrow schema for all the columns. + virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0; + +@@ -316,6 +342,43 @@ class PARQUET_EXPORT ColumnReader { + // the data available in the file. + virtual ::arrow::Status NextBatch(int64_t batch_size, + std::shared_ptr<::arrow::ChunkedArray>* out) = 0; ++ ++ /// \brief Leaf column indices covered by this (sub)tree, in leaf order. ++ /// ++ /// Used to drive per-leaf row filtering: after page-level skipping each leaf ++ /// lives in its own compressed coordinate space, so callers must reset and ++ /// skip/read each leaf independently rather than in lockstep. ++ virtual std::vector LeafColumnIndices() const { return {}; } ++ ++ /// \brief Reset the leaf identified by col_idx and reserve space for ++ /// `reserve` records (in that leaf's post-page-filter compressed space). ++ /// Must be called before SkipRecords()/ReadRecords() for that leaf, and ++ /// followed by BuildArray() to get the result. ++ virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) { ++ return ::arrow::Status::NotImplemented("ResetLeaf not implemented"); ++ } ++ ++ /// \brief Skip num_records on the leaf identified by col_idx and return the ++ /// number of records actually skipped. Returns 0 when num_records <= 0 or ++ /// col_idx does not belong to this (sub)tree. May throw ParquetException on a ++ /// decode error; callers convert it to Status at the public boundary. ++ virtual int64_t SkipRecords(int col_idx, int64_t num_records) { return 0; } ++ ++ /// \brief Read num_records on the leaf identified by col_idx and return the ++ /// number of records actually read. Values accumulate across successive calls ++ /// until BuildArray() is called. Returns 0 when num_records <= 0 or col_idx ++ /// does not belong to this (sub)tree. May throw ParquetException on a decode ++ /// error; callers convert it to Status at the public boundary. ++ virtual int64_t ReadRecords(int col_idx, int64_t num_records) { return 0; } ++ ++ /// \brief Build the Arrow array from previously loaded data. ++ /// For leaf readers, calls TransferColumnData if not already done. ++ /// For nested readers, assembles the nested array from child arrays. ++ virtual ::arrow::Status BuildArray( ++ int64_t length_upper_bound, ++ std::shared_ptr<::arrow::ChunkedArray>* out) { ++ return ::arrow::Status::NotImplemented("BuildArray not implemented"); ++ } + }; + + /// \brief Experimental helper class for bindings (like Python) that struggle +diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h +index cf9dbb8657..9216f18289 100644 +--- a/cpp/src/parquet/arrow/reader_internal.h ++++ b/cpp/src/parquet/arrow/reader_internal.h +@@ -26,6 +26,7 @@ + #include + #include + ++#include "parquet/arrow/reader.h" + #include "parquet/arrow/schema.h" + #include "parquet/column_reader.h" + #include "parquet/file_reader.h" +@@ -70,7 +71,10 @@ class FileColumnIterator { + + virtual ~FileColumnIterator() {} + +- std::unique_ptr<::parquet::PageReader> NextChunk() { ++ /// \brief Fetch the PageReader for the next row group in this iterator's ++ /// range. Virtual so subclasses can decorate the returned PageReader, e.g. ++ /// to install a data_page_filter for I/O-level page skipping. ++ virtual std::unique_ptr<::parquet::PageReader> NextChunk() { + if (row_groups_.empty()) { + return nullptr; + } +@@ -95,9 +99,6 @@ class FileColumnIterator { + std::deque row_groups_; + }; + +-using FileColumnIteratorFactory = +- std::function; +- + Status TransferColumnData(::parquet::internal::RecordReader* reader, + const std::shared_ptr<::arrow::Field>& value_field, + const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, +diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc +index ec3890a41f..943f69bb6c 100644 +--- a/cpp/src/parquet/arrow/schema.cc ++++ b/cpp/src/parquet/arrow/schema.cc +@@ -178,7 +178,7 @@ static Status GetTimestampMetadata(const ::arrow::TimestampType& type, + + // The user is explicitly asking for Impala int96 encoding, there is no + // logical type. +- if (arrow_properties.support_deprecated_int96_timestamps()) { ++ if (arrow_properties.support_deprecated_int96_timestamps() && target_unit == ::arrow::TimeUnit::NANO) { + *physical_type = ParquetType::INT96; + return Status::OK(); + } +diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc +index 4fd7ef1b47..feff99c99b 100644 +--- a/cpp/src/parquet/arrow/writer.cc ++++ b/cpp/src/parquet/arrow/writer.cc +@@ -26,6 +26,7 @@ + #include + + #include "arrow/array.h" ++#include "arrow/array/concatenate.h" + #include "arrow/extension_type.h" + #include "arrow/ipc/writer.h" + #include "arrow/record_batch.h" +@@ -142,13 +143,24 @@ class ArrowColumnWriterV2 { + leaf_idx, ctx, [&](const MultipathLevelBuilderResult& result) { + size_t visited_component_size = result.post_list_visited_elements.size(); + DCHECK_GT(visited_component_size, 0); +- if (visited_component_size != 1) { +- return Status::NotImplemented( +- "Lists with non-zero length null components are not supported"); ++ std::shared_ptr values_array; ++ if (visited_component_size == 1) { ++ const ElementRange& range = result.post_list_visited_elements[0]; ++ values_array = result.leaf_array->Slice(range.start, range.Size()); ++ } else { ++ // Multiple leaf ranges can be produced when child values are ++ // skipped, such as null fixed-size-list slots, or when ++ // list-view ranges are non-contiguous. Concatenate the slices ++ // in logical write order. ++ ::arrow::ArrayVector arrays; ++ arrays.reserve(visited_component_size); ++ for (const auto& range : result.post_list_visited_elements) { ++ DCHECK(!range.Empty()); ++ arrays.push_back(result.leaf_array->Slice(range.start, range.Size())); ++ } ++ ARROW_ASSIGN_OR_RAISE(values_array, ++ ::arrow::Concatenate(arrays, ctx->memory_pool)); + } +- const ElementRange& range = result.post_list_visited_elements[0]; +- std::shared_ptr values_array = +- result.leaf_array->Slice(range.start, range.Size()); + + return column_writer->WriteArrow(result.def_levels, result.rep_levels, + result.def_rep_level_count, *values_array, +@@ -314,6 +326,14 @@ class FileWriterImpl : public FileWriter { + return Status::OK(); + } + ++ int64_t GetBufferedSize() override { ++ if (row_group_writer_ == nullptr) { ++ return 0; ++ } ++ return row_group_writer_->total_compressed_bytes() + ++ row_group_writer_->total_compressed_bytes_written(); ++ } ++ + Status Close() override { + if (!closed_) { + // Make idempotent +@@ -418,10 +438,13 @@ class FileWriterImpl : public FileWriter { + + // Max number of rows allowed in a row group. + const int64_t max_row_group_length = this->properties().max_row_group_length(); ++ const int64_t max_row_group_size = this->properties().max_row_group_size(); // Initialize a new buffered row group writer if necessary. if (row_group_writer_ == nullptr || !row_group_writer_->buffered() || @@ -103,124 +627,223 @@ index 4a1a033a7b..0f13d05e44 100644 /// \brief Write the footer and close the file. virtual ::arrow::Status Close() = 0; virtual ~FileWriter(); -diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h -index 4d3acb491e..3906ff3c59 100644 ---- a/cpp/src/parquet/properties.h -+++ b/cpp/src/parquet/properties.h -@@ -139,6 +139,7 @@ static constexpr bool DEFAULT_IS_DICTIONARY_ENABLED = true; - static constexpr int64_t DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT = kDefaultDataPageSize; - static constexpr int64_t DEFAULT_WRITE_BATCH_SIZE = 1024; - static constexpr int64_t DEFAULT_MAX_ROW_GROUP_LENGTH = 1024 * 1024; -+static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = 128 * 1024 * 1024; - static constexpr bool DEFAULT_ARE_STATISTICS_ENABLED = true; - static constexpr int64_t DEFAULT_MAX_STATISTICS_SIZE = 4096; - static constexpr Encoding::type DEFAULT_ENCODING = Encoding::UNKNOWN; -@@ -232,6 +233,7 @@ class PARQUET_EXPORT WriterProperties { - dictionary_pagesize_limit_(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT), - write_batch_size_(DEFAULT_WRITE_BATCH_SIZE), - max_row_group_length_(DEFAULT_MAX_ROW_GROUP_LENGTH), -+ max_row_group_size_(DEFAULT_MAX_ROW_GROUP_SIZE), - pagesize_(kDefaultDataPageSize), - version_(ParquetVersion::PARQUET_2_6), - data_page_version_(ParquetDataPageVersion::V1), -@@ -244,6 +246,7 @@ class PARQUET_EXPORT WriterProperties { - dictionary_pagesize_limit_(properties.dictionary_pagesize_limit()), - write_batch_size_(properties.write_batch_size()), - max_row_group_length_(properties.max_row_group_length()), -+ max_row_group_size_(properties.max_row_group_size()), - pagesize_(properties.data_pagesize()), - version_(properties.version()), - data_page_version_(properties.data_page_version()), -@@ -321,6 +324,13 @@ class PARQUET_EXPORT WriterProperties { - return this; - } +diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc +index ebf9515f27..0abc7d2320 100644 +--- a/cpp/src/parquet/column_reader.cc ++++ b/cpp/src/parquet/column_reader.cc +@@ -208,6 +208,39 @@ ReaderProperties default_reader_properties() { + return default_reader_properties; + } -+ /// Specify the max bytes size to put in a single row group. -+ /// Default 128 M. -+ Builder* max_row_group_size(int64_t max_row_group_size) { -+ max_row_group_size_ = max_row_group_size; -+ return this; ++void PageReader::set_data_page_read_plan( ++ int64_t first_data_page_offset, ++ std::vector data_pages) { ++ if (data_page_filter_) { ++ throw ParquetException( ++ "data_page_filter and data_page_read_plan cannot be enabled together"); ++ } ++ if (first_data_page_offset < 0) { ++ throw ParquetException("Invalid negative first data page offset"); ++ } ++ ++ int64_t previous_end = first_data_page_offset; ++ int32_t previous_ordinal = -1; ++ for (const auto& page : data_pages) { ++ int64_t page_end; ++ if (page.page_ordinal < 0 || page.offset < first_data_page_offset || ++ page.compressed_page_size <= 0 || ++ AddWithOverflow(page.offset, page.compressed_page_size, &page_end)) { ++ throw ParquetException("Invalid data page read plan entry"); ++ } ++ if (page.offset < previous_end || page.page_ordinal <= previous_ordinal) { ++ throw ParquetException("Data page read plan entries must be ordered"); + } ++ previous_end = page_end; ++ previous_ordinal = page.page_ordinal; ++ } + - /// Specify the data page size. - /// Default 1MB. - Builder* data_pagesize(int64_t pg_size) { -@@ -664,7 +674,7 @@ class PARQUET_EXPORT WriterProperties { - - return std::shared_ptr(new WriterProperties( - pool_, dictionary_pagesize_limit_, write_batch_size_, max_row_group_length_, -- pagesize_, version_, created_by_, page_checksum_enabled_, -+ max_row_group_size_, pagesize_, version_, created_by_, page_checksum_enabled_, - std::move(file_encryption_properties_), default_column_properties_, - column_properties, data_page_version_, store_decimal_as_integer_, - std::move(sorting_columns_))); -@@ -675,6 +685,7 @@ class PARQUET_EXPORT WriterProperties { - int64_t dictionary_pagesize_limit_; - int64_t write_batch_size_; - int64_t max_row_group_length_; -+ int64_t max_row_group_size_; - int64_t pagesize_; - ParquetVersion::type version_; - ParquetDataPageVersion data_page_version_; -@@ -705,6 +716,8 @@ class PARQUET_EXPORT WriterProperties { ++ data_page_read_plan_enabled_ = true; ++ first_data_page_offset_ = first_data_page_offset; ++ data_page_read_plan_ = std::move(data_pages); ++ next_data_page_ = 0; ++} ++ + namespace { - inline int64_t max_row_group_length() const { return max_row_group_length_; } + // Extracts encoded statistics from V1 and V2 data page headers +@@ -430,9 +463,43 @@ std::shared_ptr SerializedPageReader::NextPage() { -+ inline int64_t max_row_group_size() const { return max_row_group_size_; } + // Loop here because there may be unhandled page types that we skip until + // finding a page that we do know what to do with +- while (seen_num_values_ < total_num_values_) { ++ while (data_page_read_plan_enabled_ || seen_num_values_ < total_num_values_) { ++ const DataPageReadPlanEntry* planned_data_page = nullptr; ++ uint32_t page_header_limit = max_page_header_size_; + - inline int64_t data_pagesize() const { return pagesize_; } ++ if (data_page_read_plan_enabled_) { ++ if (next_data_page_ >= data_page_read_plan_.size()) { ++ return nullptr; ++ } ++ ++ PARQUET_ASSIGN_OR_THROW(int64_t current_position, stream_->Tell()); ++ if (current_position < first_data_page_offset_) { ++ page_header_limit = static_cast(std::min( ++ page_header_limit, first_data_page_offset_ - current_position)); ++ } else { ++ planned_data_page = &data_page_read_plan_[next_data_page_]; ++ if (current_position > planned_data_page->offset) { ++ throw ParquetException("Data page read plan points behind stream position"); ++ } ++ PARQUET_THROW_NOT_OK( ++ stream_->Advance(planned_data_page->offset - current_position)); ++ PARQUET_ASSIGN_OR_THROW(int64_t target_position, stream_->Tell()); ++ if (target_position != planned_data_page->offset) { ++ throw ParquetException("Failed to seek to planned data page"); ++ } ++ page_ordinal_ = planned_data_page->page_ordinal; ++ page_header_limit = static_cast(std::min( ++ page_header_limit, planned_data_page->compressed_page_size)); ++ } ++ } ++ ++ if (page_header_limit == 0) { ++ throw ParquetException("No bytes available for page header"); ++ } ++ + uint32_t header_size = 0; +- uint32_t allowed_page_size = kDefaultPageHeaderSize; ++ uint32_t allowed_page_size = ++ std::min(kDefaultPageHeaderSize, page_header_limit); - inline ParquetDataPageVersion data_page_version() const { -@@ -810,7 +823,7 @@ class PARQUET_EXPORT WriterProperties { - private: - explicit WriterProperties( - MemoryPool* pool, int64_t dictionary_pagesize_limit, int64_t write_batch_size, -- int64_t max_row_group_length, int64_t pagesize, ParquetVersion::type version, -+ int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, ParquetVersion::type version, - const std::string& created_by, bool page_write_checksum_enabled, - std::shared_ptr file_encryption_properties, - const ColumnProperties& default_column_properties, -@@ -821,6 +834,7 @@ class PARQUET_EXPORT WriterProperties { - dictionary_pagesize_limit_(dictionary_pagesize_limit), - write_batch_size_(write_batch_size), - max_row_group_length_(max_row_group_length), -+ max_row_group_size_(max_row_group_size), - pagesize_(pagesize), - parquet_data_page_version_(data_page_version), - parquet_version_(version), -@@ -836,6 +850,7 @@ class PARQUET_EXPORT WriterProperties { - int64_t dictionary_pagesize_limit_; - int64_t write_batch_size_; - int64_t max_row_group_length_; -+ int64_t max_row_group_size_; - int64_t pagesize_; - ParquetDataPageVersion parquet_data_page_version_; - ParquetVersion::type parquet_version_; + // Page headers can be very large because of page statistics + // We try to deserialize a larger buffer progressively +@@ -458,11 +525,12 @@ std::shared_ptr SerializedPageReader::NextPage() { + // Failed to deserialize. Double the allowed page header size and try again + std::stringstream ss; + ss << e.what(); +- allowed_page_size *= 2; +- if (allowed_page_size > max_page_header_size_) { ++ if (allowed_page_size >= page_header_limit) { + ss << "Deserializing page header failed.\n"; + throw ParquetException(ss.str()); + } ++ allowed_page_size = ++ std::min(allowed_page_size * 2, page_header_limit); + } + } + // Advance the stream offset +@@ -474,6 +542,20 @@ std::shared_ptr SerializedPageReader::NextPage() { + throw ParquetException("Invalid page header"); + } ---- a/cpp/src/parquet/file_reader.h -+++ b/cpp/src/parquet/file_reader.h -@@ -210,6 +210,17 @@ - ::arrow::Future<> WhenBuffered(const std::vector& row_groups, - const std::vector& column_indices) const; ++ const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); ++ if (planned_data_page != nullptr) { ++ if (page_type != PageType::DATA_PAGE && page_type != PageType::DATA_PAGE_V2) { ++ throw ParquetException("Data page read plan points to a non-data page"); ++ } ++ int64_t total_compressed_size; ++ if (AddWithOverflow(static_cast(header_size), ++ static_cast(compressed_len), ++ &total_compressed_size) || ++ total_compressed_size != planned_data_page->compressed_page_size) { ++ throw ParquetException("Planned data page size does not match page header"); ++ } ++ } ++ + EncodedStatistics data_page_statistics; + if (ShouldSkipPage(&data_page_statistics)) { + PARQUET_THROW_NOT_OK(stream_->Advance(compressed_len)); +@@ -494,8 +576,6 @@ std::shared_ptr SerializedPageReader::NextPage() { + ParquetException::EofException(ss.str()); + } -+ /// Pre-buffer arbitrary byte ranges (e.g., page-level ranges from OffsetIndex). -+ /// Unlike PreBuffer(), this does NOT set the column bitmap, so -+ /// GetColumnPageReader will use CachedInputStream (page-level cache path). -+ void PreBufferRanges(const std::vector<::arrow::io::ReadRange>& ranges, -+ const ::arrow::io::IOContext& ctx, -+ const ::arrow::io::CacheOptions& options); +- const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); +- + if (properties_.page_checksum_verification() && current_page_header_.__isset.crc && + PageCanUseChecksum(page_type)) { + // verify crc +@@ -534,6 +614,9 @@ std::shared_ptr SerializedPageReader::NextPage() { + LoadEnumSafe(&dict_header.encoding), + is_sorted); + } else if (page_type == PageType::DATA_PAGE) { ++ if (planned_data_page != nullptr) { ++ ++next_data_page_; ++ } + ++page_ordinal_; + const format::DataPageHeader& header = current_page_header_.data_page_header; + page_buffer = +@@ -545,6 +628,9 @@ std::shared_ptr SerializedPageReader::NextPage() { + LoadEnumSafe(&header.repetition_level_encoding), uncompressed_len, + std::move(data_page_statistics)); + } else if (page_type == PageType::DATA_PAGE_V2) { ++ if (planned_data_page != nullptr) { ++ ++next_data_page_; ++ } + ++page_ordinal_; + const format::DataPageHeaderV2& header = current_page_header_.data_page_header_v2; + +diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h +index 29e1b2a25e..386e574644 100644 +--- a/cpp/src/parquet/column_reader.h ++++ b/cpp/src/parquet/column_reader.h +@@ -76,6 +76,18 @@ struct PARQUET_EXPORT DataPageStats { + std::optional num_rows; + }; + ++/// \brief Identifies a data page that PageReader should read directly. ++/// ++/// The offset is relative to the beginning of the column chunk stream passed to ++/// PageReader::Open. The compressed size includes both the serialized page header ++/// and the compressed page body. The ordinal is the original data page ordinal in ++/// the column chunk and does not include the dictionary page. ++struct PARQUET_EXPORT DataPageReadPlanEntry { ++ int32_t page_ordinal; ++ int64_t offset; ++ int32_t compressed_page_size; ++}; + -+ /// Wait for arbitrary byte ranges to be pre-buffered. -+ ::arrow::Future<> WhenBufferedRanges( -+ const std::vector<::arrow::io::ReadRange>& ranges) const; + class PARQUET_EXPORT LevelDecoder { + public: + LevelDecoder(); +@@ -147,9 +159,21 @@ class PARQUET_EXPORT PageReader { + // ApplicationVersion::HasCorrectStatistics(). + // \note API EXPERIMENTAL + void set_data_page_filter(DataPageFilter data_page_filter) { ++ if (data_page_read_plan_enabled_) { ++ throw ParquetException( ++ "data_page_filter and data_page_read_plan cannot be enabled together"); ++ } + data_page_filter_ = std::move(data_page_filter); + } + ++ /// Configure PageReader to jump directly to selected data pages before reading ++ /// their headers. `first_data_page_offset` and each entry offset are relative to ++ /// the beginning of the column chunk stream. Dictionary pages before ++ /// `first_data_page_offset` are still read normally. ++ // \note API EXPERIMENTAL ++ void set_data_page_read_plan(int64_t first_data_page_offset, ++ std::vector data_pages); + - private: - // Holds a pointer to an instance of Contents implementation - std::unique_ptr contents_; + // @returns: shared_ptr(nullptr) on EOS, std::shared_ptr + // containing new Page otherwise + // +@@ -162,6 +186,11 @@ class PARQUET_EXPORT PageReader { + protected: + // Callback that decides if we should skip a page or not. + DataPageFilter data_page_filter_; ++ ++ bool data_page_read_plan_enabled_ = false; ++ int64_t first_data_page_offset_ = 0; ++ std::vector data_page_read_plan_; ++ size_t next_data_page_ = 0; + }; + class PARQUET_EXPORT ColumnReader { +diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc +index 3e9eeea6c6..671ebe4644 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc -@@ -207,6 +207,117 @@ +@@ -207,6 +207,117 @@ const RowGroupMetaData* RowGroupReader::metadata() const { return contents_->met return {col_start, col_length}; } @@ -338,7 +961,7 @@ index 4d3acb491e..3906ff3c59 100644 // RowGroupReader::Contents implementation for the Parquet file specification class SerializedRowGroup : public RowGroupReader::Contents { public: -@@ -242,6 +343,11 @@ +@@ -242,6 +353,11 @@ class SerializedRowGroup : public RowGroupReader::Contents { // segments. PARQUET_ASSIGN_OR_THROW(auto buffer, cached_source_->Read(col_range)); stream = std::make_shared<::arrow::io::BufferReader>(buffer); @@ -350,7 +973,7 @@ index 4d3acb491e..3906ff3c59 100644 } else { stream = properties_.GetStream(source_, col_range.offset, col_range.length); } -@@ -417,6 +523,26 @@ +@@ -417,6 +533,26 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } @@ -377,7 +1000,7 @@ index 4d3acb491e..3906ff3c59 100644 // Metadata/footer parsing. Divided up to separate sync/async paths, and to use // exceptions for error handling (with the async path converting to Future/Status). -@@ -911,6 +1037,22 @@ +@@ -911,6 +1047,22 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, return file->WhenBuffered(row_groups, column_indices); } @@ -400,553 +1023,118 @@ index 4d3acb491e..3906ff3c59 100644 // ---------------------------------------------------------------------- // File metadata helpers -diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake ---- a/cpp/cmake_modules/ThirdpartyToolchain.cmake -+++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake -@@ -981,6 +981,11 @@ if(CMAKE_TOOLCHAIN_FILE) - list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}) - endif() - -+# Compatibility with bundled dependencies that require old CMake versions. -+if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.30") -+ list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_POLICY_VERSION_MINIMUM=3.5) -+endif() -+ - # and crosscompiling emulator (for try_run() ) - if(CMAKE_CROSSCOMPILING_EMULATOR) - string(REPLACE ";" ${EP_LIST_SEPARATOR} EP_CMAKE_CROSSCOMPILING_EMULATOR -@@ -1720,6 +1725,7 @@ macro(build_thrift) - -DWITH_JAVASCRIPT=OFF - -DWITH_LIBEVENT=OFF - -DWITH_NODEJS=OFF -+ -DWITH_OPENSSL=OFF - -DWITH_PYTHON=OFF - -DWITH_QT5=OFF - -DWITH_ZLIB=OFF) -diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.cmake ---- a/cpp/cmake_modules/BuildUtils.cmake -+++ b/cpp/cmake_modules/BuildUtils.cmake -@@ -112,7 +112,7 @@ function(arrow_create_merged_static_lib output_target) - execute_process(COMMAND ${LIBTOOL_MACOS} -V - OUTPUT_VARIABLE LIBTOOL_V_OUTPUT - OUTPUT_STRIP_TRAILING_WHITESPACE) -- if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools-([0-9.]+).*") -+ if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools(_ld)?-([0-9.]+).*") - message(FATAL_ERROR "libtool found appears to be the incompatible GNU libtool: ${LIBTOOL_MACOS}" - ) - endif() - -diff --git a/cpp/src/arrow/io/interfaces.h b/cpp/src/arrow/io/interfaces.h ---- a/cpp/src/arrow/io/interfaces.h -+++ b/cpp/src/arrow/io/interfaces.h -@@ -210,7 +210,7 @@ - /// \brief Advance or skip stream indicated number of bytes - /// \param[in] nbytes the number to move forward - /// \return Status -- Status Advance(int64_t nbytes); -+ virtual Status Advance(int64_t nbytes); - - /// \brief Return zero-copy string_view to upcoming bytes. - /// ---- a/cpp/src/parquet/arrow/reader.cc -+++ b/cpp/src/parquet/arrow/reader.cc -@@ -254,6 +254,11 @@ - return GetColumn(i, AllRowGroupsFactory(), out); - } +diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h +index b59b59f95c..657a438a3a 100644 +--- a/cpp/src/parquet/file_reader.h ++++ b/cpp/src/parquet/file_reader.h +@@ -210,6 +210,17 @@ class PARQUET_EXPORT ParquetFileReader { + ::arrow::Future<> WhenBuffered(const std::vector& row_groups, + const std::vector& column_indices) const; -+ ::arrow::Status GetColumn( -+ int i, const std::vector& column_indices, -+ FileColumnIteratorFactory iterator_factory, -+ std::unique_ptr* out) override; ++ /// Pre-buffer arbitrary byte ranges (e.g., page-level ranges from OffsetIndex). ++ /// Unlike PreBuffer(), this does NOT set the column bitmap, so ++ /// GetColumnPageReader will use CachedInputStream (page-level cache path). ++ void PreBufferRanges(const std::vector<::arrow::io::ReadRange>& ranges, ++ const ::arrow::io::IOContext& ctx, ++ const ::arrow::io::CacheOptions& options); + - Status GetSchema(std::shared_ptr<::arrow::Schema>* out) override { - return FromParquetSchema(reader_->metadata()->schema(), reader_properties_, - reader_->metadata()->key_value_metadata(), out); -@@ -493,10 +498,40 @@ - - ::arrow::Status BuildArray(int64_t length_upper_bound, - std::shared_ptr<::arrow::ChunkedArray>* out) final { -+ if (!out_) { -+ BEGIN_PARQUET_CATCH_EXCEPTIONS -+ RETURN_NOT_OK( -+ TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_)); -+ END_PARQUET_CATCH_EXCEPTIONS -+ } - *out = out_; - return Status::OK(); - } - -+ std::vector LeafColumnIndices() const final { -+ return {input_->column_index()}; -+ } -+ -+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { -+ if (col_idx != input_->column_index()) return Status::OK(); -+ BEGIN_PARQUET_CATCH_EXCEPTIONS -+ out_ = nullptr; -+ record_reader_->Reset(); -+ record_reader_->Reserve(reserve); -+ return Status::OK(); -+ END_PARQUET_CATCH_EXCEPTIONS -+ } -+ -+ int64_t SkipRecords(int col_idx, int64_t num_records) final { -+ if (col_idx != input_->column_index() || num_records <= 0) return 0; -+ return record_reader_->SkipRecords(num_records); -+ } -+ -+ int64_t ReadRecords(int col_idx, int64_t num_records) final { -+ if (col_idx != input_->column_index() || num_records <= 0) return 0; -+ return record_reader_->ReadRecords(num_records); -+ } ++ /// Wait for arbitrary byte ranges to be pre-buffered. ++ ::arrow::Future<> WhenBufferedRanges( ++ const std::vector<::arrow::io::ReadRange>& ranges) const; + - const std::shared_ptr field() override { return field_; } - private: -@@ -532,6 +567,22 @@ - return storage_reader_->LoadBatch(number_of_records); - } - -+ std::vector LeafColumnIndices() const final { -+ return storage_reader_->LeafColumnIndices(); -+ } -+ -+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { -+ return storage_reader_->ResetLeaf(col_idx, reserve); -+ } -+ -+ int64_t SkipRecords(int col_idx, int64_t num_records) final { -+ return storage_reader_->SkipRecords(col_idx, num_records); -+ } -+ -+ int64_t ReadRecords(int col_idx, int64_t num_records) final { -+ return storage_reader_->ReadRecords(col_idx, num_records); -+ } -+ - Status BuildArray(int64_t length_upper_bound, - std::shared_ptr* out) override { - std::shared_ptr storage; -@@ -576,6 +627,22 @@ - return item_reader_->LoadBatch(number_of_records); - } - -+ std::vector LeafColumnIndices() const final { -+ return item_reader_->LeafColumnIndices(); -+ } -+ -+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { -+ return item_reader_->ResetLeaf(col_idx, reserve); -+ } -+ -+ int64_t SkipRecords(int col_idx, int64_t num_records) final { -+ return item_reader_->SkipRecords(col_idx, num_records); -+ } -+ -+ int64_t ReadRecords(int col_idx, int64_t num_records) final { -+ return item_reader_->ReadRecords(col_idx, num_records); -+ } -+ - virtual ::arrow::Result> AssembleArray( - std::shared_ptr data) { - if (field_->type()->id() == ::arrow::Type::MAP) { -@@ -709,6 +776,39 @@ + // Holds a pointer to an instance of Contents implementation + std::unique_ptr contents_; +diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h +index 4d3acb491e..3906ff3c59 100644 +--- a/cpp/src/parquet/properties.h ++++ b/cpp/src/parquet/properties.h +@@ -139,6 +139,7 @@ static constexpr bool DEFAULT_IS_DICTIONARY_ENABLED = true; + static constexpr int64_t DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT = kDefaultDataPageSize; + static constexpr int64_t DEFAULT_WRITE_BATCH_SIZE = 1024; + static constexpr int64_t DEFAULT_MAX_ROW_GROUP_LENGTH = 1024 * 1024; ++static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = 128 * 1024 * 1024; + static constexpr bool DEFAULT_ARE_STATISTICS_ENABLED = true; + static constexpr int64_t DEFAULT_MAX_STATISTICS_SIZE = 4096; + static constexpr Encoding::type DEFAULT_ENCODING = Encoding::UNKNOWN; +@@ -232,6 +233,7 @@ class PARQUET_EXPORT WriterProperties { + dictionary_pagesize_limit_(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT), + write_batch_size_(DEFAULT_WRITE_BATCH_SIZE), + max_row_group_length_(DEFAULT_MAX_ROW_GROUP_LENGTH), ++ max_row_group_size_(DEFAULT_MAX_ROW_GROUP_SIZE), + pagesize_(kDefaultDataPageSize), + version_(ParquetVersion::PARQUET_2_6), + data_page_version_(ParquetDataPageVersion::V1), +@@ -244,6 +246,7 @@ class PARQUET_EXPORT WriterProperties { + dictionary_pagesize_limit_(properties.dictionary_pagesize_limit()), + write_batch_size_(properties.write_batch_size()), + max_row_group_length_(properties.max_row_group_length()), ++ max_row_group_size_(properties.max_row_group_size()), + pagesize_(properties.data_pagesize()), + version_(properties.version()), + data_page_version_(properties.data_page_version()), +@@ -321,6 +324,13 @@ class PARQUET_EXPORT WriterProperties { + return this; } - return Status::OK(); - } -+ -+ std::vector LeafColumnIndices() const override { -+ std::vector indices; -+ for (const std::unique_ptr& reader : children_) { -+ std::vector child_indices = reader->LeafColumnIndices(); -+ indices.insert(indices.end(), child_indices.begin(), child_indices.end()); -+ } -+ return indices; -+ } -+ -+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) override { -+ for (const std::unique_ptr& reader : children_) { -+ RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve)); -+ } -+ return Status::OK(); -+ } -+ -+ int64_t SkipRecords(int col_idx, int64_t num_records) override { -+ int64_t skipped = 0; -+ for (const std::unique_ptr& reader : children_) { -+ skipped += reader->SkipRecords(col_idx, num_records); -+ } -+ return skipped; -+ } -+ -+ int64_t ReadRecords(int col_idx, int64_t num_records) override { -+ int64_t read = 0; -+ for (const std::unique_ptr& reader : children_) { -+ read += reader->ReadRecords(col_idx, num_records); + ++ /// Specify the max bytes size to put in a single row group. ++ /// Default 128 M. ++ Builder* max_row_group_size(int64_t max_row_group_size) { ++ max_row_group_size_ = max_row_group_size; ++ return this; + } -+ return read; -+ } -+ - Status BuildArray(int64_t length_upper_bound, - std::shared_ptr* out) override; - Status GetDefLevels(const int16_t** data, int64_t* length) override; -@@ -1228,6 +1328,23 @@ - std::unique_ptr result; - RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); - *out = std::move(result); -+ return Status::OK(); -+} + -+::arrow::Status FileReaderImpl::GetColumn( -+ int i, const std::vector& column_indices, -+ FileColumnIteratorFactory iterator_factory, -+ std::unique_ptr* out) { -+ RETURN_NOT_OK(BoundsCheckColumn(i)); -+ auto ctx = std::make_shared(); -+ ctx->reader = reader_.get(); -+ ctx->pool = pool_; -+ ctx->iterator_factory = iterator_factory; -+ ctx->filter_leaves = true; -+ ctx->included_leaves = VectorToSharedSet(column_indices); -+ std::unique_ptr result; -+ RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); -+ *out = std::move(result); - return Status::OK(); - } - ---- a/cpp/src/parquet/arrow/reader.h -+++ b/cpp/src/parquet/arrow/reader.h -@@ -21,6 +21,7 @@ - // N.B. we don't include async_generator.h as it's relatively heavy - #include - #include -+#include - #include - - #include "parquet/file_reader.h" -@@ -48,9 +49,13 @@ + /// Specify the data page size. + /// Default 1MB. + Builder* data_pagesize(int64_t pg_size) { +@@ -664,7 +674,7 @@ class PARQUET_EXPORT WriterProperties { - class ColumnChunkReader; - class ColumnReader; -+class FileColumnIterator; - struct SchemaManifest; - class RowGroupReader; + return std::shared_ptr(new WriterProperties( + pool_, dictionary_pagesize_limit_, write_batch_size_, max_row_group_length_, +- pagesize_, version_, created_by_, page_checksum_enabled_, ++ max_row_group_size_, pagesize_, version_, created_by_, page_checksum_enabled_, + std::move(file_encryption_properties_), default_column_properties_, + column_properties, data_page_version_, store_decimal_as_integer_, + std::move(sorting_columns_))); +@@ -675,6 +685,7 @@ class PARQUET_EXPORT WriterProperties { + int64_t dictionary_pagesize_limit_; + int64_t write_batch_size_; + int64_t max_row_group_length_; ++ int64_t max_row_group_size_; + int64_t pagesize_; + ParquetVersion::type version_; + ParquetDataPageVersion data_page_version_; +@@ -705,6 +716,8 @@ class PARQUET_EXPORT WriterProperties { -+using FileColumnIteratorFactory = -+ std::function; -+ - /// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches. - /// - /// This interfaces caters for different use cases and thus provides different -@@ -136,6 +141,27 @@ - // The indicated column index is relative to the schema - virtual ::arrow::Status GetColumn(int i, std::unique_ptr* out) = 0; + inline int64_t max_row_group_length() const { return max_row_group_length_; } -+ /// \brief Return a ColumnReader with a custom FileColumnIteratorFactory -+ /// and leaf column filtering. -+ /// -+ /// This allows callers to customize page reading behavior (e.g., setting -+ /// data_page_filter for page-level skipping) and to select only specific -+ /// leaf columns within a nested field. The factory is called once per leaf -+ /// column included in column_indices. -+ /// -+ /// \param i top-level field index (same as GetColumn(int i, ...)) -+ /// \param column_indices leaf column indices to include (enables sub-column -+ /// projection within nested types) -+ /// \param iterator_factory factory to create FileColumnIterator per leaf -+ /// \param[out] out the ColumnReader (may be nullptr if all leaves are pruned) -+ virtual ::arrow::Status GetColumn( -+ int i, const std::vector& column_indices, -+ FileColumnIteratorFactory iterator_factory, -+ std::unique_ptr* out) { -+ return ::arrow::Status::NotImplemented( -+ "GetColumn with factory not implemented"); -+ } ++ inline int64_t max_row_group_size() const { return max_row_group_size_; } + - /// \brief Return arrow schema for all the columns. - virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0; + inline int64_t data_pagesize() const { return pagesize_; } -@@ -316,6 +342,43 @@ - // the data available in the file. - virtual ::arrow::Status NextBatch(int64_t batch_size, - std::shared_ptr<::arrow::ChunkedArray>* out) = 0; -+ -+ /// \brief Leaf column indices covered by this (sub)tree, in leaf order. -+ /// -+ /// Used to drive per-leaf row filtering: after page-level skipping each leaf -+ /// lives in its own compressed coordinate space, so callers must reset and -+ /// skip/read each leaf independently rather than in lockstep. -+ virtual std::vector LeafColumnIndices() const { return {}; } -+ -+ /// \brief Reset the leaf identified by col_idx and reserve space for -+ /// `reserve` records (in that leaf's post-page-filter compressed space). -+ /// Must be called before SkipRecords()/ReadRecords() for that leaf, and -+ /// followed by BuildArray() to get the result. -+ virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) { -+ return ::arrow::Status::NotImplemented("ResetLeaf not implemented"); -+ } -+ -+ /// \brief Skip num_records on the leaf identified by col_idx and return the -+ /// number of records actually skipped. Returns 0 when num_records <= 0 or -+ /// col_idx does not belong to this (sub)tree. May throw ParquetException on a -+ /// decode error; callers convert it to Status at the public boundary. -+ virtual int64_t SkipRecords(int col_idx, int64_t num_records) { return 0; } -+ -+ /// \brief Read num_records on the leaf identified by col_idx and return the -+ /// number of records actually read. Values accumulate across successive calls -+ /// until BuildArray() is called. Returns 0 when num_records <= 0 or col_idx -+ /// does not belong to this (sub)tree. May throw ParquetException on a decode -+ /// error; callers convert it to Status at the public boundary. -+ virtual int64_t ReadRecords(int col_idx, int64_t num_records) { return 0; } -+ -+ /// \brief Build the Arrow array from previously loaded data. -+ /// For leaf readers, calls TransferColumnData if not already done. -+ /// For nested readers, assembles the nested array from child arrays. -+ virtual ::arrow::Status BuildArray( -+ int64_t length_upper_bound, -+ std::shared_ptr<::arrow::ChunkedArray>* out) { -+ return ::arrow::Status::NotImplemented("BuildArray not implemented"); -+ } - }; - - /// \brief Experimental helper class for bindings (like Python) that struggle ---- a/cpp/src/parquet/arrow/reader_internal.h -+++ b/cpp/src/parquet/arrow/reader_internal.h -@@ -26,6 +26,7 @@ - #include - #include - -+#include "parquet/arrow/reader.h" - #include "parquet/arrow/schema.h" - #include "parquet/column_reader.h" - #include "parquet/file_reader.h" -@@ -70,7 +71,10 @@ - - virtual ~FileColumnIterator() {} - -- std::unique_ptr<::parquet::PageReader> NextChunk() { -+ /// \brief Fetch the PageReader for the next row group in this iterator's -+ /// range. Virtual so subclasses can decorate the returned PageReader, e.g. -+ /// to install a data_page_filter for I/O-level page skipping. -+ virtual std::unique_ptr<::parquet::PageReader> NextChunk() { - if (row_groups_.empty()) { - return nullptr; - } -@@ -95,9 +99,6 @@ - std::deque row_groups_; - }; - --using FileColumnIteratorFactory = -- std::function; -- - Status TransferColumnData(::parquet::internal::RecordReader* reader, - const std::shared_ptr<::arrow::Field>& value_field, - const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, -diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h ---- a/cpp/src/parquet/column_reader.h -+++ b/cpp/src/parquet/column_reader.h -@@ -76,6 +76,18 @@ struct PARQUET_EXPORT DataPageStats { - std::optional num_rows; - }; - -+/// \brief Identifies a data page that PageReader should read directly. -+/// -+/// The offset is relative to the beginning of the column chunk stream passed to -+/// PageReader::Open. The compressed size includes both the serialized page header -+/// and the compressed page body. The ordinal is the original data page ordinal in -+/// the column chunk and does not include the dictionary page. -+struct PARQUET_EXPORT DataPageReadPlanEntry { -+ int32_t page_ordinal; -+ int64_t offset; -+ int32_t compressed_page_size; -+}; -+ - class PARQUET_EXPORT LevelDecoder { - public: - LevelDecoder(); -@@ -147,9 +159,21 @@ class PARQUET_EXPORT PageReader { - // ApplicationVersion::HasCorrectStatistics(). - // \note API EXPERIMENTAL - void set_data_page_filter(DataPageFilter data_page_filter) { -+ if (data_page_read_plan_enabled_) { -+ throw ParquetException( -+ "data_page_filter and data_page_read_plan cannot be enabled together"); -+ } - data_page_filter_ = std::move(data_page_filter); - } - -+ /// Configure PageReader to jump directly to selected data pages before reading -+ /// their headers. `first_data_page_offset` and each entry offset are relative to -+ /// the beginning of the column chunk stream. Dictionary pages before -+ /// `first_data_page_offset` are still read normally. -+ // \note API EXPERIMENTAL -+ void set_data_page_read_plan(int64_t first_data_page_offset, -+ std::vector data_pages); -+ - // @returns: shared_ptr(nullptr) on EOS, std::shared_ptr - // containing new Page otherwise - // -@@ -162,6 +186,11 @@ class PARQUET_EXPORT PageReader { - protected: - // Callback that decides if we should skip a page or not. - DataPageFilter data_page_filter_; -+ -+ bool data_page_read_plan_enabled_ = false; -+ int64_t first_data_page_offset_ = 0; -+ std::vector data_page_read_plan_; -+ size_t next_data_page_ = 0; - }; - - class PARQUET_EXPORT ColumnReader { -diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc ---- a/cpp/src/parquet/column_reader.cc -+++ b/cpp/src/parquet/column_reader.cc -@@ -207,6 +207,39 @@ ReaderProperties default_reader_properties() { - return default_reader_properties; - } - -+void PageReader::set_data_page_read_plan( -+ int64_t first_data_page_offset, -+ std::vector data_pages) { -+ if (data_page_filter_) { -+ throw ParquetException( -+ "data_page_filter and data_page_read_plan cannot be enabled together"); -+ } -+ if (first_data_page_offset < 0) { -+ throw ParquetException("Invalid negative first data page offset"); -+ } -+ -+ int64_t previous_end = first_data_page_offset; -+ int32_t previous_ordinal = -1; -+ for (const auto& page : data_pages) { -+ int64_t page_end; -+ if (page.page_ordinal < 0 || page.offset < first_data_page_offset || -+ page.compressed_page_size <= 0 || -+ AddWithOverflow(page.offset, page.compressed_page_size, &page_end)) { -+ throw ParquetException("Invalid data page read plan entry"); -+ } -+ if (page.offset < previous_end || page.page_ordinal <= previous_ordinal) { -+ throw ParquetException("Data page read plan entries must be ordered"); -+ } -+ previous_end = page_end; -+ previous_ordinal = page.page_ordinal; -+ } -+ -+ data_page_read_plan_enabled_ = true; -+ first_data_page_offset_ = first_data_page_offset; -+ data_page_read_plan_ = std::move(data_pages); -+ next_data_page_ = 0; -+} -+ - namespace { - - // Extracts encoded statistics from V1 and V2 data page headers -@@ -430,9 +463,43 @@ std::shared_ptr SerializedPageReader::NextPage() { - - // Loop here because there may be unhandled page types that we skip until - // finding a page that we do know what to do with -- while (seen_num_values_ < total_num_values_) { -+ while (data_page_read_plan_enabled_ || seen_num_values_ < total_num_values_) { -+ const DataPageReadPlanEntry* planned_data_page = nullptr; -+ uint32_t page_header_limit = max_page_header_size_; -+ -+ if (data_page_read_plan_enabled_) { -+ if (next_data_page_ >= data_page_read_plan_.size()) { -+ return nullptr; -+ } -+ -+ PARQUET_ASSIGN_OR_THROW(int64_t current_position, stream_->Tell()); -+ if (current_position < first_data_page_offset_) { -+ page_header_limit = static_cast(std::min( -+ page_header_limit, first_data_page_offset_ - current_position)); -+ } else { -+ planned_data_page = &data_page_read_plan_[next_data_page_]; -+ if (current_position > planned_data_page->offset) { -+ throw ParquetException("Data page read plan points behind stream position"); -+ } -+ PARQUET_THROW_NOT_OK( -+ stream_->Advance(planned_data_page->offset - current_position)); -+ PARQUET_ASSIGN_OR_THROW(int64_t target_position, stream_->Tell()); -+ if (target_position != planned_data_page->offset) { -+ throw ParquetException("Failed to seek to planned data page"); -+ } -+ page_ordinal_ = planned_data_page->page_ordinal; -+ page_header_limit = static_cast(std::min( -+ page_header_limit, planned_data_page->compressed_page_size)); -+ } -+ } -+ -+ if (page_header_limit == 0) { -+ throw ParquetException("No bytes available for page header"); -+ } -+ - uint32_t header_size = 0; -- uint32_t allowed_page_size = kDefaultPageHeaderSize; -+ uint32_t allowed_page_size = -+ std::min(kDefaultPageHeaderSize, page_header_limit); - - // Page headers can be very large because of page statistics - // We try to deserialize a larger buffer progressively -@@ -458,11 +525,12 @@ std::shared_ptr SerializedPageReader::NextPage() { - // Failed to deserialize. Double the allowed page header size and try again - std::stringstream ss; - ss << e.what(); -- allowed_page_size *= 2; -- if (allowed_page_size > max_page_header_size_) { -+ if (allowed_page_size >= page_header_limit) { - ss << "Deserializing page header failed.\n"; - throw ParquetException(ss.str()); - } -+ allowed_page_size = -+ std::min(allowed_page_size * 2, page_header_limit); - } - } - // Advance the stream offset -@@ -474,6 +542,20 @@ std::shared_ptr SerializedPageReader::NextPage() { - throw ParquetException("Invalid page header"); - } - -+ const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); -+ if (planned_data_page != nullptr) { -+ if (page_type != PageType::DATA_PAGE && page_type != PageType::DATA_PAGE_V2) { -+ throw ParquetException("Data page read plan points to a non-data page"); -+ } -+ int64_t total_compressed_size; -+ if (AddWithOverflow(static_cast(header_size), -+ static_cast(compressed_len), -+ &total_compressed_size) || -+ total_compressed_size != planned_data_page->compressed_page_size) { -+ throw ParquetException("Planned data page size does not match page header"); -+ } -+ } -+ - EncodedStatistics data_page_statistics; - if (ShouldSkipPage(&data_page_statistics)) { - PARQUET_THROW_NOT_OK(stream_->Advance(compressed_len)); -@@ -494,8 +576,6 @@ std::shared_ptr SerializedPageReader::NextPage() { - ParquetException::EofException(ss.str()); - } - -- const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); -- - if (properties_.page_checksum_verification() && current_page_header_.__isset.crc && - PageCanUseChecksum(page_type)) { - // verify crc -@@ -534,6 +614,9 @@ std::shared_ptr SerializedPageReader::NextPage() { - LoadEnumSafe(&dict_header.encoding), - is_sorted); - } else if (page_type == PageType::DATA_PAGE) { -+ if (planned_data_page != nullptr) { -+ ++next_data_page_; -+ } - ++page_ordinal_; - const format::DataPageHeader& header = current_page_header_.data_page_header; - page_buffer = -@@ -545,6 +628,9 @@ std::shared_ptr SerializedPageReader::NextPage() { - LoadEnumSafe(&header.repetition_level_encoding), uncompressed_len, - std::move(data_page_statistics)); - } else if (page_type == PageType::DATA_PAGE_V2) { -+ if (planned_data_page != nullptr) { -+ ++next_data_page_; -+ } - ++page_ordinal_; - const format::DataPageHeaderV2& header = current_page_header_.data_page_header_v2; + inline ParquetDataPageVersion data_page_version() const { +@@ -810,7 +823,7 @@ class PARQUET_EXPORT WriterProperties { + private: + explicit WriterProperties( + MemoryPool* pool, int64_t dictionary_pagesize_limit, int64_t write_batch_size, +- int64_t max_row_group_length, int64_t pagesize, ParquetVersion::type version, ++ int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, ParquetVersion::type version, + const std::string& created_by, bool page_write_checksum_enabled, + std::shared_ptr file_encryption_properties, + const ColumnProperties& default_column_properties, +@@ -821,6 +834,7 @@ class PARQUET_EXPORT WriterProperties { + dictionary_pagesize_limit_(dictionary_pagesize_limit), + write_batch_size_(write_batch_size), + max_row_group_length_(max_row_group_length), ++ max_row_group_size_(max_row_group_size), + pagesize_(pagesize), + parquet_data_page_version_(data_page_version), + parquet_version_(version), +@@ -836,6 +850,7 @@ class PARQUET_EXPORT WriterProperties { + int64_t dictionary_pagesize_limit_; + int64_t write_batch_size_; + int64_t max_row_group_length_; ++ int64_t max_row_group_size_; + int64_t pagesize_; + ParquetDataPageVersion parquet_data_page_version_; + ParquetVersion::type parquet_version_; diff --git a/docs/source/user_guide/data_types.rst b/docs/source/user_guide/data_types.rst index 3d529332..add537ad 100644 --- a/docs/source/user_guide/data_types.rst +++ b/docs/source/user_guide/data_types.rst @@ -186,6 +186,25 @@ and `Arrow DataTypes `` where t is the data type of the contained elements. + * - ``VECTOR`` + - FixedSizeList + - Data type of a dense vector containing exactly ``n`` elements of type ``t``. + + ``n`` must be positive. ``t`` can be ``BOOLEAN``, ``TINYINT``, + ``SMALLINT``, ``INT``, ``BIGINT``, ``FLOAT``, or ``DOUBLE``. A VECTOR + value may be NULL, but its elements cannot be NULL. + + Paimon C++ currently supports VECTOR columns only in append-only tables + backed by Parquet data files. They use the standard Parquet LIST + representation on disk and are restored as Arrow ``FixedSizeList`` + values on read. Primary-key tables and data-evolution tables containing + VECTOR fields are rejected. VECTOR columns also cannot be partition or + bucket keys. Dedicated vector storage is not included yet. + + Paimon C++ also reads Parquet files written by Paimon Rust or Python whose + embedded Arrow schema restores VECTOR columns as ``FixedSizeList``, + including NULL vector values. + * - ``MAP`` - Map - Data type of an associative array that maps keys (including NULL) to values (including NULL). A map cannot contain duplicate keys; each key can map to at most one value. diff --git a/include/paimon/api.h b/include/paimon/api.h index d7566685..81f236bc 100644 --- a/include/paimon/api.h +++ b/include/paimon/api.h @@ -33,6 +33,7 @@ #include "paimon/record_batch.h" // IWYU pragma: export #include "paimon/result.h" // IWYU pragma: export #include "paimon/scan_context.h" // IWYU pragma: export +#include "paimon/statistics_mode.h" // IWYU pragma: export #include "paimon/status.h" // IWYU pragma: export #include "paimon/table/source/table_read.h" // IWYU pragma: export #include "paimon/table/source/table_scan.h" // IWYU pragma: export diff --git a/include/paimon/defs.h b/include/paimon/defs.h index 9fcf8e34..8062d3d2 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -50,6 +50,8 @@ enum class FieldType { STRUCT = 15, BLOB = 16, VARIANT = 17, + /// Fixed-length dense vector represented by Arrow FixedSizeList. + VECTOR = 18, UNKNOWN = 128, }; @@ -198,6 +200,10 @@ struct PAIMON_EXPORT Options { /// cache. Default value is 0. static const char SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS[]; + /// "scan.manifest-entry.lazy-decode.enabled" - Whether to deserialize only manifest entries + /// for the target bucket when rebuilding the cache. Default value is true. + static const char SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED[]; + /// "read.batch-size" - Read batch size for any file format if it supports. /// The default value is 1024. static const char READ_BATCH_SIZE[]; @@ -403,6 +409,10 @@ struct PAIMON_EXPORT Options { /// "file-index.read.enabled" - Whether enabled read file index. Default value is "true". static const char FILE_INDEX_READ_ENABLED[]; + /// "file-index.in-manifest-threshold" - The threshold to store file index bytes in the + /// manifest. Default value is 500B. + static const char FILE_INDEX_IN_MANIFEST_THRESHOLD[]; + /// "data-file.external-paths" - The external paths where the data of this table will be /// written, multiple elements separated by commas. static const char DATA_FILE_EXTERNAL_PATHS[]; @@ -552,10 +562,18 @@ struct PAIMON_EXPORT Options { /// "scan.timestamp" can be used as an alternative string input for the same mode. static const char SCAN_TIMESTAMP_MILLIS[]; + /// "realtime.enabled" - Whether real-time write, commit, and read operations are enabled. + /// Default value is "false". + static const char REALTIME_ENABLED[]; + /// "realtime.read-view-ttl" - Lifetime of a real-time memory view pinned by scan planning /// before reader creation. Default value is "5 min". static const char REALTIME_READ_VIEW_TTL[]; + /// "realtime.store.stats-mode" - Statistics collected by the default real-time store. + /// Supported values are "none" and "full". Default value is "none". + static const char REALTIME_STORE_STATS_MODE[]; + /// "scan.timestamp" - Optional timestamp string used in case of "from-timestamp" scan mode, /// as an alternative to "scan.timestamp-millis". /// It will be automatically converted to timestamp in unix milliseconds, using local time zone. diff --git a/include/paimon/factories/singleton.h b/include/paimon/factories/singleton.h index 6e12d456..a3030e5e 100644 --- a/include/paimon/factories/singleton.h +++ b/include/paimon/factories/singleton.h @@ -19,7 +19,9 @@ #pragma once +#include #include +#include #include "paimon/macros.h" #include "paimon/visibility.h" @@ -30,9 +32,9 @@ class PAIMON_EXPORT LazyInstantiation { protected: template static void Create(T*& ptr) { - T* tmp = new T; - MEMORY_BARRIER(); - ptr = tmp; + // Publication ordering is handled by the release store in + // Singleton::GetInstance(), so no barrier is needed here. + ptr = new T; static std::shared_ptr destroyer(ptr); } }; @@ -56,4 +58,35 @@ class PAIMON_EXPORT Singleton : private InstPolicy { static T* GetInstance(); }; +template +T* Singleton::GetInstance() { + static std::atomic ptr{nullptr}; + static std::mutex mutex; + T* p = ptr.load(std::memory_order_acquire); + if (PAIMON_UNLIKELY(p == nullptr)) { + std::lock_guard lg(mutex); + // Re-check under the mutex with a relaxed load; the mutex already + // synchronizes with the creating thread. + p = ptr.load(std::memory_order_relaxed); + if (p == nullptr) { + InstPolicy::Create(p); + ptr.store(p, std::memory_order_release); + } + } + return p; +} + +// FactoryCreator and IOHook are instantiated exactly once in singleton.cpp, and the +// extern declarations below suppress implicit instantiation everywhere else. The +// file-format/file-system plugins are separate shared libraries linked with +// -Bsymbolic, so a per-library copy of GetInstance()'s function-local static state +// would never be interposed: factory registrations would land in a different +// instance than lookups. Do not replace these with implicit instantiation. Types local to a single +// translation unit (e.g. test-only types) can still instantiate Singleton +// implicitly because they cannot span library boundaries. +class FactoryCreator; +class IOHook; +extern template class Singleton; +extern template class Singleton; + } // namespace paimon diff --git a/include/paimon/file_index/file_index_format.h b/include/paimon/file_index/file_index_format.h index b46dee8c..3993b624 100644 --- a/include/paimon/file_index/file_index_format.h +++ b/include/paimon/file_index/file_index_format.h @@ -19,6 +19,7 @@ #pragma once #include +#include #include #include #include @@ -32,6 +33,8 @@ struct ArrowSchema; namespace paimon { class InputStream; class MemoryPool; +class Bytes; +class OutputStream; /// Defines the on-disk format and versioning for Paimon file-level indexes. /// File index file format. Put all column and offset in the header. @@ -88,9 +91,15 @@ class MemoryPool; class PAIMON_EXPORT FileIndexFormat { public: class Reader; + class Writer; + + /// Serialized file indexes grouped as column name -> index type -> index bytes. A null bytes + /// pointer represents an empty index for that column and index type. + /// For example, indexes["col1"]["bsi"] = ; + using ColumnIndexes = std::map>>; + /// Creates a `Reader` to parse a index file (may contain multiple indexes) from the given input /// stream. - /// /// @param input_stream Input stream containing serialized index data. /// @param pool Memory pool for temporary allocations during reading. /// @return A unique pointer to a `Reader` on success, or an error if the stream is invalid @@ -98,18 +107,38 @@ class PAIMON_EXPORT FileIndexFormat { static Result> CreateReader( const std::shared_ptr& input_stream, const std::shared_ptr& pool); + /// Creates a `Writer` which serializes a complete V1 file index container. + /// + /// @param output_stream Destination stream for serialized index data. + /// @param pool Memory pool for writer-side allocations. + /// @return A unique pointer to a `Writer` on success. + static Result> CreateWriter( + const std::shared_ptr& output_stream, + const std::shared_ptr& pool); + public: static const int64_t MAGIC; static const int32_t EMPTY_INDEX_FLAG; static const int32_t V_1; }; +/// Writer for file index file. +class FileIndexFormat::Writer { + public: + virtual ~Writer() = default; + + /// Writes all column indexes. This is a terminal, one-shot operation. + virtual Status WriteColumnIndexes(const FileIndexFormat::ColumnIndexes& indexes) = 0; + + /// Flushes and closes the output stream supplied to `CreateWriter()`. + virtual Status Close() = 0; +}; + /// Reader for file index file. class FileIndexFormat::Reader { public: virtual ~Reader() = default; /// Reads index data for a specific column from the index file. - /// /// @param column_name Name of the column to retrieve index data for. /// @param arrow_schema Arrow schema that must contain a field corresponding to `column_name`. /// @return A vector of shared pointers to FileIndexReader objects, each corresponding to a diff --git a/include/paimon/file_store_commit.h b/include/paimon/file_store_commit.h index 63efb562..8af77695 100644 --- a/include/paimon/file_store_commit.h +++ b/include/paimon/file_store_commit.h @@ -79,11 +79,17 @@ class PAIMON_EXPORT FileStoreCommit { /// orders them by partition, bucket, and offset before validating continuity. The resulting /// snapshot atomically publishes the data files and the updated offset map. /// + /// If this method returns an error, the caller may retry with the same arguments. Each call + /// reloads the latest committed state. As in `FilterAndCommit`, a retry's identifier is + /// considered committed when it is not newer than the latest identifier for `commit_user`. + /// The requested offset ranges must also be covered by the latest committed progress. + /// /// @param realtime_commits Commit messages and left-closed, right-open offset ranges to /// commit. /// @param commit_identifier Identifier of the streaming commit operation. /// @param watermark Optional event-time watermark. - /// @return The id of the final snapshot produced by this commit. + /// @return The id of the latest snapshot containing the committed progress. On retry, this may + /// be a snapshot produced by a later commit and is suitable for refreshing a real-time context. virtual Result CommitWithProgress( const std::vector& realtime_commits, int64_t commit_identifier, std::optional watermark) = 0; @@ -117,6 +123,10 @@ class PAIMON_EXPORT FileStoreCommit { /// @param watermark An optional event-time watermark used to indicate the progress of data /// processing. Default is std::nullopt. /// @return Result of the operation. + /// @note A full-table overwrite clears all committed real-time progress. A partition + /// overwrite removes progress only for matching partitions. In either case, active + /// real-time writers and their `RealtimeContext` instances must be recreated before + /// further real-time operations. virtual Status Overwrite(const std::map& partition, const std::vector>& commit_messages, int64_t commit_identifier, @@ -131,6 +141,10 @@ class PAIMON_EXPORT FileStoreCommit { /// @param watermark An optional event-time watermark used to indicate the progress of data /// processing. Default is std::nullopt. /// @return Result of the operation. + /// @note A full-table overwrite clears all committed real-time progress. A partition + /// overwrite removes progress only for matching partitions. In either case, active + /// real-time writers and their `RealtimeContext` instances must be recreated before + /// further real-time operations. virtual Result FilterAndOverwrite( const std::map& partition, const std::vector>& commit_messages, @@ -157,6 +171,9 @@ class PAIMON_EXPORT FileStoreCommit { /// @param partitions A vector of partitions to be dropped. /// @param commit_identifier An identifier for the commit operation. /// @return Status indicating the success or failure of the drop partition operation. + /// @note A partition drop removes committed real-time progress only for matching partitions. + /// Active real-time writers and their `RealtimeContext` instances must be recreated before + /// further real-time operations. virtual Status DropPartition(const std::vector>& partitions, int64_t commit_identifier) = 0; @@ -165,6 +182,9 @@ class PAIMON_EXPORT FileStoreCommit { /// /// @param commit_identifier An identifier for the commit operation. /// @return Status indicating the success or failure of the truncate operation. + /// @note Truncation clears all committed real-time progress. Active real-time writers and + /// their `RealtimeContext` instances must be recreated before further real-time + /// operations. virtual Status TruncateTable(int64_t commit_identifier) = 0; /// Abort an unsuccessful commit. The data and index files described by the given commit @@ -182,6 +202,9 @@ class PAIMON_EXPORT FileStoreCommit { /// @param target_snapshot_id The snapshot id to roll back to. /// @return Result; true if the atomic commit succeeded. Returns an error status if /// there is no latest snapshot or the target snapshot does not exist. + /// @note Rollback restores the real-time progress recorded by the target snapshot. Active + /// real-time writers and their `RealtimeContext` instances must be recreated before + /// further real-time operations. virtual Result RollbackToAsLatest(int64_t target_snapshot_id) = 0; /// Configure row-id conflict checking from a specific snapshot id. diff --git a/include/paimon/file_store_write.h b/include/paimon/file_store_write.h index fdc172c3..1d7ca088 100644 --- a/include/paimon/file_store_write.h +++ b/include/paimon/file_store_write.h @@ -107,6 +107,12 @@ class PAIMON_EXPORT FileStoreWrite { /// /// The writer loads the snapshot's partition-bucket offsets and releases sealed memory that is /// fully covered by disk. Calling this method on a non-real-time writer returns an error. + /// If the snapshot overwrites table contents or moves committed progress backwards, such as + /// after a partition drop, overwrite, or rollback, this method returns an error and the caller + /// must recreate the `RealtimeContext` and writer. These operations are not fenced against an + /// active writer and do not clear its process-local state automatically. The caller must + /// coordinate them with active writers; skipping the resetting snapshot and continuing to use + /// an old context is unsupported. virtual Status RefreshCommittedSnapshot(int64_t snapshot_id); virtual std::shared_ptr GetMetrics() const = 0; diff --git a/include/paimon/format/column_stats.h b/include/paimon/format/column_stats.h index f16cb254..a4e3de48 100644 --- a/include/paimon/format/column_stats.h +++ b/include/paimon/format/column_stats.h @@ -33,8 +33,8 @@ namespace paimon { /// ColumnStats is an abstract base class that represents statistical information for data columns /// in Paimon tables. It provides min/max values and null count statistics /// -/// Only primitive data types support min/max statistics. Nested types (arrays, maps, structs) only -/// track null counts through `NestedColumnStats`. +/// Only primitive data types support min/max statistics. Nested types (arrays, vectors, maps, +/// structs) only track null counts through `NestedColumnStats`. /// /// @note This is an abstract base class. Use the static factory methods `CreateXXXColumnStats()` to /// create concrete instances for specific data types. @@ -52,7 +52,7 @@ class PAIMON_EXPORT ColumnStats { /// @name CreateXXXColumnStats() /// %Factory methods `CreateXXXColumnStats()` to create column statistics. /// - min/max/null_count for primitive data types - /// - null_count for nested data types (arrays, maps, structs) + /// - null_count for nested data types (arrays, vectors, maps, structs) /// /// @{ static std::unique_ptr CreateBooleanColumnStats(std::optional min, @@ -88,8 +88,8 @@ class PAIMON_EXPORT ColumnStats { static std::unique_ptr CreateDateColumnStats(std::optional min, std::optional max, std::optional null_count); - /// Creates column statistics for nested data types (arrays, maps, structs), which only track - /// null counts. + /// Creates column statistics for nested data types (arrays, vectors, maps, structs), which only + /// track null counts. static std::unique_ptr CreateNestedColumnStats(const FieldType& nested_type, std::optional null_count); /// @} @@ -180,7 +180,7 @@ class PAIMON_EXPORT NestedColumnStats : public ColumnStats { NestedColumnStats(const FieldType& nested_type, std::optional null_count) : nested_type_(nested_type), null_count_(null_count) { assert(nested_type == FieldType::ARRAY || nested_type == FieldType::MAP || - nested_type == FieldType::STRUCT); + nested_type == FieldType::STRUCT || nested_type == FieldType::VECTOR); } std::optional NullCount() const override { diff --git a/include/paimon/format/read_hints.h b/include/paimon/format/read_hints.h new file mode 100644 index 00000000..d60b4132 --- /dev/null +++ b/include/paimon/format/read_hints.h @@ -0,0 +1,35 @@ +/* + * 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. + */ + +#pragma once + +#include "paimon/visibility.h" + +namespace paimon { + +/// Runtime state of the framework read path, passed to format layers via +/// `ReaderBuilder::WithReadHints` so each format can adapt its internal behavior +/// (e.g. whether parquet enables its own pre-buffering). +struct PAIMON_EXPORT ReadHints { + /// Whether framework-level prefetch is enabled for this read. + bool prefetch_enabled = false; + /// Whether the shared read-ahead cache is enabled for this read. + bool read_ahead_cache_enabled = false; +}; + +} // namespace paimon diff --git a/include/paimon/format/reader_builder.h b/include/paimon/format/reader_builder.h index b0837a26..5a28077a 100644 --- a/include/paimon/format/reader_builder.h +++ b/include/paimon/format/reader_builder.h @@ -19,7 +19,9 @@ #pragma once #include +#include +#include "paimon/format/read_hints.h" #include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" #include "paimon/type_fwd.h" @@ -41,6 +43,15 @@ class PAIMON_EXPORT ReaderBuilder { return this; } + /// Inject runtime read state from the framework layer, so the format can adapt + /// its internal behavior accordingly. When present, the hints describe the + /// authoritative runtime state of this read; when absent, the format should fall + /// back to its own options. + virtual ReaderBuilder* WithReadHints(const std::optional& hints) { + (void)hints; + return this; + } + /// Build a file batch reader based on the created `InputStream`. virtual Result> Build( const std::shared_ptr& path) const = 0; diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 9bed5402..3e58b1c4 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -30,7 +30,7 @@ #include "paimon/predicate/predicate.h" #include "paimon/result.h" #include "paimon/type_fwd.h" -#include "paimon/utils/read_ahead_cache.h" +#include "paimon/utils/prefetch_cache_config.h" #include "paimon/visibility.h" namespace paimon { @@ -59,9 +59,8 @@ class PAIMON_EXPORT ReadContext { const std::shared_ptr& specific_file_system, const std::map& fs_scheme_to_identifier_map, const std::shared_ptr& realtime_context, - const std::map& options, - PrefetchCacheMode prefetch_cache_mode, const CacheConfig& cache_config, - const std::shared_ptr& cache); + const std::map& options, bool read_ahead_cache_enabled, + const CacheConfig& cache_config, const std::shared_ptr& cache); ~ReadContext(); const std::string& GetPath() const { @@ -128,8 +127,8 @@ class PAIMON_EXPORT ReadContext { return realtime_context_; } - PrefetchCacheMode GetPrefetchCacheMode() const { - return prefetch_cache_mode_; + bool ReadAheadCacheEnabled() const { + return read_ahead_cache_enabled_; } const CacheConfig& GetCacheConfig() const { @@ -175,7 +174,7 @@ class PAIMON_EXPORT ReadContext { std::map fs_scheme_to_identifier_map_; std::shared_ptr realtime_context_; std::map options_; - PrefetchCacheMode prefetch_cache_mode_; + bool read_ahead_cache_enabled_; CacheConfig cache_config_; std::shared_ptr cache_; // Owns schema resources and releases ArrowSchema::release in destructor. @@ -307,13 +306,13 @@ class PAIMON_EXPORT ReadContextBuilder { /// @return Reference to this builder for method chaining. ReadContextBuilder& EnablePrefetch(bool enabled); - /// Set prefetch cache mode for read operations. + /// Enable or disable the read-ahead cache for read operations. /// - /// A prefetch cache is used to prebuffer data ranges before they are needed, + /// A read-ahead cache is used to prebuffer data ranges before they are needed, /// which can improve read performance by reducing redundant I/O operations. - /// @param mode (default: PrefetchCacheMode::ALWAYS) + /// @param enabled Whether to enable the read-ahead cache (default: true) /// @return Reference to this builder for method chaining. - ReadContextBuilder& SetPrefetchCacheMode(PrefetchCacheMode mode); + ReadContextBuilder& SetReadAheadCacheEnabled(bool enabled); /// Set the cache configuration for prefetch read operations. /// diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 3dc257fc..153d524d 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -26,11 +26,7 @@ namespace paimon { /// Factory for Paimon's default Arrow-backed `RealtimeStore`. class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: - /// Creates an Arrow-backed store for one partition and bucket. - Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, - const std::map& options, - const std::shared_ptr& memory_pool) override; + Result> Create(RealtimeStoreCreateRequest&& request) override; }; } // namespace paimon diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index 9bc870fe..200e4ba4 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -73,6 +73,11 @@ using RealtimeOffsetMap = std::map; /// reads. `RealtimeContext` itself is not a customization interface and must not be implemented by /// applications. Customize real-time storage and retrieval through `RealtimeStoreFactory` and /// `RealtimeStore` instead. +/// +/// A context is valid only for one uninterrupted committed-progress history. Overwrite, truncate, +/// partition drop, and rollback operations do not automatically clear process-local real-time +/// state. Applications must coordinate these operations with active real-time writers and recreate +/// the `RealtimeContext` and writers before continuing. class PAIMON_EXPORT RealtimeContext { public: /// Creates a context backed by Paimon's default in-memory Arrow `RealtimeStore`. diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index 9177fc06..241413b3 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -27,10 +27,12 @@ #include #include +#include "arrow/c/abi.h" #include "paimon/reader/batch_reader.h" #include "paimon/realtime/offset_range.h" #include "paimon/record_batch.h" #include "paimon/result.h" +#include "paimon/statistics_mode.h" #include "paimon/visibility.h" struct ArrowSchema; @@ -40,10 +42,37 @@ namespace paimon { class MemoryPool; class Predicate; -/// A table record batch and its framework-assigned contiguous offset range. +enum class PAIMON_EXPORT RealtimeStoreMode { + APPEND_ONLY, + PRIMARY_KEY, +}; + +/// Parameters used by a `RealtimeStoreFactory` to create a store. +struct PAIMON_EXPORT RealtimeStoreCreateRequest { + /// Schema whose ownership is transferred to the factory. Append mode receives the complete + /// table write schema. Primary-key mode receives the prepared transport schema: + /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields]. + std::unique_ptr<::ArrowSchema> write_schema; + /// Table options available to the store implementation. + std::map options; + /// Memory pool for allocations retained by the store. + std::shared_ptr memory_pool; + /// Partition values identifying the store. + std::map partition; + /// Bucket identifying the store within its partition. + int32_t bucket = -1; + /// Table mode implemented by the store. + RealtimeStoreMode mode = RealtimeStoreMode::APPEND_ONLY; + /// Statistics collected by append-only stores. + StatisticsMode statistics_mode = StatisticsMode::NONE; +}; + +/// A record batch and its framework-assigned contiguous offset range. /// -/// The batch contains only table write fields. Row `i` is associated with -/// `offset_range.begin + i`; the offset is progress metadata and is not a table field. +/// Append-mode batches contain table write fields, and row `i` has offset +/// `offset_range.begin + i`. Primary-key batches use the prepared transport schema, are sorted +/// by full primary key then sequence number, and retain the original offset in +/// `_REALTIME_OFFSET`. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -78,7 +107,10 @@ class PAIMON_EXPORT RealtimeReadView { /// Parameters used by a `RealtimeStore` to create readers for a query. struct PAIMON_EXPORT RealtimeQueryContext { - /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added. + /// Append mode receives the requested output fields before the mandatory leading + /// `_VALUE_KIND` field is added. Primary-key mode receives the complete prepared schema. + /// This schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must + /// import or copy it synchronously. ::ArrowSchema* read_schema; /// Predicate using field indexes from `read_schema`. std::shared_ptr predicate; @@ -115,9 +147,10 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// - /// Concatenating the returned readers must produce every sealed row exactly once and in write - /// order. Each output batch contains `_VALUE_KIND` followed by all fields from the factory's - /// `write_schema`. + /// The returned readers collectively expose every sealed row exactly once. Append-mode readers + /// preserve write order and contain `_VALUE_KIND` followed by table write fields. Primary-key + /// readers use the prepared transport schema; each reader's complete stream is sorted by full + /// primary key then sequence number. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -127,13 +160,15 @@ class PAIMON_EXPORT RealtimeStore { /// also provide a consistent snapshot when a write or seal is in progress. virtual Result> AcquireReadView() = 0; - /// Creates readers over rows in `view` whose offsets are greater than or equal to - /// `offset_begin`. + /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater than + /// or equal to `offset_begin`; primary-key mode ignores `offset_begin`. /// - /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by - /// `context.read_schema` except a duplicate `_VALUE_KIND`. Concatenating all returned readers - /// must produce every matching row once. Paimon retains `view` for the lifetime of the - /// resulting framework reader. + /// Append-mode batches contain `_VALUE_KIND` followed by the requested fields except a + /// duplicate `_VALUE_KIND`, and collectively expose every matching row exactly once. + /// Primary-key batches use the prepared transport schema and may contain multiple mutations + /// per key; each reader's complete stream is sorted by full primary key then sequence number, + /// and the readers collectively expose every raw mutation exactly once. Paimon retains `view` + /// for the lifetime of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; @@ -156,15 +191,9 @@ class PAIMON_EXPORT RealtimeStoreFactory { public: virtual ~RealtimeStoreFactory() = default; - /// Creates a store configured with the supplied schema, options, and memory pool. - /// @param write_schema Complete table write schema whose ownership is transferred to the - /// factory. The factory may consume it or retain it in the created store. - /// @param options Effective table options available to the store. - /// @param memory_pool Memory pool provided by the write context. - virtual Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, - const std::map& options, - const std::shared_ptr& memory_pool) = 0; + /// Creates a store for the requested table mode and partition-bucket. + /// The factory consumes `request`, including ownership of `request.write_schema`. + virtual Result> Create(RealtimeStoreCreateRequest&& request) = 0; }; } // namespace paimon diff --git a/include/paimon/scan_context.h b/include/paimon/scan_context.h index dc780ad7..9c0b9d47 100644 --- a/include/paimon/scan_context.h +++ b/include/paimon/scan_context.h @@ -170,7 +170,7 @@ class PAIMON_EXPORT ScanContextBuilder { ScanContextBuilder& SetGlobalIndexResult( const std::shared_ptr& global_index_result); - /// Enables process-local union reads with the memory indexers owned by `realtime_context`. + /// Enables process-local union reads with the real-time stores owned by `realtime_context`. ScanContextBuilder& WithRealtimeContext( const std::shared_ptr& realtime_context); diff --git a/include/paimon/statistics_mode.h b/include/paimon/statistics_mode.h new file mode 100644 index 00000000..0be0e910 --- /dev/null +++ b/include/paimon/statistics_mode.h @@ -0,0 +1,32 @@ +/* + * 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. + */ + +#pragma once + +namespace paimon { + +/// Controls the amount of statistics collected for metadata pruning. +enum class StatisticsMode { + /// Do not collect statistics. + NONE, + /// Collect statistics for all supported fields. + FULL, +}; + +} // namespace paimon diff --git a/include/paimon/utils/prefetch_cache_config.h b/include/paimon/utils/prefetch_cache_config.h new file mode 100644 index 00000000..4bbf1ecd --- /dev/null +++ b/include/paimon/utils/prefetch_cache_config.h @@ -0,0 +1,76 @@ +/* + * 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. + */ + +// Adapted from Apache ORC +// https://github.com/apache/orc/blob/main/c%2B%2B/src/io/Cache.hh + +#pragma once + +#include + +#include "paimon/visibility.h" + +namespace paimon { + +/// Configuration parameters for the read-ahead cache behavior. +/// +/// This struct controls various limits and prefetching strategies used by +/// ReadAheadCache to balance memory usage, I/O efficiency, and latency hiding. +class PAIMON_EXPORT CacheConfig { + public: + CacheConfig(); + CacheConfig(uint64_t range_size_limit, uint64_t hole_size_limit, uint64_t pre_buffer_limit); + + /// Returns the maximum allowed size (in bytes) for a single cached range. + uint64_t GetRangeSizeLimit() const { + return range_size_limit_; + } + + /// Sets the maximum allowed size (in bytes) for a single cached range. + void SetRangeSizeLimit(uint64_t range_size_limit) { + range_size_limit_ = range_size_limit; + } + + /// Returns the maximum gap size (in bytes) considered mergeable between adjacent ranges. + uint64_t GetHoleSizeLimit() const { + return hole_size_limit_; + } + + /// Sets the maximum gap size (in bytes) considered mergeable between adjacent ranges. + void SetHoleSizeLimit(uint64_t hole_size_limit) { + hole_size_limit_ = hole_size_limit; + } + + /// Returns the maximum size to pre-buffer ahead of the current read position. + uint64_t GetPreBufferLimit() const { + return pre_buffer_limit_; + } + + /// Sets the maximum size to pre-buffer ahead of the current read position. + void SetPreBufferLimit(uint64_t pre_buffer_limit) { + pre_buffer_limit_ = pre_buffer_limit; + } + + private: + uint64_t range_size_limit_; + uint64_t hole_size_limit_; + uint64_t pre_buffer_limit_; +}; + +} // namespace paimon diff --git a/include/paimon/utils/special_field_ids.h b/include/paimon/utils/special_field_ids.h index 829f2988..5219d72d 100644 --- a/include/paimon/utils/special_field_ids.h +++ b/include/paimon/utils/special_field_ids.h @@ -42,6 +42,8 @@ class SpecialFieldIds { /// Special field ID reserved for index score. Value: CPP_FIELD_ID_END - 1 inline static constexpr int32_t INDEX_SCORE = CPP_FIELD_ID_END - 1; + /// Special field ID reserved for realtime offset. Value: CPP_FIELD_ID_END - 2 + inline static constexpr int32_t REALTIME_OFFSET = CPP_FIELD_ID_END - 2; /// Lowest field ID reserved for system fields; IDs at or above it are excluded from the /// highest field ID of a schema. Value: INT32_MAX / 2 diff --git a/include/paimon/write_context.h b/include/paimon/write_context.h index fa549eef..a9dd367e 100644 --- a/include/paimon/write_context.h +++ b/include/paimon/write_context.h @@ -218,7 +218,7 @@ class PAIMON_EXPORT WriteContextBuilder { WriteContextBuilder& WithFileSystem(const std::shared_ptr& file_system); /// Enables the real-time write path with the provided shared context. - /// @param realtime_context Non-null context that owns the real-time indexers. + /// @param realtime_context Non-null context that owns the real-time stores. /// @return Reference to this builder for method chaining. WriteContextBuilder& WithRealtimeContext( const std::shared_ptr& realtime_context); diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index b0fe91b0..49871672 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -85,6 +85,7 @@ set(PAIMON_COMMON_SRCS common/global_index/global_indexer_factory.cpp common/io/buffered_input_stream.cpp common/io/byte_array_input_stream.cpp + common/io/byte_array_output_stream.cpp common/io/data_input_stream.cpp common/io/data_output_stream.cpp common/io/memory_segment_output_stream.cpp @@ -157,6 +158,7 @@ set(PAIMON_COMMON_SRCS common/utils/arrow/arrow_output_stream_adapter.cpp common/utils/arrow/arrow_utils.cpp common/utils/arrow/mem_utils.cpp + common/utils/arrow/vector_utils.cpp common/utils/binary_row_partition_computer.cpp common/utils/bit_set.cpp common/utils/bloom_filter.cpp @@ -269,12 +271,15 @@ set(PAIMON_CORE_SRCS core/io/data_file_meta.cpp core/io/data_file_meta_serializer.cpp core/io/data_file_path_factory.cpp + core/io/data_file_index_writer.cpp + core/io/file_index_options.cpp core/io/append_data_file_writer_factory.cpp core/io/blob_data_file_writer_factory.cpp core/io/data_file_writer_factory.cpp core/io/data_file_writer.cpp core/io/field_mapping_reader.cpp core/io/complete_row_tracking_fields_reader.cpp + core/io/vector_file_batch_reader.cpp core/io/file_index_evaluator.cpp core/io/key_value_data_file_record_reader.cpp core/io/key_value_data_file_writer_factory.cpp @@ -373,9 +378,12 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/prepared_key_value_reader.cpp + core/realtime/primary_key_realtime_store.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp + core/realtime/realtime_primary_key_writer.cpp core/postpone/postpone_bucket_writer.cpp core/schema/arrow_schema_validator.cpp core/schema/schema_manager.cpp @@ -575,6 +583,7 @@ if(PAIMON_BUILD_TESTS) common/global_index/rangebitmap/range_bitmap_global_index_test.cpp common/global_index/wrap/file_index_reader_wrapper_test.cpp common/io/byte_array_input_stream_test.cpp + common/io/byte_array_output_stream_test.cpp common/io/data_input_output_stream_test.cpp common/io/buffered_input_stream_test.cpp common/io/memory_segment_output_stream_test.cpp @@ -611,6 +620,7 @@ if(PAIMON_BUILD_TESTS) common/utils/row_range_index_test.cpp common/utils/var_length_int_utils_test.cpp common/utils/arrow/arrow_utils_test.cpp + common/utils/arrow/vector_utils_test.cpp common/utils/arrow/arrow_stream_adapter_test.cpp common/utils/arrow/mem_utils_test.cpp common/utils/arrow/status_utils_test.cpp @@ -644,7 +654,9 @@ if(PAIMON_BUILD_TESTS) common/utils/range_helper_test.cpp common/utils/read_ahead_cache_test.cpp common/io/cache/lru_cache_test.cpp + common/io/cache/cache_manager_test.cpp common/utils/byte_range_combiner_test.cpp + common/utils/saturating_cast_test.cpp common/utils/scope_guard_test.cpp common/utils/sensitive_config_utils_test.cpp common/utils/serialization_utils_test.cpp @@ -675,6 +687,7 @@ if(PAIMON_BUILD_TESTS) add_paimon_test(common_factories_test SOURCES + common/factories/singleton_test.cpp common/factories/factory_creator_test.cpp common/factories/io_hook_test.cpp STATIC_LINK_LIBS @@ -747,7 +760,10 @@ if(PAIMON_BUILD_TESTS) core/io/key_value_in_memory_record_reader_test.cpp core/io/merged_key_value_record_reader_test.cpp core/io/complete_row_tracking_fields_reader_test.cpp + core/io/vector_file_batch_reader_test.cpp core/io/data_file_meta_test.cpp + core/io/data_file_index_writer_test.cpp + core/io/file_index_options_test.cpp core/io/file_index_evaluator_test.cpp core/io/single_file_writer_test.cpp core/io/rolling_blob_file_writer_test.cpp @@ -767,6 +783,7 @@ if(PAIMON_BUILD_TESTS) core/manifest/index_manifest_file_handler_test.cpp core/memory/writer_memory_manager_test.cpp core/realtime/arrow_realtime_store_test.cpp + core/realtime/primary_key_realtime_store_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp @@ -895,6 +912,7 @@ if(PAIMON_BUILD_TESTS) core/utils/file_utils_test.cpp core/utils/manifest_meta_reader_test.cpp core/utils/offset_row_test.cpp + core/utils/partition_utils_test.cpp core/utils/partition_path_utils_test.cpp core/utils/snapshot_manager_test.cpp core/utils/tag_manager_test.cpp diff --git a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp b/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp index b7c84095..b1cdf963 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp @@ -32,7 +32,6 @@ #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" -#include "paimon/core/casting/casting_utils.h" #include "paimon/core/utils/nested_projection_utils.h" namespace paimon { @@ -110,6 +109,56 @@ class SharedSelectedKeysReadPlan : public MapFieldReadPlan { std::vector selected_keys_; }; +Result> MaskSinglePhysicalColumn( + const std::shared_ptr& physical_struct_array, + const std::shared_ptr& field_mapping_array, + const std::shared_ptr& field_mapping_values, + const std::shared_ptr& physical_column_array, int32_t physical_column_id, + int32_t field_id, const std::string& field_name, arrow::MemoryPool* arrow_pool) { + int64_t row_count = physical_struct_array->length(); + if (physical_column_array->length() != row_count) { + return Status::Invalid("shared-shredding physical column length does not match row count"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr validity, + arrow::AllocateEmptyBitmap(row_count, arrow_pool)); + int64_t valid_count = 0; + for (int64_t row = 0; row < row_count; ++row) { + if (physical_struct_array->IsNull(row)) { + continue; + } + if (field_mapping_array->IsNull(row)) { + return Status::Invalid(fmt::format( + "__field_mapping cannot be null in non-null shared-shredding row for field {}", + field_name)); + } + int32_t mapping_offset = field_mapping_array->value_offset(row); + int32_t mapping_length = field_mapping_array->value_length(row); + if (physical_column_id < 0 || physical_column_id >= mapping_length) { + return Status::Invalid("physical column id is out of __field_mapping range"); + } + int32_t mapping_index = mapping_offset + physical_column_id; + if (field_mapping_values->IsNull(mapping_index)) { + return Status::Invalid("__field_mapping element cannot be null"); + } + if (field_mapping_values->Value(mapping_index) != field_id || + physical_column_array->IsNull(row)) { + continue; + } + arrow::bit_util::SetBit(validity->mutable_data(), row); + ++valid_count; + } + + // Replace only the top-level validity; offsets, values, and nested children stay shared. + std::shared_ptr result_data = physical_column_array->data()->Copy(); + if (result_data->buffers.empty()) { + return Status::Invalid("shared-shredding physical column has no validity buffer slot"); + } + int64_t null_count = row_count - valid_count; + result_data->buffers[0] = null_count == 0 ? nullptr : std::move(validity); + result_data->SetNullCount(null_count); + return arrow::MakeArray(std::move(result_data)); +} + class DefaultSelectedKeysReadPlan : public MapFieldReadPlan { public: DefaultSelectedKeysReadPlan(const std::shared_ptr& logical_field, @@ -386,12 +435,10 @@ Result> FullMapReadPlan::Materialize( std::shared_ptr overflow_array; CollectPhysicalColumns(physical_struct_array, &physical_column_name_to_array, &overflow_array); for (auto& [_, physical_column_array] : physical_column_name_to_array) { - if (physical_column_array->type_id() == arrow::Type::DICTIONARY) { - PAIMON_ASSIGN_OR_RAISE( - physical_column_array, - CastingUtils::Cast(physical_column_array, logical_map_type_->item_type(), - arrow::compute::CastOptions::Safe(), arrow_pool)); - } + PAIMON_ASSIGN_OR_RAISE( + physical_column_array, + NestedProjectionUtils::AlignArrayToReadType( + physical_column_array, logical_map_type_->item_type(), arrow_pool)); } std::shared_ptr overflow_keys; @@ -406,12 +453,9 @@ Result> FullMapReadPlan::Materialize( if (!overflow_items) { return Status::Invalid("__overflow map item array is null"); } - if (overflow_items->type_id() == arrow::Type::DICTIONARY) { - PAIMON_ASSIGN_OR_RAISE( - overflow_items, - CastingUtils::Cast(overflow_items, logical_map_type_->item_type(), - arrow::compute::CastOptions::Safe(), arrow_pool)); - } + PAIMON_ASSIGN_OR_RAISE(overflow_items, + NestedProjectionUtils::AlignArrayToReadType( + overflow_items, logical_map_type_->item_type(), arrow_pool)); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::unique_ptr map_builder_base, @@ -530,12 +574,9 @@ Result> SharedSelectedKeysReadPlan::Materialize( std::shared_ptr overflow_array; CollectPhysicalColumns(physical_struct_array, &physical_column_name_to_array, &overflow_array); for (auto& [_, physical_column_array] : physical_column_name_to_array) { - if (physical_column_array->type_id() == arrow::Type::DICTIONARY) { - PAIMON_ASSIGN_OR_RAISE( - physical_column_array, - CastingUtils::Cast(physical_column_array, value_type, - arrow::compute::CastOptions::Safe(), arrow_pool)); - } + PAIMON_ASSIGN_OR_RAISE(physical_column_array, + NestedProjectionUtils::AlignArrayToReadType(physical_column_array, + value_type, arrow_pool)); } std::shared_ptr overflow_keys; @@ -550,65 +591,89 @@ Result> SharedSelectedKeysReadPlan::Materialize( if (!overflow_items) { return Status::Invalid("__overflow map item array is null"); } - if (overflow_items->type_id() == arrow::Type::DICTIONARY) { - PAIMON_ASSIGN_OR_RAISE( - overflow_items, - CastingUtils::Cast(overflow_items, value_type, arrow::compute::CastOptions::Safe(), - arrow_pool)); - } + PAIMON_ASSIGN_OR_RAISE(overflow_items, NestedProjectionUtils::AlignArrayToReadType( + overflow_items, value_type, arrow_pool)); } - std::unique_ptr access_builder_base; - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(access_builder_base, - arrow::MakeBuilder(LogicalField()->type(), arrow_pool)); - if (!access_builder_base || !access_builder_base->type() || - access_builder_base->type()->id() != arrow::Type::STRUCT) { - return Status::Invalid( - fmt::format("selected-key MAP field {} is not a STRUCT", LogicalField()->name())); - } - auto* access_builder = checked_cast(access_builder_base.get()); - PAIMON_RETURN_NOT_OK_FROM_ARROW(access_builder->Reserve(physical_struct_array->length())); - - for (int64_t row = 0; row < physical_struct_array->length(); ++row) { - if (physical_struct_array->IsNull(row)) { - PAIMON_RETURN_NOT_OK_FROM_ARROW(access_builder->AppendNull()); + int64_t row_count = physical_struct_array->length(); + arrow::ArrayVector selected_key_arrays; + selected_key_arrays.reserve(selected_keys_.size()); + for (int32_t key_index = 0; key_index < selected_keys_type->num_fields(); ++key_index) { + const SelectedKey& selected_key = selected_keys_[key_index]; + if (selected_key.field_id < 0) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr null_array, + arrow::MakeArrayOfNull(selected_keys_type->field(key_index)->type(), row_count, + arrow_pool)); + selected_key_arrays.push_back(std::move(null_array)); continue; } - if (field_mapping_array->IsNull(row)) { - return Status::Invalid(fmt::format( - "__field_mapping cannot be null in non-null shared-shredding row for field {}", - LogicalField()->name())); + + if (selected_key.candidate_columns.size() == 1 && !selected_key.may_use_overflow) { + int32_t physical_column_id = selected_key.candidate_columns[0]; + std::string physical_column_name = + MapSharedShreddingDefine::PhysicalColumnName(physical_column_id); + auto physical_column_iter = physical_column_name_to_array.find(physical_column_name); + if (physical_column_iter == physical_column_name_to_array.end()) { + return Status::Invalid( + fmt::format("cannot find selected physical column {} for field {}", + physical_column_name, LogicalField()->name())); + } + const std::shared_ptr& physical_column_array = + physical_column_iter->second; + if (physical_column_array->offset() == 0) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr masked_array, + MaskSinglePhysicalColumn(physical_struct_array, field_mapping_array, + field_mapping_values, physical_column_array, + physical_column_id, selected_key.field_id, + LogicalField()->name(), arrow_pool)); + selected_key_arrays.push_back(std::move(masked_array)); + continue; + } else { + return Status::Invalid("paimon only supports arrays with zero offset"); + } } - int32_t mapping_offset = field_mapping_array->value_offset(row); - PAIMON_RETURN_NOT_OK_FROM_ARROW(access_builder->Append()); - for (int32_t key_index = 0; key_index < selected_keys_type->num_fields(); ++key_index) { - arrow::ArrayBuilder* value_builder = access_builder->field_builder(key_index); - const SelectedKey& selected_key = selected_keys_[key_index]; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::unique_ptr value_builder, + arrow::MakeBuilder(selected_keys_type->field(key_index)->type(), arrow_pool)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Reserve(row_count)); + for (int64_t row = 0; row < row_count; ++row) { + if (physical_struct_array->IsNull(row)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->AppendNull()); + continue; + } + if (field_mapping_array->IsNull(row)) { + return Status::Invalid(fmt::format( + "__field_mapping cannot be null in non-null shared-shredding row for field {}", + LogicalField()->name())); + } + int32_t mapping_offset = field_mapping_array->value_offset(row); bool appended = false; - if (selected_key.field_id >= 0) { - for (int32_t physical_column_id : selected_key.candidate_columns) { - int32_t mapping_index = mapping_offset + physical_column_id; - if (field_mapping_values->IsNull(mapping_index)) { - return Status::Invalid("__field_mapping element cannot be null"); - } - if (field_mapping_values->Value(mapping_index) != selected_key.field_id) { - continue; - } - std::string physical_column_name = - MapSharedShreddingDefine::PhysicalColumnName(physical_column_id); - auto physical_column_iter = - physical_column_name_to_array.find(physical_column_name); - if (physical_column_iter == physical_column_name_to_array.end()) { - return Status::Invalid( - fmt::format("cannot find selected physical column {} for field {}", - physical_column_name, LogicalField()->name())); - } - PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->AppendArraySlice( - *physical_column_iter->second->data(), row, 1)); - appended = true; - break; + for (int32_t physical_column_id : selected_key.candidate_columns) { + int32_t mapping_index = mapping_offset + physical_column_id; + if (field_mapping_values->IsNull(mapping_index)) { + return Status::Invalid("__field_mapping element cannot be null"); + } + if (field_mapping_values->Value(mapping_index) != selected_key.field_id) { + continue; } + std::string physical_column_name = + MapSharedShreddingDefine::PhysicalColumnName(physical_column_id); + auto physical_column_iter = + physical_column_name_to_array.find(physical_column_name); + if (physical_column_iter == physical_column_name_to_array.end()) { + return Status::Invalid( + fmt::format("cannot find selected physical column {} for field {}", + physical_column_name, LogicalField()->name())); + } + const std::shared_ptr& physical_column_array = + physical_column_iter->second; + PAIMON_RETURN_NOT_OK_FROM_ARROW( + value_builder->AppendArraySlice(*physical_column_array->data(), row, 1)); + appended = true; + break; } if (!appended && selected_key.may_use_overflow && overflow_array && @@ -630,9 +695,24 @@ Result> SharedSelectedKeysReadPlan::Materialize( PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->AppendNull()); } } + std::shared_ptr selected_key_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Finish(&selected_key_array)); + selected_key_arrays.push_back(std::move(selected_key_array)); + } + + std::shared_ptr parent_validity; + int64_t parent_null_count = physical_struct_array->null_count(); + if (parent_null_count > 0) { + if (physical_struct_array->offset() == 0) { + parent_validity = physical_struct_array->null_bitmap(); + } else { + return Status::Invalid("paimon only supports arrays with zero offset"); + } } - std::shared_ptr result; - PAIMON_RETURN_NOT_OK_FROM_ARROW(access_builder->Finish(&result)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr result, + arrow::StructArray::Make(selected_key_arrays, selected_keys_type->fields(), + std::move(parent_validity), parent_null_count)); return result; } @@ -646,14 +726,10 @@ Result> DefaultSelectedKeysReadPlan::Materialize( } auto map_array = checked_pointer_cast(physical_array); auto selected_keys_type = checked_pointer_cast(LogicalField()->type()); - auto physical_map_type = checked_pointer_cast(PhysicalReadField()->type()); std::shared_ptr items = map_array->items(); - if (items->type_id() == arrow::Type::DICTIONARY) { - PAIMON_ASSIGN_OR_RAISE(items, - CastingUtils::Cast(items, physical_map_type->item_type(), - arrow::compute::CastOptions::Safe(), arrow_pool)); - } + PAIMON_ASSIGN_OR_RAISE(items, NestedProjectionUtils::AlignArrayToReadType( + items, selected_keys_type->field(0)->type(), arrow_pool)); std::shared_ptr keys = map_array->keys(); std::unique_ptr access_builder_base; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(access_builder_base, diff --git a/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp index 80f5046b..b2dd9374 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp @@ -333,6 +333,96 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjection) { AssertChunkedArrayEquals(expected, actual); } +TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionSharesValueBuffers) { + ASSERT_OK_AND_ASSIGN(auto physical_array, PhysicalArray()); + auto physical_root = checked_pointer_cast(physical_array); + auto physical_tags = + checked_pointer_cast(physical_root->GetFieldByName("tags")); + auto physical_column = + physical_tags->GetFieldByName(MapSharedShreddingDefine::PhysicalColumnName(1)); + + auto selected_type = arrow::struct_( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64()), + arrow::field("e", arrow::int64()), arrow::field("missing", arrow::int64())}); + auto selected_field = arrow::field( + "tags", selected_type, /*nullable=*/true, + arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,b,e,missing"})); + ASSERT_OK_AND_ASSIGN( + auto field_read_plan, + MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan(selected_field, TagsMeta())); + ASSERT_OK_AND_ASSIGN(auto result, + field_read_plan->Materialize(physical_tags, arrow::default_memory_pool())); + auto result_struct = checked_pointer_cast(result); + + auto expected = arrow::ipc::internal::json::ArrayFromJSON(selected_type, R"([ + [10, 20, null, null], + [40, null, null, null], + null, + [80, null, 70, null] + ])") + .ValueOrDie(); + ASSERT_TRUE(expected->Equals(result)) << "Expected:\n" + << expected->ToString() << "\nActual:\n" + << result->ToString(); + ASSERT_EQ(physical_column->data()->buffers[1], result_struct->field(1)->data()->buffers[1]); + ASSERT_EQ(physical_column->data()->buffers[1], result_struct->field(2)->data()->buffers[1]); + ASSERT_NE(physical_column->data()->buffers[0], result_struct->field(1)->data()->buffers[0]); + ASSERT_EQ(physical_tags->data()->buffers[0], result_struct->data()->buffers[0]); +} + +TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionSharesNestedValueBuffers) { + auto item_type = arrow::list(arrow::int64()); + auto logical_schema = + arrow::schema({arrow::field("id", arrow::int32()), + arrow::field("tags", arrow::map(arrow::utf8(), item_type))}); + ASSERT_OK_AND_ASSIGN(auto physical_schema, MapSharedShreddingUtils::LogicalToPhysicalSchema( + logical_schema, {{"tags", 1}})); + auto physical_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(physical_schema->fields()), R"([ + [1, [[0], [1, 2], null]], + [2, [[1], [3, 4, 5], null]], + [3, null], + [4, [[0], null, null]] + ])") + .ValueOrDie(); + auto physical_root = checked_pointer_cast(physical_array); + auto physical_tags = + checked_pointer_cast(physical_root->GetFieldByName("tags")); + auto physical_column = + physical_tags->GetFieldByName(MapSharedShreddingDefine::PhysicalColumnName(0)); + + MapSharedShreddingFieldMeta meta; + meta.name_to_id = {{"a", 0}, {"b", 1}}; + meta.field_to_columns = {{0, {0}}, {1, {0}}}; + meta.num_columns = 1; + meta.max_row_width = 1; + auto selected_type = arrow::struct_({arrow::field("b", item_type)}); + auto selected_field = + arrow::field("tags", selected_type, /*nullable=*/true, + arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"b"})); + ASSERT_OK_AND_ASSIGN( + auto field_read_plan, + MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan(selected_field, meta)); + ASSERT_OK_AND_ASSIGN(auto result, + field_read_plan->Materialize(physical_tags, arrow::default_memory_pool())); + auto result_struct = checked_pointer_cast(result); + auto result_list = checked_pointer_cast(result_struct->field(0)); + + auto expected = arrow::ipc::internal::json::ArrayFromJSON(selected_type, R"([ + [null], + [[3, 4, 5]], + null, + [null] + ])") + .ValueOrDie(); + ASSERT_TRUE(expected->Equals(result)) << "Expected:\n" + << expected->ToString() << "\nActual:\n" + << result->ToString(); + ASSERT_EQ(physical_column->data()->buffers[1], result_list->data()->buffers[1]); + ASSERT_EQ(physical_column->data()->child_data[0]->buffers[1], + result_list->data()->child_data[0]->buffers[1]); +} + TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionFromDefaultMap) { auto map_type = checked_pointer_cast( arrow::map(arrow::utf8(), arrow::field("value", arrow::int64()))); @@ -674,6 +764,66 @@ TEST_F(MapSharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) { AssertChunkedArrayEquals(expected, actual); } +TEST_F(MapSharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringListValue) { + std::shared_ptr logical_schema = arrow::schema({ + arrow::field("id", arrow::int32()), + arrow::field("tags", arrow::map(arrow::utf8(), arrow::list(arrow::utf8()))), + }); + auto options = options_; + std::string format = "orc"; + options[Options::FILE_FORMAT] = format; + options["orc.dictionary-key-size-threshold"] = "1"; + ASSERT_OK_AND_ASSIGN(auto table_schema, + TableSchema::Create(TableSchema::FIRST_SCHEMA_ID, logical_schema, + /*partition_keys=*/{}, /*primary_keys=*/{}, options)); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); + auto path_factory = CreatePathFactory(dir->Str(), format, core_options); + auto compact_manager = std::make_shared(); + ASSERT_OK_AND_ASSIGN( + auto writer, + CreateAppendOnlyWriter(core_options, /*schema_id=*/0, logical_schema, + /*write_cols=*/std::nullopt, + /*max_sequence_number=*/-1, path_factory, compact_manager)); + auto batch = CreateBatch(logical_schema, R"([ + [1, [["a", ["red", "blue"]], ["b", ["blue"]]]], + [2, [["c", ["green"]], ["a", ["red", null, "blue"]], ["b", ["blue"]]]], + [3, null], + [4, [["d", ["yellow"]], ["e", ["blue"]], ["c", [null]], ["a", ["red"]]]] + ])"); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(auto inc, writer->PrepareCommit(/*wait_compaction=*/true)); + ASSERT_OK(writer->Close()); + + std::string data_file_path = + path_factory->ToPath(inc.GetNewFilesIncrement().NewFiles()[0]->file_name); + std::map reader_options = {{"orc.read.enable-lazy-decoding", "true"}}; + auto reader = WrapReader(OpenFormatReader(data_file_path, format, reader_options), + /*selected_keys_str=*/"a,c"); + + auto read_metadata = std::make_shared(); + read_metadata->Append("paimon.map.selected-keys", "a,c"); + arrow::FieldVector read_fields = logical_schema->fields(); + read_fields[1] = read_fields[1]->WithMetadata(read_metadata); + auto read_schema = ExportSchema(arrow::schema(std::move(read_fields))); + ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get())); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_(logical_schema->fields()), {R"([ + [1, [["a", ["red", "blue"]]]], + [2, [["a", ["red", null, "blue"]], ["c", ["green"]]]], + [3, null], + [4, [["a", ["red"]], ["c", [null]]]] + ])"}, + &expected) + .ok()); + AssertChunkedArrayEquals(expected, actual); +} + TEST_F(MapSharedShreddingFileReaderTest, TestReadsRealFormatFile) { // TODO(lisizhuo.lsz): support other format auto options = options_; diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index bac4f16f..36b9c1c9 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -59,6 +59,8 @@ const char Options::SCAN_SNAPSHOT_ID[] = "scan.snapshot-id"; const char Options::SCAN_MODE[] = "scan.mode"; const char Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS[] = "scan.manifest-entry-cache.max-snapshots"; +const char Options::SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED[] = + "scan.manifest-entry.lazy-decode.enabled"; const char Options::READ_BATCH_SIZE[] = "read.batch-size"; const char Options::WRITE_BATCH_SIZE[] = "write.batch-size"; const char Options::WRITE_BUFFER_SIZE[] = "write-buffer-size"; @@ -100,6 +102,7 @@ const char Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_SEQUENCE_GROUP[] = const char Options::SCAN_FALLBACK_BRANCH[] = "scan.fallback-branch"; const char Options::BRANCH[] = "branch"; const char Options::FILE_INDEX_READ_ENABLED[] = "file-index.read.enabled"; +const char Options::FILE_INDEX_IN_MANIFEST_THRESHOLD[] = "file-index.in-manifest-threshold"; const char Options::DATA_FILE_EXTERNAL_PATHS[] = "data-file.external-paths"; const char Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY[] = "data-file.external-paths.strategy"; const char Options::DATA_FILE_PREFIX[] = "data-file.prefix"; @@ -147,7 +150,9 @@ const char Options::AGGREGATION_REMOVE_RECORD_ON_DELETE[] = "aggregation.remove- const char Options::TABLE_READ_SEQUENCE_NUMBER_ENABLED[] = "table-read.sequence-number.enabled"; const char Options::KEY_VALUE_SEQUENCE_NUMBER_ENABLED[] = "key-value.sequence_number.enabled"; const char Options::SCAN_TIMESTAMP_MILLIS[] = "scan.timestamp-millis"; +const char Options::REALTIME_ENABLED[] = "realtime.enabled"; const char Options::REALTIME_READ_VIEW_TTL[] = "realtime.read-view-ttl"; +const char Options::REALTIME_STORE_STATS_MODE[] = "realtime.store.stats-mode"; const char Options::SCAN_TIMESTAMP[] = "scan.timestamp"; const char Options::SCAN_TAG_NAME[] = "scan.tag-name"; const char Options::WRITE_ONLY[] = "write-only"; diff --git a/src/paimon/common/factories/io_hook.cpp b/src/paimon/common/factories/io_hook.cpp index a0576b46..394dc8ea 100644 --- a/src/paimon/common/factories/io_hook.cpp +++ b/src/paimon/common/factories/io_hook.cpp @@ -19,9 +19,12 @@ #include "paimon/common/factories/io_hook.h" #include +#include +#include #include #include "fmt/format.h" +#include "paimon/macros.h" #include "paimon/status.h" namespace paimon { @@ -29,42 +32,66 @@ namespace paimon { class IOHook::Impl { public: Status Try(const std::string& path) { - if (io_count_.fetch_add(1) < pos_.load()) { - return Status::OK(); - } else { - switch (mode_) { - case IOHook::Mode::SILENT: - return Status::OK(); - case IOHook::Mode::RETURN_ERROR: - return Status::IOError(fmt::format( - "io hook triggered io error at position {}, path {}", pos_.load(), path)); - case IOHook::Mode::THROW_EXCEPTION: - throw std::runtime_error(fmt::format( - "io hook throw io exception at position {}, path {}", pos_.load(), path)); - return Status::OK(); - default: - return Status::OK(); - } + // Fast path: the hook is disabled, which is always the case in production; + // writers (Reset()/Clear()) only exist in tests. This keeps Try() a single + // atomic load on the IO path instead of a shared_mutex acquisition per IO. + if (PAIMON_UNLIKELY(armed_.load(std::memory_order_acquire))) { + return TryArmed(path); } + return Status::OK(); } inline void Reset(int64_t pos, IOHook::Mode mode) { + std::unique_lock lock(mutex_); + mode_ = mode; pos_ = pos; io_count_ = 0; - mode_ = mode; + // Arm only after the configuration is complete: TryArmed() reads mode_/pos_ + // under mutex_, which synchronizes with this store, so an observed armed state + // always implies a complete configuration. + armed_.store(true, std::memory_order_release); } int64_t IOCount() const { + std::shared_lock lock(mutex_); return io_count_.load(); } void Clear() { - Reset(-1, IOHook::Mode::SILENT); + std::unique_lock lock(mutex_); + // Disarm first so IO threads stop taking the lock as soon as possible. + armed_.store(false, std::memory_order_release); + mode_ = IOHook::Mode::SILENT; + pos_ = -1; + io_count_ = 0; } private: + Status TryArmed(const std::string& path) { + std::shared_lock lock(mutex_); + if (io_count_.fetch_add(1) < pos_) { + return Status::OK(); + } else { + switch (mode_) { + case IOHook::Mode::SILENT: + return Status::OK(); + case IOHook::Mode::RETURN_ERROR: + return Status::IOError(fmt::format( + "io hook triggered io error at position {}, path {}", pos_, path)); + case IOHook::Mode::THROW_EXCEPTION: + throw std::runtime_error(fmt::format( + "io hook throw io exception at position {}, path {}", pos_, path)); + return Status::OK(); + default: + return Status::OK(); + } + } + } + + mutable std::shared_mutex mutex_; + std::atomic armed_ = {false}; std::atomic io_count_ = {0}; - std::atomic pos_ = {-1}; + int64_t pos_ = -1; IOHook::Mode mode_ = IOHook::Mode::SILENT; }; diff --git a/src/paimon/common/factories/io_hook.h b/src/paimon/common/factories/io_hook.h index e0a2f68b..0c66381b 100644 --- a/src/paimon/common/factories/io_hook.h +++ b/src/paimon/common/factories/io_hook.h @@ -45,7 +45,8 @@ class PAIMON_EXPORT IOHook : public Singleton { }; /// Reset the IO exception position and behavior mode to handle the exception. - /// IOCount will be reset to 0. + /// IOCount will be reset to 0. Arms the hook: Try() switches from its lock-free + /// disabled fast path to the synchronized armed path. /// /// @params pos The position where the IO exception occurs. /// @params mode The mode of behavior for handling the exception. @@ -56,12 +57,14 @@ class PAIMON_EXPORT IOHook : public Singleton { Status Try(const std::string& path); /// Get the count of IO operations that have already occurred. + /// IOs are only counted while the hook is armed (after Reset(), before Clear()); + /// the disabled fast path does not count. /// /// @return The number of IO operations executed. int64_t IOCount() const; /// Clear the state of the IOHook, including resetting IO count and - /// any stored exception state. + /// any stored exception state. Disarms the hook back to the lock-free fast path. void Clear(); private: diff --git a/src/paimon/common/factories/io_hook_test.cpp b/src/paimon/common/factories/io_hook_test.cpp index 9bbb1b34..653dc73b 100644 --- a/src/paimon/common/factories/io_hook_test.cpp +++ b/src/paimon/common/factories/io_hook_test.cpp @@ -19,9 +19,13 @@ #include "paimon/common/factories/io_hook.h" +#include #include +#include +#include #include "gtest/gtest.h" +#include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -64,4 +68,81 @@ TEST(IOHookTest, TestThrowExceptionMode) { hook->Clear(); } +// The disabled state is the production default: Try() must take the lock-free fast +// path, always return OK, and not count IOs (see IOCount()'s contract). Clear() first +// so the test does not depend on execution order. +TEST(IOHookTest, TestDisabledFastPath) { + auto hook = IOHook::GetInstance(); + hook->Clear(); + ASSERT_OK(hook->Try("path")); + ASSERT_OK(hook->Try("path")); + ASSERT_EQ(0, hook->IOCount()); + + // Re-arming and disarming must restore the exact disabled behavior. + hook->Reset(0, IOHook::Mode::RETURN_ERROR); + ASSERT_NOK(hook->Try("path")); + ASSERT_EQ(1, hook->IOCount()); + hook->Clear(); + ASSERT_OK(hook->Try("path")); + ASSERT_OK(hook->Try("path")); + ASSERT_EQ(0, hook->IOCount()); +} + +// Regression test for torn IOHook configurations: Reset()/Clear() run on one thread +// while other threads call Try() concurrently. A shared start barrier releases all +// threads together, and the reset thread keeps hammering until every worker has +// finished, so overlap is structural rather than timing-dependent. The continuous +// arm/disarm cycling also keeps workers switching between the disabled fast path and +// the synchronized armed path. Under a ThreadSanitizer build this deterministically +// reports any unsynchronized access; functionally every Try() must return OK. +TEST(IOHookTest, TestConcurrentResetAndTry) { + auto hook = IOHook::GetInstance(); + + constexpr int32_t kTryIterations = 50000; + constexpr int32_t kNumWorkers = 4; + + std::atomic start{false}; + std::atomic workers_done{0}; + std::atomic observed_error{false}; + + std::thread reset_thread([hook, &start, &workers_done]() { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + while (workers_done.load(std::memory_order_relaxed) < kNumWorkers) { + hook->Reset(INT64_MAX, IOHook::Mode::RETURN_ERROR); + hook->Clear(); + } + }); + + std::vector workers; + workers.reserve(kNumWorkers); + for (int32_t t = 0; t < kNumWorkers; t++) { + workers.emplace_back([hook, &start, &workers_done, &observed_error]() { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (int32_t i = 0; i < kTryIterations; i++) { + Status status = hook->Try("concurrent_path"); + // Reset() arms an unreachable position, while Clear() uses SILENT mode, + // so both complete states return OK. An IOError exposes a torn state. + if (!status.ok()) { + observed_error.store(true, std::memory_order_relaxed); + } + } + workers_done.fetch_add(1, std::memory_order_relaxed); + }); + } + + start.store(true, std::memory_order_release); + reset_thread.join(); + for (auto& worker : workers) { + worker.join(); + } + + ASSERT_FALSE(observed_error.load(std::memory_order_relaxed)); + // Leave the process-wide singleton in its default SILENT state for later tests. + hook->Clear(); +} + } // namespace paimon::test diff --git a/src/paimon/common/factories/singleton.cpp b/src/paimon/common/factories/singleton.cpp index a9732259..2a5daa89 100644 --- a/src/paimon/common/factories/singleton.cpp +++ b/src/paimon/common/factories/singleton.cpp @@ -19,26 +19,14 @@ #include "paimon/factories/singleton.h" -#include - #include "paimon/common/factories/io_hook.h" #include "paimon/factories/factory_creator.h" namespace paimon { -template -T* Singleton::GetInstance() { - static T* ptr; - static std::mutex mutex; - if (PAIMON_UNLIKELY(!ptr)) { - std::lock_guard lg(mutex); - if (!ptr) { - InstPolicy::Create(ptr); - } - } - return const_cast(ptr); -} - +// The single definition point for the two cross-library singletons. See the +// extern template declarations in singleton.h for why implicit instantiation +// must stay suppressed for these types. template class Singleton; template class Singleton; diff --git a/src/paimon/common/factories/singleton_test.cpp b/src/paimon/common/factories/singleton_test.cpp new file mode 100644 index 00000000..efaf921c --- /dev/null +++ b/src/paimon/common/factories/singleton_test.cpp @@ -0,0 +1,108 @@ +/* + * 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 "paimon/factories/singleton.h" + +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +namespace paimon::test { + +namespace { + +constexpr int32_t kNumThreads = 32; + +// Runs `worker(i)` on kNumThreads threads that are all blocked on a shared start +// flag and released at (nearly) the same time, so that they race on the first +// Singleton::GetInstance() publication. Joins all threads before returning. +template +void RunStorm(const Worker& worker) { + std::atomic start{false}; + std::vector threads; + threads.reserve(kNumThreads); + for (int32_t i = 0; i < kNumThreads; ++i) { + threads.emplace_back([&start, &worker, i]() { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + worker(i); + }); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } +} + +// Local to this translation unit, so nothing else in the test binary can have +// instantiated Singleton before this test runs: the storm +// below is guaranteed to race on the *first* publication regardless of link order, +// --gtest_shuffle, or --gtest_filter. GetInstance() is defined in the header, so a +// translation-unit-local type can instantiate it. +class FirstPublicationTarget { + public: + FirstPublicationTarget() { + for (size_t i = 0; i < payload_.size(); ++i) { + payload_[i] = kMagic ^ (i * 0x9E3779B97F4A7C15ULL); + } + } + + // The publication race let a reader observe the instance pointer before the + // constructor's stores were visible; this checks every word the ctor wrote. + bool IsFullyConstructed() const { + for (size_t i = 0; i < payload_.size(); ++i) { + if (payload_[i] != (kMagic ^ (i * 0x9E3779B97F4A7C15ULL))) { + return false; + } + } + return true; + } + + private: + static constexpr uint64_t kMagic = 0xA5A5F00D12345678ULL; + std::array payload_{}; +}; + +} // namespace + +// Regression gate for the Singleton double-checked-locking publication race: 32 +// threads race the first GetInstance() of a type local to this file, so the gate +// cannot silently degrade into exercising only the already-published fast path. +TEST(SingletonTest, TestConcurrentFirstPublication) { + std::array instances{}; + std::array fully_constructed{}; + RunStorm([&instances, &fully_constructed](int32_t i) { + instances[i] = Singleton::GetInstance(); + fully_constructed[i] = instances[i]->IsFullyConstructed(); + }); + + FirstPublicationTarget* expected = instances[0]; + ASSERT_NE(expected, nullptr); + for (int32_t i = 0; i < kNumThreads; ++i) { + ASSERT_EQ(expected, instances[i]); + ASSERT_TRUE(fully_constructed[i]); + } +} + +} // namespace paimon::test diff --git a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp index d068d3c2..1082e669 100644 --- a/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp +++ b/src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp @@ -30,6 +30,7 @@ #include "gtest/gtest.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/mock/mock_file_batch_reader.h" @@ -37,7 +38,6 @@ #include "paimon/testing/mock/mock_format_reader_builder.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" -#include "paimon/utils/read_ahead_cache.h" namespace arrow { class Array; @@ -97,7 +97,7 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test, prefetch_batch_count, batch_size, prefetch_batch_count * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, - /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, CacheConfig(), pool_)); + /*read_ahead_cache_enabled=*/true, CacheConfig(), pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/common/file_index/file_index_format.cpp b/src/paimon/common/file_index/file_index_format.cpp index 85500845..fab5c7a7 100644 --- a/src/paimon/common/file_index/file_index_format.cpp +++ b/src/paimon/common/file_index/file_index_format.cpp @@ -27,7 +27,9 @@ #include "arrow/type.h" #include "fmt/format.h" #include "paimon/common/file_index/empty/empty_file_index_reader.h" +#include "paimon/common/io/data_output_stream.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/file_index/file_indexer.h" #include "paimon/file_index/file_indexer_factory.h" #include "paimon/io/byte_array_input_stream.h" @@ -39,6 +41,128 @@ namespace paimon { class InputStream; class MemoryPool; +class FileIndexFormatWriterImpl : public FileIndexFormat::Writer { + public: + explicit FileIndexFormatWriterImpl(const std::shared_ptr& output_stream) + : output_stream_(output_stream) { + assert(output_stream_); + } + + Status WriteColumnIndexes(const FileIndexFormat::ColumnIndexes& indexes) override { + if (written_) { + return Status::Invalid("File index column indexes have already been written"); + } + + PAIMON_RETURN_NOT_OK(WriteHead(indexes)); + // Write body. + DataOutputStream data_output(output_stream_); + for (const auto& [column_name, column_indexes] : indexes) { + for (const auto& [index_type, bytes] : column_indexes) { + if (bytes) { + PAIMON_RETURN_NOT_OK(data_output.WriteBytes(bytes)); + } + } + } + written_ = true; + return Status::OK(); + } + + Status Close() override { + if (closed_) { + return Status::OK(); + } + closed_ = true; + PAIMON_RETURN_NOT_OK(output_stream_->Flush()); + return output_stream_->Close(); + } + + private: + static constexpr int32_t kRedundantLength = 0; + + static Result CalculateHeadLength(const FileIndexFormat::ColumnIndexes& indexes) { + // magic(8), version(4), header length(4), and column count(4). + int64_t head_length = 8 + 4 + 4 + 4; + int64_t body_length = 0; + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(indexes.size(), "file index column count")); + for (const auto& [column_name, column_indexes] : indexes) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(column_name.size(), + "file index column name length")); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(column_indexes.size(), "column index count")); + // column name(2 + N) + index count(4) + head_length += 2 + static_cast(column_name.size()) + 4; + for (const auto& [index_type, bytes] : column_indexes) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(index_type.size(), + "file index type name length")); + // index type(2 + N) + body offset(4) + body length(4) + head_length += 2 + static_cast(index_type.size()) + 4 + 4; + if (bytes) { + PAIMON_RETURN_NOT_OK(AddChecked(bytes->size(), "index body", &body_length)); + } + } + } + + head_length += 4; // The trailing redundant-length field(4). + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(head_length, "file index header length")); + int64_t container_length = head_length + body_length; + PAIMON_RETURN_NOT_OK(ValidateValueInRange(container_length, "file index size")); + return static_cast(head_length); + } + + Status WriteHead(const FileIndexFormat::ColumnIndexes& indexes) { + PAIMON_ASSIGN_OR_RAISE(int32_t head_length, CalculateHeadLength(indexes)); + DataOutputStream data_output(output_stream_); + // Write magic. + PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::MAGIC)); + // Write version. + PAIMON_RETURN_NOT_OK(data_output.WriteValue(FileIndexFormat::V_1)); + // Write head length. + PAIMON_RETURN_NOT_OK(data_output.WriteValue(head_length)); + // Write column count. + PAIMON_RETURN_NOT_OK(data_output.WriteValue(static_cast(indexes.size()))); + + int64_t body_offset = head_length; + for (const auto& [column_name, column_indexes] : indexes) { + // Write column name. + PAIMON_RETURN_NOT_OK(data_output.WriteString(column_name)); + // Write index count for the column. + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(column_indexes.size()))); + for (const auto& [index_type, bytes] : column_indexes) { + // Write index type. + PAIMON_RETURN_NOT_OK(data_output.WriteString(index_type)); + // Write body offset and length. + if (bytes) { + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(body_offset))); + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(static_cast(bytes->size()))); + body_offset += static_cast(bytes->size()); + } else { + PAIMON_RETURN_NOT_OK( + data_output.WriteValue(FileIndexFormat::EMPTY_INDEX_FLAG)); + PAIMON_RETURN_NOT_OK(data_output.WriteValue(0)); + } + } + } + // Write redundant length for future format extensions. + return data_output.WriteValue(kRedundantLength); + } + + template + static Status AddChecked(T value, const char* name, int64_t* total) { + PAIMON_RETURN_NOT_OK(ValidateValueInRange(value, name)); + *total += static_cast(value); + return ValidateValueInRange(*total, name); + } + + std::shared_ptr output_stream_; + bool written_ = false; + bool closed_ = false; +}; + class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { public: using HeaderType = @@ -153,4 +277,15 @@ Result> FileIndexFormat::CreateReader( const std::shared_ptr& input_stream, const std::shared_ptr& pool) { return FileIndexFormatReaderImpl::Create(input_stream, pool); } + +Result> FileIndexFormat::CreateWriter( + const std::shared_ptr& output_stream, const std::shared_ptr& pool) { + if (!output_stream) { + return Status::Invalid("File index output stream cannot be null"); + } + if (!pool) { + return Status::Invalid("File index memory pool cannot be null"); + } + return std::make_unique(output_stream); +} } // namespace paimon diff --git a/src/paimon/common/file_index/file_index_format_test.cpp b/src/paimon/common/file_index/file_index_format_test.cpp index 7851d57e..40989f7e 100644 --- a/src/paimon/common/file_index/file_index_format_test.cpp +++ b/src/paimon/common/file_index/file_index_format_test.cpp @@ -24,17 +24,20 @@ #include "paimon/common/file_index/bloomfilter/bloom_filter_file_index.h" #include "paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h" #include "paimon/common/file_index/empty/empty_file_index_reader.h" +#include "paimon/common/io/byte_array_output_stream.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/file_index/file_index_result.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/literal.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { + class FileIndexFormatTest : public ::testing::Test { public: void SetUp() override { @@ -55,14 +58,25 @@ class FileIndexFormatTest : public ::testing::Test { std::shared_ptr pool_; }; -TEST_F(FileIndexFormatTest, TestCreateEmptyFileIndexReader) { +TEST_F(FileIndexFormatTest, TestWriteAndReadEmptyIndexGoldenBytes) { + // the expected bytes are generated from Java Paimon + std::vector expected = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, 0, 0, 0, 47, + 0, 0, 0, 1, 0, 2, 99, 49, 0, 0, 0, 1, 0, 5, 101, 109, + 112, 116, 121, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0}; + FileIndexFormat::ColumnIndexes indexes; + indexes["c1"]["empty"] = nullptr; + auto segment_output = std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + auto output = std::make_shared(std::move(segment_output)); + + ASSERT_OK_AND_ASSIGN(auto writer, FileIndexFormat::CreateWriter(output, pool_)); + ASSERT_OK(writer->WriteColumnIndexes(indexes)); + ASSERT_OK(writer->Close()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, output->Finish(pool_.get())); + + ASSERT_EQ(expected, std::vector(actual->data(), actual->data() + actual->size())); auto schema = arrow::schema({arrow::field("c1", arrow::utf8())}); - std::vector index_file_bytes = {0, 5, 78, 78, -48, 26, 53, -82, 0, 0, 0, 1, - 0, 0, 0, 47, 0, 0, 0, 1, 0, 2, 99, 49, - 0, 0, 0, 1, 0, 5, 101, 109, 112, 116, 121, -1, - -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0}; - auto input_stream = - std::make_shared(index_file_bytes.data(), index_file_bytes.size()); + auto input_stream = std::make_shared(actual->data(), actual->size()); ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input_stream, pool_)); ASSERT_OK_AND_ASSIGN(auto index_file_readers, reader->ReadColumnIndex("c1", CreateArrowSchema(schema).get())); diff --git a/src/paimon/common/io/byte_array_output_stream.cpp b/src/paimon/common/io/byte_array_output_stream.cpp new file mode 100644 index 00000000..bc6d1f1a --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream.cpp @@ -0,0 +1,80 @@ +/* + * 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 "paimon/common/io/byte_array_output_stream.h" + +#include +#include +#include +#include +#include + +#include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/utils/math.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon { + +ByteArrayOutputStream::ByteArrayOutputStream(std::unique_ptr&& output) + : output_(std::move(output)) { + assert(output_); +} + +Result ByteArrayOutputStream::Write(const char* buffer, int64_t size) { + if (closed_) { + return Status::Invalid("Byte array output stream is closed"); + } + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(size, "write length")); + if (buffer == nullptr && size > 0) { + return Status::Invalid("Write buffer must not be null when size is positive"); + } + int64_t remaining = size; + while (remaining > 0) { + uint32_t to_write = static_cast(std::min( + remaining, static_cast(std::numeric_limits::max()))); + output_->Write(buffer, to_write); + buffer += to_write; + remaining -= to_write; + } + return size; +} + +Status ByteArrayOutputStream::Close() { + closed_ = true; + return Status::OK(); +} + +Result> ByteArrayOutputStream::Finish(MemoryPool* pool) { + assert(pool); + PAIMON_RETURN_NOT_OK(Close()); + if (result_) { + return result_; + } + // TODO(jinli.zjw): Support int64_t lengths in MemorySegmentUtils::CopyToBytes and remove this + // limit. + const int64_t size = output_->CurrentSize(); + PAIMON_RETURN_NOT_OK(ValidateValueInRange(size, "byte array output stream size")); + const std::vector& segments = output_->Segments(); + result_ = std::make_shared(static_cast(size), pool); + MemorySegmentUtils::CopyToBytes(segments, /*offset=*/0, result_.get(), + /*bytes_offset=*/0, static_cast(size)); + return result_; +} + +} // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream.h b/src/paimon/common/io/byte_array_output_stream.h new file mode 100644 index 00000000..9b87ca42 --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream.h @@ -0,0 +1,69 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/fs/file_system.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon { + +class Bytes; +class MemoryPool; + +/// An in-memory output stream backed by segments allocated from a Paimon MemoryPool. +class ByteArrayOutputStream : public OutputStream { + public: + /// Takes ownership of an initialized segmented output stream. + explicit ByteArrayOutputStream(std::unique_ptr&& output); + + ~ByteArrayOutputStream() override = default; + + Result Write(const char* buffer, int64_t size) override; + + Status Flush() override { + return Status::OK(); + } + + Result GetPos() const override { + return output_->CurrentSize(); + } + + Result GetUri() const override { + return std::string(); + } + + Status Close() override; + + /// Closes the stream and returns its contents as an exactly-sized contiguous byte array. + /// @note The caller must keep `pool` alive until the returned bytes are destroyed. + Result> Finish(MemoryPool* pool); + + private: + std::unique_ptr output_; + std::shared_ptr result_; + bool closed_ = false; +}; + +} // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream_test.cpp b/src/paimon/common/io/byte_array_output_stream_test.cpp new file mode 100644 index 00000000..bd185095 --- /dev/null +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -0,0 +1,86 @@ +/* + * 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 "paimon/common/io/byte_array_output_stream.h" + +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { + std::shared_ptr pool = GetMemoryPool(); + auto output = std::make_unique(/*segment_size=*/2, pool); + std::shared_ptr stream = + std::make_shared(std::move(output)); + ASSERT_GT(pool->CurrentUsage(), 0); + ASSERT_OK_AND_ASSIGN(int64_t first_write, stream->Write("ab", 2)); + ASSERT_EQ(2, first_write); + ASSERT_OK_AND_ASSIGN(int64_t second_write, stream->Write("cdef", 4)); + ASSERT_EQ(4, second_write); + ASSERT_OK_AND_ASSIGN(int64_t position, stream->GetPos()); + ASSERT_EQ(6, position); + ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + ASSERT_EQ("abcdef", std::string(result->data(), result->size())); + ASSERT_NOK_WITH_MSG(stream->Write("x", 1), "closed"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated, stream->Finish(pool.get())); + ASSERT_EQ(result, repeated); + stream.reset(); + ASSERT_EQ(6, pool->CurrentUsage()); +} + +TEST(ByteArrayOutputStreamTest, TestWriteValidation) { + std::shared_ptr pool = GetDefaultPool(); + auto output = std::make_unique(/*segment_size=*/8, pool); + std::shared_ptr stream = + std::make_shared(std::move(output)); + ASSERT_NOK(stream->Write(nullptr, 1)); + ASSERT_NOK(stream->Write("", -1)); + ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write(nullptr, 0)); + ASSERT_EQ(0, written); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + ASSERT_EQ(0, result->size()); +} + +TEST(ByteArrayOutputStreamTest, TestCallerKeepsMemoryPoolAlive) { + std::shared_ptr pool = GetMemoryPool(); + auto output = std::make_unique(/*segment_size=*/8, pool); + std::shared_ptr stream = + std::make_shared(std::move(output)); + ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write("data", 4)); + ASSERT_EQ(4, written); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + + stream.reset(); + ASSERT_GT(pool->CurrentUsage(), 0); + ASSERT_EQ("data", std::string(result->data(), result->size())); + + result.reset(); + ASSERT_EQ(0, pool->CurrentUsage()); +} + +} // namespace paimon::test diff --git a/src/paimon/common/io/cache/cache_manager.h b/src/paimon/common/io/cache/cache_manager.h index f899d46c..6fafefb5 100644 --- a/src/paimon/common/io/cache/cache_manager.h +++ b/src/paimon/common/io/cache/cache_manager.h @@ -25,6 +25,7 @@ #include "paimon/cache/cache.h" #include "paimon/common/io/cache/cache_key.h" #include "paimon/common/io/cache/lru_cache.h" +#include "paimon/common/utils/saturating_cast.h" #include "paimon/memory/memory_segment.h" #include "paimon/result.h" @@ -59,9 +60,13 @@ class PAIMON_EXPORT CacheManager { /// @param high_priority_pool_ratio Ratio of capacity reserved for index cache [0.0, 1.0). /// If 0, index and data share the same cache. CacheManager(int64_t max_memory_bytes, double high_priority_pool_ratio) { - auto index_cache_bytes = static_cast(max_memory_bytes * high_priority_pool_ratio); + // Both factors are config-validated non-negative values, so the products are finite; + // saturation is only a defense against the undefined double->int64_t conversion when + // max_memory_bytes is close enough to INT64_MAX that the product rounds to 2^63. + auto index_cache_bytes = + SaturatingDoubleToInteger(max_memory_bytes * high_priority_pool_ratio); auto data_cache_bytes = - static_cast(max_memory_bytes * (1.0 - high_priority_pool_ratio)); + SaturatingDoubleToInteger(max_memory_bytes * (1.0 - high_priority_pool_ratio)); data_cache_ = std::make_shared(data_cache_bytes); if (high_priority_pool_ratio == 0.0) { index_cache_ = data_cache_; diff --git a/src/paimon/common/io/cache/cache_manager_test.cpp b/src/paimon/common/io/cache/cache_manager_test.cpp new file mode 100644 index 00000000..4d7c180a --- /dev/null +++ b/src/paimon/common/io/cache/cache_manager_test.cpp @@ -0,0 +1,155 @@ +/* + * 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 "paimon/common/io/cache/cache_manager.h" + +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/cache/cache.h" +#include "paimon/common/io/cache/cache_key.h" +#include "paimon/common/io/cache/lru_cache.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/memory/memory_segment.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class CacheManagerTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + } + + std::shared_ptr MakeKey(int64_t position, bool is_index = false) const { + return CacheKey::ForPosition("test_file", position, 64, is_index); + } + + MemorySegment MakeSegment(int32_t size, char fill_byte) const { + auto segment = MemorySegment::AllocateHeapMemory(size, pool_.get()); + std::memset(segment.MutableData(), fill_byte, size); + return segment; + } + + std::shared_ptr DataLru(const CacheManager& manager) const { + return std::dynamic_pointer_cast(manager.DataCache()); + } + + std::shared_ptr IndexLru(const CacheManager& manager) const { + return std::dynamic_pointer_cast(manager.IndexCache()); + } + + private: + std::shared_ptr pool_; +}; + +/// Regression test for the double->int64_t conversions in the CacheManager constructor: +/// (double)INT64_MAX rounds to 2^63, which is not representable as int64_t, so casting the +/// product back is undefined behavior (x86 cvttsd2si yields INT64_MIN, aarch64 fcvtzs +/// saturates to INT64_MAX). The conversion must saturate, keeping the capacity non-negative. +TEST_F(CacheManagerTest, TestCapacitySaturatesAtInt64Max) { + CacheManager manager(std::numeric_limits::max(), /*high_priority_pool_ratio=*/0.0); + + std::shared_ptr data_lru = DataLru(manager); + ASSERT_NE(data_lru, nullptr); + ASSERT_GE(data_lru->GetMaxWeight(), 0); + ASSERT_EQ(data_lru->GetMaxWeight(), std::numeric_limits::max()); + + // A ratio of 0.0 means index and data share the same cache. + ASSERT_EQ(manager.DataCache(), manager.IndexCache()); + + // The saturated capacity accepts entries instead of rejecting every insert. + std::shared_ptr key = MakeKey(0); + auto reader = [&](const std::shared_ptr&) -> Result { + return MakeSegment(64, 'A'); + }; + ASSERT_OK_AND_ASSIGN(MemorySegment segment, manager.GetPage(key, reader, {})); + ASSERT_EQ(segment.Size(), 64); + ASSERT_EQ(segment.Get(0), 'A'); +} + +/// Verifies the exact capacity split between the data and index caches for a normal +/// configuration, plus a Get/Invalidate smoke path through CacheManager::GetPage. +TEST_F(CacheManagerTest, TestNormalSplitAndSmokePath) { + CacheManager manager(/*max_memory_bytes=*/1024, /*high_priority_pool_ratio=*/0.5); + + std::shared_ptr data_lru = DataLru(manager); + std::shared_ptr index_lru = IndexLru(manager); + ASSERT_NE(data_lru, nullptr); + ASSERT_NE(index_lru, nullptr); + ASSERT_EQ(data_lru->GetMaxWeight(), 512); + ASSERT_EQ(index_lru->GetMaxWeight(), 512); + + std::shared_ptr key = MakeKey(0); + int32_t reader_calls = 0; + auto reader = [&](const std::shared_ptr&) -> Result { + reader_calls++; + return MakeSegment(128, 'B'); + }; + + // The first GetPage is a miss and invokes the reader; the second is a cache hit. + ASSERT_OK_AND_ASSIGN(MemorySegment first, manager.GetPage(key, reader, {})); + ASSERT_EQ(first.Get(0), 'B'); + ASSERT_EQ(reader_calls, 1); + ASSERT_OK_AND_ASSIGN(MemorySegment second, manager.GetPage(key, reader, {})); + ASSERT_EQ(second.Get(0), 'B'); + ASSERT_EQ(reader_calls, 1); + + // After InvalidPage the reader is invoked again. + manager.InvalidPage(key); + ASSERT_OK_AND_ASSIGN(MemorySegment third, manager.GetPage(key, reader, {})); + ASSERT_EQ(third.Get(0), 'B'); + ASSERT_EQ(reader_calls, 2); +} + +/// Verifies weight-based eviction through GetPage: inserting beyond the data cache capacity +/// evicts the least recently used page and runs its eviction callback. +TEST_F(CacheManagerTest, TestGetPageEviction) { + // The data cache capacity is 512 * (1.0 - 0.5) = 256 bytes. + CacheManager manager(/*max_memory_bytes=*/512, /*high_priority_pool_ratio=*/0.5); + + std::vector evicted; + auto callback_for = [&evicted](int64_t position) -> CacheCallback { + return + [&evicted, position](const std::shared_ptr&) { evicted.push_back(position); }; + }; + auto reader = [&](const std::shared_ptr&) -> Result { + return MakeSegment(128, 'C'); + }; + + std::shared_ptr key0 = MakeKey(0); + std::shared_ptr key1 = MakeKey(1); + std::shared_ptr key2 = MakeKey(2); + ASSERT_OK_AND_ASSIGN(MemorySegment segment0, manager.GetPage(key0, reader, callback_for(0))); + ASSERT_EQ(segment0.Get(0), 'C'); + ASSERT_OK_AND_ASSIGN(MemorySegment segment1, manager.GetPage(key1, reader, callback_for(1))); + ASSERT_EQ(segment1.Get(0), 'C'); + ASSERT_TRUE(evicted.empty()); + + // 128 + 128 + 128 > 256: inserting key2 evicts key0, the least recently used page. + ASSERT_OK_AND_ASSIGN(MemorySegment segment2, manager.GetPage(key2, reader, callback_for(2))); + ASSERT_EQ(segment2.Get(0), 'C'); + ASSERT_EQ(evicted, std::vector({0})); + ASSERT_EQ(manager.DataCache()->Size(), 2); +} + +} // namespace paimon::test diff --git a/src/paimon/common/io/cache_input_stream.h b/src/paimon/common/io/cache_input_stream.h index 9ccbf260..015b082c 100644 --- a/src/paimon/common/io/cache_input_stream.h +++ b/src/paimon/common/io/cache_input_stream.h @@ -18,13 +18,12 @@ #pragma once -#include #include #include #include "paimon/common/utils/math.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/fs/file_system.h" -#include "paimon/utils/read_ahead_cache.h" namespace paimon { @@ -48,10 +47,9 @@ class CacheInputStream : public InputStream { PAIMON_RETURN_NOT_OK(ValidateValueInRange(offset, "read offset")); PAIMON_RETURN_NOT_OK(ValidateValueInRange(size, "read size")); ByteRange range{static_cast(offset), static_cast(size)}; - PAIMON_ASSIGN_OR_RAISE(ByteSlice slice, cache_->Read(range)); - if (slice.buffer) { - std::memcpy(buffer, slice.buffer->data() + slice.offset, slice.length); - return slice.length; + PAIMON_ASSIGN_OR_RAISE(bool hit, cache_->Read(range, buffer)); + if (hit) { + return size; } } return input_stream_->Read(buffer, size, offset); @@ -70,14 +68,12 @@ class CacheInputStream : public InputStream { return; } ByteRange range{static_cast(offset), static_cast(size)}; - Result slice = cache_->Read(range); - if (!slice.ok()) { - callback(slice.status()); + Result hit = cache_->Read(range, buffer); + if (!hit.ok()) { + callback(hit.status()); return; } - if (slice.value().buffer) { - std::memcpy(buffer, slice.value().buffer->data() + slice.value().offset, - slice.value().length); + if (hit.value()) { callback(Status::OK()); return; } diff --git a/src/paimon/common/io/cache_input_stream_test.cpp b/src/paimon/common/io/cache_input_stream_test.cpp index d4a61854..0b14dc33 100644 --- a/src/paimon/common/io/cache_input_stream_test.cpp +++ b/src/paimon/common/io/cache_input_stream_test.cpp @@ -26,6 +26,7 @@ #include "gtest/gtest.h" #include "paimon/common/factories/io_hook.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/fs/file_system.h" #include "paimon/fs/file_system_factory.h" @@ -33,7 +34,6 @@ #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/testharness.h" -#include "paimon/utils/read_ahead_cache.h" namespace paimon::test { @@ -60,7 +60,7 @@ class CacheInputStreamTest : public ::testing::Test { std::shared_ptr CreateCache(std::vector ranges) { auto stream = OpenFile(); - CacheConfig config(/*buffer_size_limit=*/1024 * 1024, /*range_size_limit=*/1024, + CacheConfig config(/*range_size_limit=*/1024, /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024 * 1024); auto cache = std::make_shared(std::move(stream), config, pool_); EXPECT_OK(cache->Init(std::move(ranges))); @@ -204,7 +204,7 @@ TEST_F(CacheInputStreamTest, TestReadAsyncCacheReadError) { ASSERT_OK_AND_ASSIGN(auto fs, FileSystemFactory::Get("local", file_path_, {})); ASSERT_OK_AND_ASSIGN(auto cache_stream, fs->Open(file_path_)); ASSERT_OK_AND_ASSIGN(auto underlying, fs->Open(file_path_)); - CacheConfig config(/*buffer_size_limit=*/1024 * 1024, /*range_size_limit=*/1024, + CacheConfig config(/*range_size_limit=*/1024, /*hole_size_limit=*/0, /*pre_buffer_limit=*/1024 * 1024); auto cache = std::make_shared(std::move(cache_stream), config, pool_); ASSERT_OK(cache->Init(std::vector{{0, 10}})); diff --git a/src/paimon/common/io/data_input_output_stream_test.cpp b/src/paimon/common/io/data_input_output_stream_test.cpp index 4e606370..0a5dd574 100644 --- a/src/paimon/common/io/data_input_output_stream_test.cpp +++ b/src/paimon/common/io/data_input_output_stream_test.cpp @@ -79,12 +79,8 @@ class DataInputOutputStreamTest : public ::testing::Test, (void)data_output_stream->WriteValue(static_cast(9223372036854775805)); // 8 bytes (void)data_output_stream->WriteValue(true); // 1 byte std::string str1 = "This is a very very very long sentence."; - if constexpr (std::is_same_v) { - (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len - } else { - (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len - } - std::string str2 = "我是一个粉刷匠～"; // 24 bytes + (void)data_output_stream->WriteString(str1); // 39 bytes + 2 bytes len + std::string str2 = "我是一个粉刷匠～"; // 24 bytes auto bytes = std::make_shared(str2, pool_.get()); (void)data_output_stream->WriteBytes(bytes); } diff --git a/src/paimon/common/io/memory_segment_output_stream.cpp b/src/paimon/common/io/memory_segment_output_stream.cpp index 5355f72b..2d0d274a 100644 --- a/src/paimon/common/io/memory_segment_output_stream.cpp +++ b/src/paimon/common/io/memory_segment_output_stream.cpp @@ -54,11 +54,7 @@ void MemorySegmentOutputStream::WriteString(const std::string& str) { } void MemorySegmentOutputStream::Write(const char* data, uint32_t size) { - auto bytes = std::make_shared(size, pool_.get()); - if (size != 0) { - memcpy(bytes->data(), data, size); - } - auto segment = MemorySegment::Wrap(bytes); + MemorySegment segment = MemorySegment::WrapView(data, size); Write(segment, 0, segment.Size()); } diff --git a/src/paimon/common/io/memory_segment_output_stream_test.cpp b/src/paimon/common/io/memory_segment_output_stream_test.cpp index 61fbfe30..69c207ce 100644 --- a/src/paimon/common/io/memory_segment_output_stream_test.cpp +++ b/src/paimon/common/io/memory_segment_output_stream_test.cpp @@ -82,4 +82,16 @@ TEST_P(MemorySegmentOutputStreamTest, TestSimple) { ASSERT_EQ(out.CurrentSize(), input_stream->GetPos().value()); } +TEST(MemorySegmentOutputStreamTest, TestRawWriteDoesNotAllocateTemporaryBuffer) { + std::shared_ptr pool = GetMemoryPool(); + MemorySegmentOutputStream out(/*segment_size=*/8, pool); + uint64_t allocated_before_write = pool->CurrentUsage(); + + out.Write("abc", 3); + + ASSERT_EQ(allocated_before_write, pool->CurrentUsage()); + ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); + ASSERT_EQ(3, out.CurrentSize()); +} + } // namespace paimon::test diff --git a/src/paimon/common/predicate/literal_converter.cpp b/src/paimon/common/predicate/literal_converter.cpp index 19410768..102348d7 100644 --- a/src/paimon/common/predicate/literal_converter.cpp +++ b/src/paimon/common/predicate/literal_converter.cpp @@ -166,6 +166,7 @@ Result LiteralConverter::ConvertLiteralsFromRow( case FieldType::DATE: return Literal(FieldType::DATE, row.GetInt(field_idx)); case FieldType::ARRAY: + case FieldType::VECTOR: case FieldType::MAP: case FieldType::STRUCT: default: diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index c4465179..12e38966 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -31,10 +31,10 @@ #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/format/reader_builder.h" #include "paimon/fs/file_system.h" -#include "paimon/utils/read_ahead_cache.h" namespace arrow { class Schema; @@ -60,7 +60,7 @@ Result> PrefetchFileBatchReaderImpl const std::shared_ptr& fs, uint32_t prefetch_max_parallel_num, int32_t batch_size, uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, bool initialize_read_ranges, - PrefetchCacheMode prefetch_cache_mode, const CacheConfig& cache_config, + bool read_ahead_cache_enabled, const CacheConfig& cache_config, const std::shared_ptr& pool) { if (prefetch_max_parallel_num == 0) { return Status::Invalid("prefetch max parallel num should be greater than 0."); @@ -82,7 +82,7 @@ Result> PrefetchFileBatchReaderImpl } std::shared_ptr cache; - if (prefetch_cache_mode != PrefetchCacheMode::NEVER) { + if (read_ahead_cache_enabled) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, fs->Open(FileStatus(data_file_path, data_file_size))); cache = std::make_shared(input_stream, cache_config, pool); @@ -119,9 +119,9 @@ Result> PrefetchFileBatchReaderImpl } uint32_t prefetch_queue_capacity = prefetch_batch_count / readers.size(); - auto reader = std::unique_ptr(new PrefetchFileBatchReaderImpl( - readers, batch_size, prefetch_queue_capacity, enable_adaptive_prefetch_strategy, executor, - cache, prefetch_cache_mode)); + auto reader = std::unique_ptr( + new PrefetchFileBatchReaderImpl(readers, batch_size, prefetch_queue_capacity, + enable_adaptive_prefetch_strategy, executor, cache)); if (initialize_read_ranges) { // normally initialize read ranges should be false, as set read schema will refresh read // ranges, and set read schema will always be called before read. @@ -133,13 +133,11 @@ Result> PrefetchFileBatchReaderImpl PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl( const std::vector>& readers, int32_t batch_size, uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, - const std::shared_ptr& executor, const std::shared_ptr& cache, - PrefetchCacheMode cache_mode) + const std::shared_ptr& executor, const std::shared_ptr& cache) : readers_(std::move(readers)), batch_size_(batch_size), executor_(executor), cache_(cache), - cache_mode_(cache_mode), prefetch_queue_capacity_(prefetch_queue_capacity), enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy) { for (size_t i = 0; i < readers_.size(); i++) { @@ -158,6 +156,9 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( ::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) { PAIMON_RETURN_NOT_OK(CleanUp()); + if (cache_) { + cache_->Reset(); + } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, arrow::ImportSchema(read_schema)); for (const auto& reader : readers_) { @@ -172,6 +173,9 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( Status PrefetchFileBatchReaderImpl::RefreshReadRanges() { PAIMON_RETURN_NOT_OK(CleanUp()); + if (cache_) { + cache_->Reset(); + } return RefreshReadRangesAfterCleanUp(); } @@ -294,35 +298,14 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { reader_is_working_[i] = false; } is_shutdown_ = false; - if (cache_) { - cache_->Reset(); - } SetReadStatus(Status::OK()); return Status::OK(); } -bool PrefetchFileBatchReaderImpl::NeedInitCache() const { - switch (cache_mode_) { - case PrefetchCacheMode::NEVER: - return false; - case PrefetchCacheMode::EXCLUDE_PREDICATE: - return predicate_ == nullptr; - case PrefetchCacheMode::EXCLUDE_BITMAP: - return selection_bitmap_ == std::nullopt; - case PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE: - return predicate_ == nullptr && selection_bitmap_ == std::nullopt; - case PrefetchCacheMode::ALWAYS: - return true; - default: - assert(false); - return true; - } -} - void PrefetchFileBatchReaderImpl::Workloop() { std::vector> futures; futures.resize(readers_.size()); - if (cache_ && NeedInitCache()) { + if (cache_) { auto read_ranges = readers_[0]->PreBufferRange(); if (read_ranges.ok()) { std::vector ranges; @@ -332,6 +315,11 @@ void PrefetchFileBatchReaderImpl::Workloop() { auto s = cache_->Init(std::move(ranges)); if (!s.ok()) { SetReadStatus(s); + } else { + // Init() only registers the ranges, so without this the first + // cache fetch races the readers' first reads instead of running + // ahead of them. + cache_->Warmup(); } } else { SetReadStatus(read_ranges.status()); @@ -622,7 +610,15 @@ Status PrefetchFileBatchReaderImpl::SeekToRow(uint64_t row_number) { } std::shared_ptr PrefetchFileBatchReaderImpl::GetReaderMetrics() const { - return MetricsImpl::CollectReadMetrics(readers_); + auto res_metrics = MetricsImpl::CollectReadMetrics(readers_); + if (cache_) { + // The shared read-ahead cache serves reads of all sub-readers, so its + // hit/miss counters are file-level and merge into the reader metrics. + std::shared_ptr cache_metrics = std::make_shared(); + cache_->CollectMetrics(&cache_metrics); + res_metrics->Merge(cache_metrics); + } + return res_metrics; } Result> PrefetchFileBatchReaderImpl::GetFileSchema() const { @@ -676,7 +672,17 @@ Result> PrefetchFileBatchReaderImpl::EofRange() co } void PrefetchFileBatchReaderImpl::Close() { + // CleanUp() no longer resets the read-ahead cache: ConcatBatchReader closes file readers as + // soon as they reach EOF, and the cache hit/miss counters must remain readable through + // GetReaderMetrics() after that. The cache is reset only when the reader is reused via + // SetReadSchema()/RefreshReadRanges(). (void)CleanUp(); + if (cache_) { + // Free the prefetched buffers of this file right away (ConcatBatchReader keeps + // closed file readers alive until the whole scan finishes), but keep the + // counters for GetReaderMetrics(). + cache_->ReleaseBuffers(); + } for (const auto& reader : readers_) { reader->Close(); } diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h index 78cfbb5f..c21856d0 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -35,12 +35,12 @@ #include #include "arrow/c/abi.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/common/utils/threadsafe_queue.h" #include "paimon/reader/batch_reader.h" #include "paimon/reader/prefetch_file_batch_reader.h" #include "paimon/result.h" #include "paimon/status.h" -#include "paimon/utils/read_ahead_cache.h" #include "paimon/utils/roaring_bitmap32.h" struct ArrowSchema; @@ -60,8 +60,8 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { const ReaderBuilder* reader_builder, const std::shared_ptr& fs, uint32_t prefetch_max_parallel_num, int32_t batch_size, uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy, const std::shared_ptr& executor, - bool initialize_read_ranges, PrefetchCacheMode prefetch_cache_mode, - const CacheConfig& cache_config, const std::shared_ptr& pool); + bool initialize_read_ranges, bool read_ahead_cache_enabled, const CacheConfig& cache_config, + const std::shared_ptr& pool); ~PrefetchFileBatchReaderImpl() override; @@ -113,8 +113,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { PrefetchFileBatchReaderImpl( const std::vector>& readers, int32_t batch_size, uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy, - const std::shared_ptr& executor, const std::shared_ptr& cache, - PrefetchCacheMode cache_mode); + const std::shared_ptr& executor, const std::shared_ptr& cache); Status CleanUp(); void Workloop(); @@ -143,7 +142,6 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { const std::pair& read_range) const; Status HandleReadResult(size_t reader_idx, const std::pair& read_range, FileBatchReader::ReadBatchWithBitmap&& read_batch_with_bitmap); - bool NeedInitCache() const; private: std::vector> readers_; @@ -162,7 +160,6 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { std::condition_variable cv_; std::shared_ptr executor_; std::shared_ptr cache_; - PrefetchCacheMode cache_mode_; mutable std::shared_mutex rw_mutex_; std::unique_ptr background_thread_; diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp index b3828f7e..192028ac 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp @@ -27,6 +27,7 @@ #include "gtest/gtest.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/executor.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" @@ -39,7 +40,6 @@ #include "paimon/testing/mock/mock_format_reader_builder.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" -#include "paimon/utils/read_ahead_cache.h" namespace paimon::test { @@ -115,7 +115,7 @@ class ControlledMockFormatReaderBuilder : public ReaderBuilder { struct TestParam { std::string file_format; - PrefetchCacheMode cache_mode; + bool read_ahead_cache_enabled; }; class PrefetchFileBatchReaderImplTest : public ::testing::Test, @@ -194,7 +194,7 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, const std::string& file_format_str, const arrow::Schema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap, int32_t batch_size, - int32_t prefetch_max_parallel_num, PrefetchCacheMode cache_mode) const { + int32_t prefetch_max_parallel_num, bool read_ahead_cache_enabled) const { EXPECT_OK_AND_ASSIGN(std::unique_ptr file_format, FileFormatFactory::Get(file_format_str, {})); EXPECT_OK_AND_ASSIGN(auto reader_builder, file_format->CreateReaderBuilder(batch_size)); @@ -209,7 +209,8 @@ class PrefetchFileBatchReaderImplTest : public ::testing::Test, data_file_path, data_file_status.GetLen(), reader_builder.get(), local_fs_, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor, - /*initialize_read_ranges=*/false, cache_mode, CacheConfig(), GetDefaultPool())); + /*initialize_read_ranges=*/false, read_ahead_cache_enabled, CacheConfig(), + GetDefaultPool())); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); EXPECT_TRUE(arrow_status.ok()); @@ -275,18 +276,11 @@ CollectResultAndRowIds(FileBatchReader* reader) { } std::vector PrepareTestParam() { - std::vector values = { - TestParam{"parquet", PrefetchCacheMode::ALWAYS}, - TestParam{"parquet", PrefetchCacheMode::EXCLUDE_BITMAP}, - TestParam{"parquet", PrefetchCacheMode::EXCLUDE_PREDICATE}, - TestParam{"parquet", PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE}, - TestParam{"parquet", PrefetchCacheMode::NEVER}}; + std::vector values = {TestParam{"parquet", /*read_ahead_cache_enabled=*/true}, + TestParam{"parquet", /*read_ahead_cache_enabled=*/false}}; #ifdef PAIMON_ENABLE_ORC - values.emplace_back(TestParam{"orc", PrefetchCacheMode::ALWAYS}); - values.emplace_back(TestParam{"orc", PrefetchCacheMode::EXCLUDE_BITMAP}); - values.emplace_back(TestParam{"orc", PrefetchCacheMode::EXCLUDE_PREDICATE}); - values.emplace_back(TestParam{"orc", PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE}); - values.emplace_back(TestParam{"orc", PrefetchCacheMode::NEVER}); + values.emplace_back(TestParam{"orc", /*read_ahead_cache_enabled=*/true}); + values.emplace_back(TestParam{"orc", /*read_ahead_cache_enabled=*/false}); #endif return values; } @@ -300,13 +294,12 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) { for (auto prefetch_max_parallel_num : {1, 2, 3, 5, 8, 10}) { MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + auto reader, PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); @@ -323,14 +316,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { int32_t prefetch_max_parallel_num = 12; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); // simulate read limits, only read 8 batches for (int32_t i = 0; i < 8; i++) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, @@ -353,14 +345,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { int32_t prefetch_max_parallel_num = 12; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); @@ -430,14 +421,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRanges) { int32_t batch_size = 30; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -460,15 +450,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, RefreshReadRangesDisablePrefetchByAdapti /*need_prefetch=*/true, /*set_read_ranges_statuses=*/{}); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, - /*prefetch_batch_count=*/2, - /*enable_adaptive_prefetch_strategy=*/true, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, + /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/true, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_FALSE(reader->NeedPrefetch()); } @@ -478,14 +467,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { int32_t batch_size = 30; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -522,14 +510,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail /*set_read_ranges_statuses=*/ {Status::IOError("set read ranges failed"), Status::IOError("set read ranges failed")}); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -539,43 +526,23 @@ TEST_F(PrefetchFileBatchReaderImplTest, SetReadRangesReturnErrorWhenPushDownFail ASSERT_TRUE(status.IsIOError()); } -TEST_F(PrefetchFileBatchReaderImplTest, NeedInitCacheNeverMode) { - auto data_array = PrepareArray(10); - int32_t batch_size = 5; - int32_t prefetch_max_parallel_num = 1; - MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::NEVER, - CacheConfig(), GetDefaultPool())); - - auto prefetch_reader = dynamic_cast(reader.get()); - ASSERT_FALSE(prefetch_reader->NeedInitCache()); -} - TEST_F(PrefetchFileBatchReaderImplTest, WorkloopSetReadStatusWhenCacheInitFailed) { auto data_array = PrepareArray(10); int32_t batch_size = 5; int32_t prefetch_max_parallel_num = 1; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); CacheConfig invalid_cache_config( - /*buffer_size_limit=*/512 * 1024, /*range_size_limit=*/4 * 1024, /*hole_size_limit=*/8 * 1024, /*pre_buffer_limit=*/128 * 1024); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - invalid_cache_config, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + invalid_cache_config, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -589,14 +556,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenShutdown) { int32_t batch_size = 5; int32_t prefetch_max_parallel_num = 1; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -608,14 +574,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, DoReadBatchReturnOkWhenNoCurrentReadRang int32_t batch_size = 5; int32_t prefetch_max_parallel_num = 1; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/false, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -627,14 +592,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) { int32_t batch_size = 150; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -648,14 +612,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { dynamic_cast(prefetch_reader->readers_[i].get()) @@ -694,14 +657,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -730,14 +692,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -750,14 +711,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 6; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); auto row_ids = array_and_row_ids.second; @@ -776,14 +736,13 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { int32_t batch_size = 10; int32_t prefetch_max_parallel_num = 3; MockFormatReaderBuilder reader_builder(data_array, data_type_, batch_size); - ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -797,16 +756,16 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, /*prefetch_max_parallel_num=*/0, batch_size, 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + GetDefaultPool())); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, prefetch_max_parallel_num, /*batch_size=*/-1, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + GetDefaultPool())); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -814,33 +773,32 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, /*executor=*/nullptr, /*initialize_read_ranges=*/true, - /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, CacheConfig(), GetDefaultPool())); + /*read_ahead_cache_enabled=*/true, CacheConfig(), GetDefaultPool())); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( data_file_path, /*data_file_size=*/0, /*reader_builder=*/nullptr, mock_fs_, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + GetDefaultPool())); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( data_file_path, /*data_file_size=*/0, &reader_builder, /*fs=*/nullptr, prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, CacheConfig(), + GetDefaultPool())); } { ASSERT_OK_AND_ASSIGN( - auto reader, - PrefetchFileBatchReaderImpl::Create( - data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, - prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + auto reader, PrefetchFileBatchReaderImpl::Create( + data_file_path, /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, batch_size, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/true, + CacheConfig(), GetDefaultPool())); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -850,7 +808,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { /// [30,60) will be filtered out. /// The read range is [0,30), [30,60), [60,90). So, expected results is [0,30), [60,90) TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCompleteFiltering) { - auto [file_format, cache_mode] = GetParam(); + auto [file_format, read_ahead_cache_enabled] = GetParam(); auto data_array = PrepareArray(90); int32_t batch_size = 10; PrepareTestData(file_format, data_array, /*stripe_row_count=*/30, /*row_index_stride=*/30); @@ -866,7 +824,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, /*selection_bitmap=*/std::nullopt, /*batch_size=*/batch_size, /*prefetch_max_parallel_num=*/3, - cache_mode); + read_ahead_cache_enabled); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); arrow::ArrayVector expected_array_vector; @@ -883,7 +841,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom /// The read range is [0,30), [30,60), [60,90). TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithOrcPredicatePushdownWithRowGroupGranularity) { - auto [file_format, cache_mode] = GetParam(); + auto [file_format, read_ahead_cache_enabled] = GetParam(); auto data_array = PrepareArray(90); int32_t batch_size = 10; PrepareTestData(file_format, data_array, /*stripe_row_count=*/30, /*row_index_stride=*/10); @@ -900,7 +858,7 @@ TEST_P(PrefetchFileBatchReaderImplTest, auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, /*selection_bitmap=*/std::nullopt, /*batch_size=*/batch_size, /*prefetch_max_parallel_num=*/3, - cache_mode); + read_ahead_cache_enabled); ASSERT_OK(reader->RefreshReadRanges()); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); @@ -925,14 +883,14 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { MockFormatReaderBuilder reader_builder(data_array, data_type_, bitmap, /*read_batch_size=*/100); int32_t prefetch_max_parallel_num = 3; - ASSERT_OK_AND_ASSIGN(auto reader, PrefetchFileBatchReaderImpl::Create( - /*data_file_path=*/"", /*data_file_size=*/0, - &reader_builder, mock_fs_, prefetch_max_parallel_num, - /*batch_size=*/100, prefetch_max_parallel_num * 2, - /*enable_adaptive_prefetch_strategy=*/false, executor_, - /*initialize_read_ranges=*/true, - /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, - CacheConfig(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, mock_fs_, + prefetch_max_parallel_num, + /*batch_size=*/100, prefetch_max_parallel_num * 2, + /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, + /*read_ahead_cache_enabled=*/true, CacheConfig(), GetDefaultPool())); ASSERT_OK_AND_ASSIGN(auto result_chunk_array, ReadResultCollector::CollectResult(reader.get())); ASSERT_OK_AND_ASSIGN(auto data_batch, ReadResultCollector::GetReadBatch(data_array)); @@ -946,7 +904,7 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { } TEST_P(PrefetchFileBatchReaderImplTest, TestRowMapping) { - auto [file_format, cache_mode] = GetParam(); + auto [file_format, read_ahead_cache_enabled] = GetParam(); auto data_array = PrepareArray(90); PrepareTestData(file_format, data_array, /*stripe_row_count=*/30, /*row_index_stride=*/10); auto schema = arrow::schema(fields_); @@ -959,10 +917,10 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestRowMapping) { Literal(70l), Literal(79l)), })); - auto reader = - PreparePrefetchReader(file_format, schema.get(), predicate, - /*selection_bitmap=*/std::nullopt, - /*batch_size=*/10, /*prefetch_max_parallel_num=*/3, cache_mode); + auto reader = PreparePrefetchReader(file_format, schema.get(), predicate, + /*selection_bitmap=*/std::nullopt, + /*batch_size=*/10, /*prefetch_max_parallel_num=*/3, + read_ahead_cache_enabled); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, paimon::test::ReadResultCollector::CollectResultOneBatch(reader.get())); diff --git a/src/paimon/common/sst/sst_file_writer.cpp b/src/paimon/common/sst/sst_file_writer.cpp index ec736e33..f2b3b2de 100644 --- a/src/paimon/common/sst/sst_file_writer.cpp +++ b/src/paimon/common/sst/sst_file_writer.cpp @@ -20,6 +20,7 @@ #include "paimon/common/utils/crc32c.h" #include "paimon/common/utils/murmurhash_utils.h" +#include "paimon/common/utils/saturating_cast.h" namespace paimon { SstFileWriter::SstFileWriter(const std::shared_ptr& out, @@ -27,8 +28,10 @@ SstFileWriter::SstFileWriter(const std::shared_ptr& out, const std::shared_ptr& factory, const std::shared_ptr& pool) : pool_(pool), out_(out), bloom_filter_(bloom_filter), block_size_(block_size) { + // block_size * 1.1 exceeds INT32_MAX for block_size above ~1.9GB; saturate instead of + // relying on the undefined double->int32_t conversion. data_block_writer_ = - std::make_unique(static_cast(block_size * 1.1), pool); + std::make_unique(SaturatingDoubleToInteger(block_size * 1.1), pool); index_block_writer_ = std::make_unique(BlockHandle::MAX_ENCODED_LENGTH * 1024, pool); compression_type_ = factory->GetCompressionType(); diff --git a/src/paimon/common/table/special_fields.h b/src/paimon/common/table/special_fields.h index 74b95b19..3279bfed 100644 --- a/src/paimon/common/table/special_fields.h +++ b/src/paimon/common/table/special_fields.h @@ -66,6 +66,13 @@ struct SpecialFields { return data_field; } + static const DataField& RealtimeOffset() { + static const DataField data_field = + DataField(SpecialFieldIds::REALTIME_OFFSET, + arrow::field("_REALTIME_OFFSET", arrow::int64(), false)); + return data_field; + } + static bool IsSystemField(const std::string& field_name) { if (StringUtils::StartsWith(field_name, KEY_FIELD_PREFIX)) { return true; diff --git a/src/paimon/common/table/special_fields_test.cpp b/src/paimon/common/table/special_fields_test.cpp index 68e805fd..b61d289b 100644 --- a/src/paimon/common/table/special_fields_test.cpp +++ b/src/paimon/common/table/special_fields_test.cpp @@ -55,6 +55,13 @@ TEST(SpecialFieldsTest, TestIndexScore) { ASSERT_EQ(SpecialFields::IndexScore().Type()->id(), arrow::Type::FLOAT); } +TEST(SpecialFieldsTest, TestRealtimeOffset) { + ASSERT_EQ(SpecialFields::RealtimeOffset().Id(), SpecialFieldIds::REALTIME_OFFSET); + ASSERT_EQ(SpecialFields::RealtimeOffset().Name(), "_REALTIME_OFFSET"); + ASSERT_EQ(SpecialFields::RealtimeOffset().Type()->id(), arrow::Type::INT64); + ASSERT_FALSE(SpecialFields::RealtimeOffset().Nullable()); +} + TEST(SpecialFieldsTest, TestKeyValueSpecialFieldCount) { ASSERT_EQ(SpecialFields::KEY_VALUE_SPECIAL_FIELD_COUNT, 2); } @@ -66,6 +73,7 @@ TEST(SpecialFieldsTest, TestIsSystemField) { ASSERT_TRUE(SpecialFields::IsSystemField("rowkind")); ASSERT_TRUE(SpecialFields::IsSystemField("_ROW_ID")); ASSERT_TRUE(SpecialFields::IsSystemField("_INDEX_SCORE")); + ASSERT_FALSE(SpecialFields::IsSystemField("_REALTIME_OFFSET")); ASSERT_TRUE(SpecialFields::IsSystemField("_KEY_0")); } diff --git a/src/paimon/common/types/data_type.cpp b/src/paimon/common/types/data_type.cpp index 9b6d3c90..2bf5d73c 100644 --- a/src/paimon/common/types/data_type.cpp +++ b/src/paimon/common/types/data_type.cpp @@ -29,6 +29,7 @@ #include "paimon/common/types/array_type.h" #include "paimon/common/types/map_type.h" #include "paimon/common/types/row_type.h" +#include "paimon/common/types/vector_type.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/decimal_utils.h" @@ -52,6 +53,8 @@ std::unique_ptr DataType::Create( return std::make_unique(type, nullable, metadata); case arrow::Type::type::LIST: return std::make_unique(type, nullable, metadata); + case arrow::Type::type::FIXED_SIZE_LIST: + return std::make_unique(type, nullable, metadata); case arrow::Type::type::STRUCT: if (VariantTypeUtils::IsVariantMetadata(metadata)) { // A variant field is physically a struct but is a scalar diff --git a/src/paimon/common/types/data_type_json_parser.cpp b/src/paimon/common/types/data_type_json_parser.cpp index 33308ab6..e95582a1 100644 --- a/src/paimon/common/types/data_type_json_parser.cpp +++ b/src/paimon/common/types/data_type_json_parser.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,7 @@ #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/types/vector_type.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/rapidjson_util.h" #include "paimon/common/utils/string_utils.h" @@ -148,6 +150,7 @@ enum class Keyword : int32_t { ROW, BLOB, VARIANT, + VECTOR, // NULL is keyword in c++ NULL_, RAW, @@ -197,6 +200,7 @@ const std::map& Keywords() { {"ROW", Keyword::ROW}, {"BLOB", Keyword::BLOB}, {"VARIANT", Keyword::VARIANT}, + {"VECTOR", Keyword::VECTOR}, {"NULL", Keyword::NULL_}, {"RAW", Keyword::RAW}, {"LEGACY", Keyword::LEGACY}, @@ -249,6 +253,7 @@ class TokenParser { Result> ParseDoubleType(); Result> ParseTimestampType(); Result> ParseTimestampLtzType(); + Result> ParseVectorType(); Result ParseOptionalPrecision(int32_t default_precision); private: @@ -526,6 +531,8 @@ Result> TokenParser::ParseTypeByKeyword( return ParseTimestampType(); case Keyword::TIMESTAMP_LTZ: return ParseTimestampLtzType(); + case Keyword::VECTOR: + return ParseVectorType(); default: return Status::Invalid(fmt::format("Unsupported type: {}", GetToken().value)); } @@ -607,6 +614,31 @@ Result> TokenParser::ParseTimestampLtzType() { return ts_type; } +Result> TokenParser::ParseVectorType() { + PAIMON_RETURN_NOT_OK(NextToken(TokenType::BEGIN_SUBTYPE)); + bool element_nullable = true; + AtomicTypeAttributes element_attributes; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element_type, + ParseTypeWithNullability(&element_nullable, &element_attributes)); + if (element_attributes.is_blob || element_attributes.is_variant || + !VectorType::IsValidElementType(element_type)) { + return Status::Invalid( + fmt::format("Invalid element type for vector: {}", element_type->ToString())); + } + PAIMON_RETURN_NOT_OK(NextToken(TokenType::LIST_SEPARATOR)); + PAIMON_RETURN_NOT_OK(NextToken(TokenType::LITERAL_INT)); + const std::string& length_token = GetToken().value; + std::optional length = StringUtils::StringToValue(length_token); + if (!length || length.value() < 1) { + return Status::Invalid( + fmt::format("Vector length must be between 1 and {} (both inclusive), but was {}", + std::numeric_limits::max(), length_token)); + } + PAIMON_RETURN_NOT_OK(NextToken(TokenType::END_SUBTYPE)); + return arrow::fixed_size_list(arrow::field("item", element_type, element_nullable), + length.value()); +} + Result TokenParser::ParseOptionalPrecision(int32_t default_precision) { auto precision = default_precision; if (HasNextToken({TokenType::BEGIN_PARAMETER})) { @@ -659,6 +691,8 @@ Result> DataTypeJsonParser::ParseComplexTypeField( if (StringUtils::StartsWith(type_str, "ARRAY")) { return ParseArrayType(name, type_json_value, nullable); + } else if (StringUtils::StartsWith(type_str, "VECTOR")) { + return ParseVectorType(name, type_json_value, nullable); } else if (StringUtils::StartsWith(type_str, "MAP")) { return ParseMapType(name, type_json_value, nullable); } else if (StringUtils::StartsWith(type_str, "ROW")) { @@ -681,6 +715,27 @@ Result> DataTypeJsonParser::ParseArrayType( return arrow::field(name, arrow::list(element_field), nullable); } +Result> DataTypeJsonParser::ParseVectorType( + const std::string& name, const rapidjson::Value& type_json_value, bool nullable) { + if (!type_json_value.HasMember("element") || !type_json_value.HasMember("length")) { + return Status::Invalid("vector data type must have element and length"); + } + if (!type_json_value["length"].IsInt()) { + return Status::Invalid("vector length must be an integer"); + } + int32_t length = type_json_value["length"].GetInt(); + if (length < 1) { + return Status::Invalid("Vector length must be between 1 and 2147483647 (both inclusive)"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element_field, + ParseType("item", type_json_value["element"])); + if (!VectorType::IsValidElementType(element_field->type())) { + return Status::Invalid( + fmt::format("Invalid element type for vector: {}", element_field->type()->ToString())); + } + return arrow::field(name, arrow::fixed_size_list(element_field, length), nullable); +} + Result> DataTypeJsonParser::ParseMapType( const std::string& name, const rapidjson::Value& type_json_value, bool nullable) { if (!type_json_value.HasMember("key") || !type_json_value.HasMember("value")) { diff --git a/src/paimon/common/types/data_type_json_parser.h b/src/paimon/common/types/data_type_json_parser.h index 92cb5d55..2134236a 100644 --- a/src/paimon/common/types/data_type_json_parser.h +++ b/src/paimon/common/types/data_type_json_parser.h @@ -50,6 +50,8 @@ class DataTypeJsonParser { static Result> ParseArrayType( const std::string& name, const rapidjson::Value& type_json_value, bool nullable); + static Result> ParseVectorType( + const std::string& name, const rapidjson::Value& type_json_value, bool nullable); static Result> ParseMapType( const std::string& name, const rapidjson::Value& type_json_value, bool nullable); static Result> ParseRowType( diff --git a/src/paimon/common/types/data_type_json_parser_test.cpp b/src/paimon/common/types/data_type_json_parser_test.cpp index 5026db02..e5dfbc21 100644 --- a/src/paimon/common/types/data_type_json_parser_test.cpp +++ b/src/paimon/common/types/data_type_json_parser_test.cpp @@ -48,6 +48,56 @@ TEST(DataTypeJsonParserTest, ParseTypeArrayTypeSuccess) { ASSERT_NE(field, nullptr); } +TEST(DataTypeJsonParserTest, ParseVectorTypeSuccess) { + const char* json = R"({ + "type": "VECTOR NOT NULL", + "element": "FLOAT", + "length": 3 + })"; + rapidjson::Document doc; + doc.Parse(json); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr field, + DataTypeJsonParser::ParseType("embedding", doc)); + ASSERT_FALSE(field->nullable()); + ASSERT_EQ(field->type()->id(), arrow::Type::FIXED_SIZE_LIST); + auto vector_type = checked_pointer_cast(field->type()); + ASSERT_EQ(vector_type->list_size(), 3); + ASSERT_TRUE(vector_type->value_type()->Equals(arrow::float32())); + + rapidjson::Document sql_doc; + rapidjson::Value sql_value("VECTOR", sql_doc.GetAllocator()); + ASSERT_OK_AND_ASSIGN(field, DataTypeJsonParser::ParseType("embedding", sql_value)); + vector_type = checked_pointer_cast(field->type()); + ASSERT_TRUE(field->nullable()); + ASSERT_EQ(vector_type->list_size(), 5); + ASSERT_FALSE(vector_type->value_field()->nullable()); + ASSERT_TRUE(vector_type->value_type()->Equals(arrow::int64())); +} + +TEST(DataTypeJsonParserTest, ParseVectorTypeFailure) { + for (const char* json : { + R"({"type":"VECTOR","element":"FLOAT","length":0})", + R"({"type":"VECTOR","element":"STRING","length":3})", + R"({"type":"VECTOR","element":"FLOAT"})", + R"({"type":"VECTOR","element":"FLOAT","length":"3"})", + }) { + rapidjson::Document doc; + doc.Parse(json); + ASSERT_NOK(DataTypeJsonParser::ParseType("embedding", doc)); + } + + rapidjson::Document sql_doc; + rapidjson::Value sql_value("VECTOR", sql_doc.GetAllocator()); + ASSERT_NOK_WITH_MSG(DataTypeJsonParser::ParseType("embedding", sql_value), + "Invalid element type for vector"); + sql_value.SetString("VECTOR", sql_doc.GetAllocator()); + ASSERT_OK(DataTypeJsonParser::ParseType("embedding", sql_value)); + sql_value.SetString("VECTOR", sql_doc.GetAllocator()); + ASSERT_NOK_WITH_MSG(DataTypeJsonParser::ParseType("embedding", sql_value), + "Vector length must be between 1 and 2147483647"); +} + TEST(DataTypeJsonParserTest, ParseTypeMapTypeSuccess) { const std::string name = "map_field"; const char* json = R"({ diff --git a/src/paimon/common/types/data_type_test.cpp b/src/paimon/common/types/data_type_test.cpp index 9568c3ff..d6eacdc1 100644 --- a/src/paimon/common/types/data_type_test.cpp +++ b/src/paimon/common/types/data_type_test.cpp @@ -148,4 +148,18 @@ TEST(DataTypeTest, NestedTypeSerializationUsesChildMetadata) { R"({"type":"ARRAY","element":"INT"})"); } +TEST(DataTypeTest, VectorTypeSerialization) { + auto vector_field = arrow::field( + "embedding", arrow::fixed_size_list(arrow::field("item", arrow::float32()), 3), false); + auto data_type = + DataType::Create(vector_field->type(), vector_field->nullable(), vector_field->metadata()); + rapidjson::Document doc; + auto value = data_type->ToJson(&doc.GetAllocator()); + rapidjson::StringBuffer buffer; + rapidjson::Writer writer(buffer); + value.Accept(writer); + ASSERT_EQ(std::string(buffer.GetString()), + R"({"type":"VECTOR NOT NULL","element":"FLOAT","length":3})"); +} + } // namespace paimon::test diff --git a/src/paimon/common/types/vector_type.h b/src/paimon/common/types/vector_type.h new file mode 100644 index 00000000..9c55c165 --- /dev/null +++ b/src/paimon/common/types/vector_type.h @@ -0,0 +1,76 @@ +/* + * 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. + */ + +#pragma once + +#include +#include + +#include "arrow/api.h" +#include "paimon/common/types/data_type.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/rapidjson_util.h" + +namespace paimon { + +/// Fixed-size VECTOR logical type backed by Arrow FixedSizeList. +class VectorType : public DataType { + public: + static constexpr char TYPE[] = "VECTOR"; + + VectorType(const std::shared_ptr& type, bool nullable, + const std::shared_ptr& metadata) + : DataType(type, nullable, metadata) {} + + static bool IsValidElementType(const std::shared_ptr& type) { + switch (type->id()) { + case arrow::Type::BOOL: + case arrow::Type::INT8: + case arrow::Type::INT16: + case arrow::Type::INT32: + case arrow::Type::INT64: + case arrow::Type::FLOAT: + case arrow::Type::DOUBLE: + return true; + default: + return false; + } + } + + rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const + noexcept(false) override { + rapidjson::Value obj(rapidjson::kObjectType); + obj.AddMember( + rapidjson::StringRef("type"), + RapidJsonUtil::SerializeValue(WithNullable(std::string(TYPE)), allocator).Move(), + *allocator); + auto* type = checked_cast(type_.get()); + auto value_field = type->value_field(); + std::shared_ptr data_type = + DataType::Create(value_field->type(), value_field->nullable(), value_field->metadata()); + obj.AddMember(rapidjson::StringRef("element"), + RapidJsonUtil::SerializeValue(*data_type, allocator).Move(), *allocator); + obj.AddMember(rapidjson::StringRef("length"), + RapidJsonUtil::SerializeValue(type->list_size(), allocator).Move(), + *allocator); + return obj; + } +}; + +} // namespace paimon diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index 707e888f..f29e1d11 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -30,6 +30,7 @@ #include "arrow/util/compression.h" #include "fmt/format.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/arrow/vector_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/string_utils.h" @@ -160,6 +161,28 @@ Result> RebaseListLike( return rebased; } +/// Rebases a fixed size list array, whose child holds `list_size` values per row. +Result> RebaseFixedSizeList( + const std::shared_ptr& data, arrow::MemoryPool* pool) { + if (data->child_data.size() != 1) { + return CopyToZeroOffset(data, pool); + } + const int64_t list_size = + checked_cast(*data->type).list_size(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr validity, + RebaseValidityBitmap(*data, pool)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr child_slice, + data->child_data[0]->SliceSafe(data->offset * list_size, data->length * list_size)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr child, + RebaseToZeroOffset(child_slice, pool)); + std::shared_ptr rebased = + arrow::ArrayData::Make(data->type, data->length, data->null_count.load(), /*offset=*/0); + rebased->buffers = {std::move(validity)}; + rebased->child_data = {std::move(child)}; + return rebased; +} + /// Rebases a struct array, whose slices keep full length children. Result> RebaseStruct( const std::shared_ptr& data, arrow::MemoryPool* pool) { @@ -233,6 +256,8 @@ Result> RebaseToZeroOffset( return RebaseListLike(data, pool); case arrow::Type::LARGE_LIST: return RebaseListLike(data, pool); + case arrow::Type::FIXED_SIZE_LIST: + return RebaseFixedSizeList(data, pool); case arrow::Type::STRUCT: return RebaseStruct(data, pool); case arrow::Type::DICTIONARY: @@ -320,6 +345,11 @@ void ArrowUtils::TraverseArray(const std::shared_ptr& array) { TraverseArray(list_array->values()); return; } + case arrow::Type::type::FIXED_SIZE_LIST: { + auto* vector_array = checked_cast(array.get()); + TraverseArray(vector_array->values()); + return; + } default: return; } @@ -330,6 +360,13 @@ bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr& ty if (type->id() != other_type->id() || type->num_fields() != other_type->num_fields()) { return false; } + if (type->id() == arrow::Type::FIXED_SIZE_LIST) { + const auto& vector_type = checked_cast(*type); + const auto& other_vector_type = checked_cast(*other_type); + if (vector_type.list_size() != other_vector_type.list_size()) { + return false; + } + } for (int32_t i = 0; i < type->num_fields(); ++i) { const auto& field = type->field(i); const auto& other_field = other_type->field(i); @@ -363,6 +400,12 @@ Status ArrowUtils::InnerCheckNullabilityMatch(const std::shared_ptr(data); PAIMON_RETURN_NOT_OK( InnerCheckNullabilityMatch(list_type->value_field(), list_array->values())); + } else if (type->id() == arrow::Type::FIXED_SIZE_LIST) { + Status status = VectorUtils::ValidateVectorElements(*data); + if (!status.ok()) { + return Status::Invalid( + fmt::format("VECTOR field {} is invalid: {}", field->name(), status.message())); + } } else if (type->id() == arrow::Type::MAP) { auto map_type = checked_pointer_cast(field->type()); auto map_array = checked_pointer_cast(data); diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp b/src/paimon/common/utils/arrow/arrow_utils_test.cpp index 3680291c..4e1fdaa0 100644 --- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp @@ -249,6 +249,42 @@ TEST(ArrowUtilsTest, TestCheckNullableMatchWithList) { } } +TEST(ArrowUtilsTest, TestCheckNullableMatchRejectsNullVectorElement) { + auto vector_type = arrow::fixed_size_list(arrow::float32(), 3); + auto vector_field = arrow::field("embedding", vector_type); + arrow::FloatBuilder values_builder; + ASSERT_TRUE(values_builder.Append(1.0f).ok()); + ASSERT_TRUE(values_builder.AppendNull().ok()); + ASSERT_TRUE(values_builder.Append(3.0f).ok()); + std::shared_ptr values = values_builder.Finish().ValueOrDie(); + auto vector_data = arrow::ArrayData::Make(vector_type, 1, {nullptr}, {values->data()}, 0); + auto vector_array = arrow::MakeArray(vector_data); + auto struct_array = arrow::StructArray::Make({vector_array}, {vector_field}).ValueOrDie(); + + ASSERT_NOK_WITH_MSG( + ArrowUtils::CheckNullabilityMatch(arrow::schema({vector_field}), struct_array), + "VECTOR field embedding is invalid: VECTOR cannot contain null elements"); +} + +// Arrow accepts a FixedSizeList whose child is shorter than `length * list_size` when importing +// it over the C data interface, so the nullability check must reject it rather than scan past the +// end of the child. +TEST(ArrowUtilsTest, TestCheckNullableMatchRejectsTruncatedVector) { + auto vector_type = arrow::fixed_size_list(arrow::float32(), 3); + auto vector_field = arrow::field("embedding", vector_type); + arrow::FloatBuilder values_builder; + ASSERT_TRUE(values_builder.AppendValues({1.0f, 2.0f, 3.0f}).ok()); + std::shared_ptr values = values_builder.Finish().ValueOrDie(); + auto vector_data = arrow::ArrayData::Make(vector_type, /*length=*/2, {nullptr}, + {values->data()}, /*null_count=*/0); + auto vector_array = arrow::MakeArray(vector_data); + auto struct_array = arrow::StructArray::Make({vector_array}, {vector_field}).ValueOrDie(); + + ASSERT_NOK_WITH_MSG( + ArrowUtils::CheckNullabilityMatch(arrow::schema({vector_field}), struct_array), + "VECTOR field embedding is invalid: VECTOR holds 3 elements while 2 rows of dimension 3"); +} + TEST(ArrowUtilsTest, TestCheckNullableMatchWithMap) { auto key_field = arrow::field("key", arrow::int32(), /*nullable=*/false); auto value_field = arrow::field("value", arrow::int32(), /*nullable=*/true); @@ -523,6 +559,9 @@ std::vector NormalizeCases() { {arrow::list(arrow::utf8()), R"([["a"], null, ["bb", "ccc"], [], ["d"], null, ["e", "f"], [], ["g"], ["h"]])", {{{0}, 2}}}, + {arrow::fixed_size_list(arrow::int32(), 2), + "[[0, 1], null, [2, 3], [4, 5], [6, 7], null, [8, 9], [10, 11], [12, 13], [14, 15]]", + {{{0}, 1}}}, {arrow::struct_({int_field, text_field}), R"([{"a": 0, "b": "x"}, null, {"a": 2, "b": null}, {"a": null, "b": "yyy"}, {"a": 4, "b": "z"}, {"a": 5, "b": ""}, null, {"a": 7, "b": "w"}, @@ -758,6 +797,14 @@ TEST(ArrowUtilsTest, TestEqualsIgnoreNullable) { ASSERT_TRUE(ArrowUtils::EqualsIgnoreNullable(struct_type1, struct_type3)); ASSERT_FALSE(ArrowUtils::EqualsIgnoreNullable(struct_type1, struct_type4)); } + { + auto vector3 = arrow::fixed_size_list(arrow::float32(), 3); + auto vector3_non_null = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), false), 3); + auto vector5 = arrow::fixed_size_list(arrow::float32(), 5); + ASSERT_TRUE(ArrowUtils::EqualsIgnoreNullable(vector3, vector3_non_null)); + ASSERT_FALSE(ArrowUtils::EqualsIgnoreNullable(vector3, vector5)); + } { // test complex auto key_field = arrow::field("key", arrow::int32(), /*nullable=*/false); diff --git a/src/paimon/common/utils/arrow/vector_utils.cpp b/src/paimon/common/utils/arrow/vector_utils.cpp new file mode 100644 index 00000000..e5cc8396 --- /dev/null +++ b/src/paimon/common/utils/arrow/vector_utils.cpp @@ -0,0 +1,128 @@ +/* + * 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 "paimon/common/utils/arrow/vector_utils.h" + +#include + +#include "arrow/array.h" +#include "arrow/array/array_nested.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/utils/checked_cast.h" + +namespace paimon { +namespace { + +Status ValidateListVector(const arrow::ListArray& array) { + if (array.values()->null_count() == 0) { + return Status::OK(); + } + for (int64_t i = 0; i < array.length(); ++i) { + if (array.IsNull(i)) { + continue; + } + int64_t value_offset = array.value_offset(i); + int64_t value_length = array.value_length(i); + for (int64_t j = 0; j < value_length; ++j) { + if (array.values()->IsNull(value_offset + j)) { + return Status::Invalid(fmt::format( + "VECTOR cannot contain null elements, found one at row {} position {}", i, j)); + } + } + } + return Status::OK(); +} + +Status ValidateFixedSizeListVector(const arrow::FixedSizeListArray& array) { + const auto& vector_type = checked_cast(*array.type()); + int32_t vector_length = vector_type.list_size(); + const std::shared_ptr& values = array.values(); + // Arrow does not check this when importing an array over the C data interface, so the + // element scan below would otherwise read past the end of the values array. + if (values->length() < (array.offset() + array.length()) * vector_length) { + return Status::Invalid(fmt::format( + "VECTOR holds {} elements while {} rows of dimension {} require {}", values->length(), + array.length(), vector_length, (array.offset() + array.length()) * vector_length)); + } + if (values->null_count() == 0) { + return Status::OK(); + } + for (int64_t i = 0; i < array.length(); ++i) { + if (array.IsNull(i)) { + continue; + } + int64_t value_offset = (array.offset() + i) * vector_length; + for (int32_t j = 0; j < vector_length; ++j) { + if (values->IsNull(value_offset + j)) { + return Status::Invalid(fmt::format( + "VECTOR cannot contain null elements, found one at row {} position {}", i, j)); + } + } + } + return Status::OK(); +} + +} // namespace + +bool VectorUtils::ContainsVectorType(const std::shared_ptr& type) { + if (!type) { + return false; + } + if (type->id() == arrow::Type::FIXED_SIZE_LIST) { + return true; + } + for (const auto& field : type->fields()) { + if (ContainsVectorType(field->type())) { + return true; + } + } + return false; +} + +bool VectorUtils::ContainsVectorField(const std::shared_ptr& field) { + return field != nullptr && ContainsVectorType(field->type()); +} + +bool VectorUtils::ContainsVector(const std::shared_ptr& schema) { + if (!schema) { + return false; + } + for (const auto& field : schema->fields()) { + if (ContainsVectorField(field)) { + return true; + } + } + return false; +} + +Status VectorUtils::ValidateVectorElements(const arrow::Array& array) { + switch (array.type_id()) { + case arrow::Type::LIST: + return ValidateListVector(checked_cast(array)); + case arrow::Type::FIXED_SIZE_LIST: + return ValidateFixedSizeListVector( + checked_cast(array)); + default: + return Status::Invalid( + fmt::format("Cannot validate VECTOR values of type {}", array.type()->ToString())); + } +} + +} // namespace paimon diff --git a/src/paimon/common/utils/arrow/vector_utils.h b/src/paimon/common/utils/arrow/vector_utils.h new file mode 100644 index 00000000..0031a5de --- /dev/null +++ b/src/paimon/common/utils/arrow/vector_utils.h @@ -0,0 +1,54 @@ +/* + * 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. + */ + +#pragma once + +#include + +#include "paimon/status.h" +#include "paimon/visibility.h" + +namespace arrow { +class Array; +class DataType; +class Field; +class Schema; +} // namespace arrow + +namespace paimon { + +/// Helpers shared by the schema, read and write paths handling VECTOR values, which are +/// represented as Arrow FixedSizeList. +class PAIMON_EXPORT VectorUtils { + public: + VectorUtils() = delete; + ~VectorUtils() = delete; + + static bool ContainsVectorType(const std::shared_ptr& type); + + static bool ContainsVectorField(const std::shared_ptr& field); + + static bool ContainsVector(const std::shared_ptr& schema); + + /// Rejects VECTOR values whose elements are not fully materialized or contain nulls. + /// `array` must be the List or FixedSizeList array holding the VECTOR values. + static Status ValidateVectorElements(const arrow::Array& array); +}; + +} // namespace paimon diff --git a/src/paimon/common/utils/arrow/vector_utils_test.cpp b/src/paimon/common/utils/arrow/vector_utils_test.cpp new file mode 100644 index 00000000..1cce6853 --- /dev/null +++ b/src/paimon/common/utils/arrow/vector_utils_test.cpp @@ -0,0 +1,100 @@ +/* + * 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 "paimon/common/utils/arrow/vector_utils.h" + +#include + +#include "arrow/api.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr ArrayFromJSON(const std::shared_ptr& type, + const std::string& json) { + arrow::Result> result = + arrow::ipc::internal::json::ArrayFromJSON(type, json); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return std::move(result).ValueOrDie(); +} + +} // namespace + +TEST(VectorUtilsTest, TestContainsVector) { + auto vector_type = arrow::fixed_size_list(arrow::float32(), 3); + ASSERT_TRUE(VectorUtils::ContainsVectorType(vector_type)); + ASSERT_TRUE(VectorUtils::ContainsVectorType(arrow::list(vector_type))); + ASSERT_TRUE(VectorUtils::ContainsVectorType(arrow::map(arrow::utf8(), vector_type))); + ASSERT_TRUE(VectorUtils::ContainsVectorType(arrow::struct_({arrow::field("v", vector_type)}))); + ASSERT_FALSE(VectorUtils::ContainsVectorType(arrow::list(arrow::float32()))); + ASSERT_FALSE(VectorUtils::ContainsVectorType(nullptr)); + + ASSERT_TRUE(VectorUtils::ContainsVectorField(arrow::field("v", arrow::list(vector_type)))); + ASSERT_FALSE(VectorUtils::ContainsVectorField(arrow::field("v", arrow::int32()))); + ASSERT_FALSE(VectorUtils::ContainsVectorField(nullptr)); + + ASSERT_TRUE(VectorUtils::ContainsVector( + arrow::schema({arrow::field("id", arrow::int32()), arrow::field("v", vector_type)}))); + ASSERT_FALSE(VectorUtils::ContainsVector(arrow::schema({arrow::field("id", arrow::int32())}))); + ASSERT_FALSE(VectorUtils::ContainsVector(nullptr)); +} + +TEST(VectorUtilsTest, TestValidateVectorElements) { + auto vector_type = arrow::fixed_size_list(arrow::float32(), 3); + ASSERT_OK(VectorUtils::ValidateVectorElements( + *ArrayFromJSON(vector_type, R"([[1.0, 2.0, 3.0], null, [4.0, 5.0, 6.0]])"))); + ASSERT_NOK_WITH_MSG(VectorUtils::ValidateVectorElements( + *ArrayFromJSON(vector_type, R"([[1.0, 2.0, 3.0], [4.0, null, 6.0]])")), + "VECTOR cannot contain null elements, found one at row 1 position 1"); + + // A sliced array must be validated against its own rows only. + std::shared_ptr sliced = + ArrayFromJSON(vector_type, R"([[1.0, null, 3.0], [4.0, 5.0, 6.0]])")->Slice(1, 1); + ASSERT_OK(VectorUtils::ValidateVectorElements(*sliced)); + + auto list_type = arrow::list(arrow::float32()); + ASSERT_OK(VectorUtils::ValidateVectorElements( + *ArrayFromJSON(list_type, R"([[1.0, 2.0, 3.0], null])"))); + ASSERT_NOK_WITH_MSG( + VectorUtils::ValidateVectorElements(*ArrayFromJSON(list_type, R"([[1.0, null, 3.0]])")), + "VECTOR cannot contain null elements, found one at row 0 position 1"); + + ASSERT_NOK_WITH_MSG( + VectorUtils::ValidateVectorElements(*ArrayFromJSON(arrow::int32(), "[1, 2]")), + "Cannot validate VECTOR values of type int32"); +} + +// Arrow does not check that a FixedSizeList child holds `length * list_size` values when +// importing an array over the C data interface, so the element scan must reject it instead of +// reading past the end of the child. +TEST(VectorUtilsTest, TestValidateVectorElementsRejectsTruncatedValues) { + auto vector_type = arrow::fixed_size_list(arrow::float32(), 3); + std::shared_ptr values = ArrayFromJSON(arrow::float32(), "[1.0, null, 3.0]"); + auto truncated = arrow::MakeArray(arrow::ArrayData::Make(vector_type, /*length=*/2, {nullptr}, + {values->data()}, + /*null_count=*/0)); + + ASSERT_NOK_WITH_MSG(VectorUtils::ValidateVectorElements(*truncated), + "VECTOR holds 3 elements while 2 rows of dimension 3 require 6"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/utils/binary_row_partition_computer.cpp b/src/paimon/common/utils/binary_row_partition_computer.cpp index 43ec7d40..ec773e32 100644 --- a/src/paimon/common/utils/binary_row_partition_computer.cpp +++ b/src/paimon/common/utils/binary_row_partition_computer.cpp @@ -77,31 +77,68 @@ Result> BinaryRowPartitionComputer:: Result BinaryRowPartitionComputer::ToBinaryRow( const std::map& partition) const { + return ConvertToBinaryRow(partition, /*included_fields=*/nullptr); +} + +Result BinaryRowPartitionComputer::ConvertToBinaryRow( + const std::map& partition, std::vector* included_fields) const { BinaryRow binary_row(partition_converters_.size()); BinaryRowWriter writer(&binary_row, /*initial_size=*/0, memory_pool_.get()); - for (size_t field_idx = 0; field_idx < partition_converters_.size(); field_idx++) { - const auto& partition_extractor = partition_converters_[field_idx]; - const auto& partition_key = partition_extractor.partition_key; - const auto& to_binary_row = partition_extractor.converter; - auto input_iter = partition.find(partition_key); + if (included_fields != nullptr) { + included_fields->assign(partition_converters_.size(), false); + } + for (size_t field_idx = 0; field_idx < partition_converters_.size(); ++field_idx) { + const PartitionConverter& partition_converter = partition_converters_[field_idx]; + auto input_iter = partition.find(partition_converter.partition_key); if (input_iter == partition.end()) { + if (included_fields != nullptr) { + writer.SetNullAt(field_idx); + continue; + } return Status::Invalid( fmt::format("can not find partition key '{}' in input partition '{}'", - partition_key, partition)); + partition_converter.partition_key, partition)); } - const auto& value_str = input_iter->second; - if (value_str == default_part_value_) { + if (included_fields != nullptr) { + (*included_fields)[field_idx] = true; + } + if (input_iter->second == default_part_value_) { // TODO(yonghao.fyh): when support decimal/ timestamp in partition, use // WriteTimestamp(null) for non compact precision writer.SetNullAt(field_idx); } else { - PAIMON_RETURN_NOT_OK(to_binary_row(value_str, field_idx, &writer)); + PAIMON_RETURN_NOT_OK( + partition_converter.converter(input_iter->second, field_idx, &writer)); } } writer.Complete(); return binary_row; } +Result> BinaryRowPartitionComputer::NormalizePartitionSpec( + const std::map& partition) const { + for (const auto& [partition_key, _] : partition) { + if (std::find(partition_keys_.begin(), partition_keys_.end(), partition_key) == + partition_keys_.end()) { + return Status::Invalid( + fmt::format("field {} does not exist in partition keys", partition_key)); + } + } + + std::vector included_fields; + PAIMON_ASSIGN_OR_RAISE(BinaryRow binary_row, ConvertToBinaryRow(partition, &included_fields)); + + std::vector> normalized_values; + PAIMON_ASSIGN_OR_RAISE(normalized_values, GeneratePartitionVector(binary_row)); + std::map normalized_partition; + for (size_t field_idx = 0; field_idx < normalized_values.size(); ++field_idx) { + if (included_fields[field_idx]) { + normalized_partition.insert(normalized_values[field_idx]); + } + } + return normalized_partition; +} + Result>> BinaryRowPartitionComputer::GeneratePartitionVector(const BinaryRow& partition) const { if (static_cast(partition.GetFieldCount()) != partition_converters_.size()) { diff --git a/src/paimon/common/utils/binary_row_partition_computer.h b/src/paimon/common/utils/binary_row_partition_computer.h index 63aa7549..2e96ad6f 100644 --- a/src/paimon/common/utils/binary_row_partition_computer.h +++ b/src/paimon/common/utils/binary_row_partition_computer.h @@ -55,6 +55,8 @@ class BinaryRowPartitionComputer { bool legacy_partition_name_enabled, const std::shared_ptr& memory_pool); Result ToBinaryRow(const std::map& partition) const; + Result> NormalizePartitionSpec( + const std::map& partition) const; Result>> GeneratePartitionVector( const BinaryRow& partition) const; const std::vector& GetPartitionKeys() const { @@ -73,6 +75,10 @@ class BinaryRowPartitionComputer { const std::vector& partition_converters, const std::shared_ptr& memory_pool); + /// A non-null `included_fields` enables partial partitions and records present fields. + Result ConvertToBinaryRow(const std::map& partition, + std::vector* included_fields) const; + static Result GetTypeFromArrowSchema( const std::shared_ptr& schema, const std::string& field_name); diff --git a/src/paimon/common/utils/binary_row_partition_computer_test.cpp b/src/paimon/common/utils/binary_row_partition_computer_test.cpp index 974c4d7a..18c42178 100644 --- a/src/paimon/common/utils/binary_row_partition_computer_test.cpp +++ b/src/paimon/common/utils/binary_row_partition_computer_test.cpp @@ -268,6 +268,36 @@ TEST(BinaryRowPartitionComputerTest, TestNullOrWhitespaceOnlyStr) { ASSERT_EQ(partition_key_values, expected); } +TEST(BinaryRowPartitionComputerTest, TestNormalizePartialPartitionSpec) { + using PartitionSpec = std::map; + + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr schema = + arrow::schema({arrow::field("pt", arrow::date32()), arrow::field("region", arrow::utf8())}); + std::vector partition_keys = {"pt", "region"}; + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr legacy_computer, + BinaryRowPartitionComputer::Create(partition_keys, schema, "__DEFAULT_PARTITION__", + /*legacy_partition_name_enabled=*/true, pool)); + ASSERT_OK_AND_ASSIGN(PartitionSpec legacy_partition, + legacy_computer->NormalizePartitionSpec({{"pt", "2024-01-01"}})); + PartitionSpec expected_legacy_partition = {{"pt", "19723"}}; + ASSERT_EQ(expected_legacy_partition, legacy_partition); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr non_legacy_computer, + BinaryRowPartitionComputer::Create(partition_keys, schema, "__DEFAULT_PARTITION__", + /*legacy_partition_name_enabled=*/false, pool)); + ASSERT_OK_AND_ASSIGN(PartitionSpec non_legacy_partition, + non_legacy_computer->NormalizePartitionSpec({{"pt", "2024-01-01"}})); + PartitionSpec expected_non_legacy_partition = {{"pt", "2024-01-01"}}; + ASSERT_EQ(expected_non_legacy_partition, non_legacy_partition); + + ASSERT_NOK_WITH_MSG(non_legacy_computer->NormalizePartitionSpec({{"unknown", "value"}}), + "field unknown does not exist in partition keys"); +} + TEST(BinaryRowPartitionComputerTest, TestPartToSimpleString) { auto pool = GetDefaultPool(); { diff --git a/src/paimon/common/utils/byte_range_combiner.cpp b/src/paimon/common/utils/byte_range_combiner.cpp index 42877003..c306c100 100644 --- a/src/paimon/common/utils/byte_range_combiner.cpp +++ b/src/paimon/common/utils/byte_range_combiner.cpp @@ -23,6 +23,7 @@ #include #include +#include #include "fmt/format.h" @@ -39,6 +40,20 @@ Result> ByteRangeCombiner::CoalesceByteRanges( return ranges; } + // Reject ranges that exceed the int64 bound before any offset + length arithmetic + // below. Such ranges can originate from corrupt file metadata (e.g. negative signed + // values cast to uint64_t) and would otherwise wrap around or make the splitting + // loop run until memory is exhausted. + constexpr auto kMaxRangeValue = static_cast(std::numeric_limits::max()); + for (const auto& range : ranges) { + if (range.offset > kMaxRangeValue || range.length > kMaxRangeValue || + range.offset + range.length > kMaxRangeValue) { + return Status::Invalid( + fmt::format("byte range (offset={}, length={}) exceeds the int64 bound", + range.offset, range.length)); + } + } + std::vector adjusted_ranges; for (const auto& range : ranges) { uint64_t range_start = range.offset; diff --git a/src/paimon/common/utils/byte_range_combiner.h b/src/paimon/common/utils/byte_range_combiner.h index 599ca59c..90942719 100644 --- a/src/paimon/common/utils/byte_range_combiner.h +++ b/src/paimon/common/utils/byte_range_combiner.h @@ -23,8 +23,8 @@ #include +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/result.h" -#include "paimon/utils/read_ahead_cache.h" namespace paimon { diff --git a/src/paimon/common/utils/byte_range_combiner_test.cpp b/src/paimon/common/utils/byte_range_combiner_test.cpp index 19d739a4..5a3cc1c6 100644 --- a/src/paimon/common/utils/byte_range_combiner_test.cpp +++ b/src/paimon/common/utils/byte_range_combiner_test.cpp @@ -21,9 +21,11 @@ #include "paimon/common/utils/byte_range_combiner.h" +#include + #include "gtest/gtest.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/testing/utils/testharness.h" -#include "paimon/utils/read_ahead_cache.h" namespace paimon::test { @@ -75,4 +77,26 @@ TEST(ByteRangeCombinerTest, TestBasics) { check({{20, 5}, {20, 5}, {21, 2}}, {{20, 5}}); } +// Ranges beyond the int64 bound (e.g. negative signed metadata values cast to uint64_t) +// must be rejected before the unchecked offset + length arithmetic, which would otherwise +// wrap around or spin the splitting loop until memory is exhausted. +TEST(ByteRangeCombinerTest, TestRejectsRangesBeyondInt64Bound) { + constexpr auto kInt64Max = static_cast(std::numeric_limits::max()); + auto check_invalid = [](std::vector ranges) -> void { + ASSERT_NOK_WITH_MSG( + ByteRangeCombiner::CoalesceByteRanges(std::move(ranges), /*hole_size_limit=*/9, + /*range_size_limit=*/99), + "exceeds the int64 bound"); + }; + + // Offset beyond int64 (negative int64 cast to uint64_t lands here). + check_invalid({{kInt64Max + 1, 1}}); + // Length beyond int64 (e.g. -1 cast to uint64_t), which would explode the split loop. + check_invalid({{0, std::numeric_limits::max()}}); + // Both in range individually, but the end position overflows the int64 bound. + check_invalid({{kInt64Max - 10, 20}}); + // One bad range among valid ones still fails the whole batch. + check_invalid({{100, 10}, {0, std::numeric_limits::max()}}); +} + } // namespace paimon::test diff --git a/src/paimon/common/utils/field_type_utils.h b/src/paimon/common/utils/field_type_utils.h index 72246769..d2786e26 100644 --- a/src/paimon/common/utils/field_type_utils.h +++ b/src/paimon/common/utils/field_type_utils.h @@ -93,6 +93,8 @@ class FieldTypeUtils { return FieldType::MAP; case arrow::Type::type::STRUCT: return FieldType::STRUCT; + case arrow::Type::type::FIXED_SIZE_LIST: + return FieldType::VECTOR; default: return Status::Invalid( fmt::format("Not support arrow type {}", static_cast(arrow_type))); @@ -135,6 +137,8 @@ class FieldTypeUtils { return "STRUCT"; case FieldType::VARIANT: return "VARIANT"; + case FieldType::VECTOR: + return "VECTOR"; default: return "UNKNOWN, type id:" + std::to_string(static_cast(type)); } diff --git a/src/paimon/common/utils/field_type_utils_test.cpp b/src/paimon/common/utils/field_type_utils_test.cpp index 50f60237..c70edb09 100644 --- a/src/paimon/common/utils/field_type_utils_test.cpp +++ b/src/paimon/common/utils/field_type_utils_test.cpp @@ -101,6 +101,10 @@ TEST(FieldTypeUtilsTest, ConvertToFieldType) { ASSERT_OK_AND_ASSIGN(result, FieldTypeUtils::ConvertToFieldType(arrow::Type::type::STRUCT)); ASSERT_EQ(result, FieldType::STRUCT); + ASSERT_OK_AND_ASSIGN(result, + FieldTypeUtils::ConvertToFieldType(arrow::Type::type::FIXED_SIZE_LIST)); + ASSERT_EQ(result, FieldType::VECTOR); + // Test unsupported Arrow type ASSERT_NOK(FieldTypeUtils::ConvertToFieldType(arrow::Type::type::UINT8)); } @@ -124,6 +128,7 @@ TEST(FieldTypeUtilsTest, FieldTypeToString) { ASSERT_EQ(FieldTypeUtils::FieldTypeToString(FieldType::ARRAY), "ARRAY"); ASSERT_EQ(FieldTypeUtils::FieldTypeToString(FieldType::MAP), "MAP"); ASSERT_EQ(FieldTypeUtils::FieldTypeToString(FieldType::STRUCT), "STRUCT"); + ASSERT_EQ(FieldTypeUtils::FieldTypeToString(FieldType::VECTOR), "VECTOR"); // Test UNKNOWN type auto unknown_type = static_cast(128); diff --git a/src/paimon/common/utils/read_ahead_cache.cpp b/src/paimon/common/utils/read_ahead_cache.cpp index b0001189..a74b3520 100644 --- a/src/paimon/common/utils/read_ahead_cache.cpp +++ b/src/paimon/common/utils/read_ahead_cache.cpp @@ -20,15 +20,19 @@ // Adapted from Apache ORC // https://github.com/apache/orc/blob/main/c%2B%2B/src/io/Cache.cc -#include "paimon/utils/read_ahead_cache.h" +#include "paimon/common/utils/read_ahead_cache.h" #include +#include #include +#include #include #include #include "paimon/common/utils/byte_range_combiner.h" #include "paimon/common/utils/math.h" +#include "paimon/memory/bytes.h" +#include "paimon/metrics.h" namespace paimon { @@ -47,18 +51,53 @@ struct RangeCacheEntry { } }; -CacheConfig::CacheConfig(uint64_t buffer_size_limit, uint64_t range_size_limit, - uint64_t hole_size_limit, uint64_t pre_buffer_limit) - : buffer_size_limit_(buffer_size_limit), - range_size_limit_(range_size_limit), +// Everything needed to dispatch the prefetch IO of an entry AFTER the entry +// has been published into entries_: the promise resolves the entry's future +// and the buffer capture keeps the destination alive for the async IO. +struct PendingFetch { + ByteRange range; + std::shared_ptr buffer; + std::shared_ptr> promise; +}; + +namespace { + +// Copy the requested window out of the covering entries into dest. The +// entries must fully cover the range and their futures must be resolved. +void CopyRangeFromEntries(const std::vector& covering, const ByteRange& range, + char* dest) { + size_t pos = 0; + for (const auto& entry : covering) { + const uint64_t entry_end = entry.range.offset + entry.range.length; + const uint64_t copy_begin = std::max(range.offset, entry.range.offset); + const uint64_t copy_end = std::min(range.offset + range.length, entry_end); + const auto copy_len = static_cast(copy_end - copy_begin); + std::memcpy(dest + pos, entry.buffer->data() + (copy_begin - entry.range.offset), copy_len); + pos += copy_len; + } +} + +} // namespace + +CacheConfig::CacheConfig(uint64_t range_size_limit, uint64_t hole_size_limit, + uint64_t pre_buffer_limit) + : range_size_limit_(range_size_limit), hole_size_limit_(hole_size_limit), pre_buffer_limit_(pre_buffer_limit) {} CacheConfig::CacheConfig() - : CacheConfig(/*buffer_size_limit=*/512 * 1024 * 1024, - /*range_size_limit=*/16 * 1024 * 1024, + // Aligned with the reader's request granularity and with realistic data + // file sizes: + // - range_size_limit matches the parquet reader's 32 MiB request blocks + // (Arrow ReadRangeCache's own range limit); a smaller limit cuts entries + // below the request size, so a request can never be served from one piece. + // - pre_buffer_limit must exceed the LARGEST single read a reader issues + // (coalesced column-chunk reads of ~128 MiB were observed): fetches are + // only dispatched up to this window, so a request reaching past it can + // never be served and falls back to a second fetch of the same bytes. + : CacheConfig(/*range_size_limit=*/32 * 1024 * 1024, /*hole_size_limit=*/8 * 1024, - /*pre_buffer_limit=*/128 * 1024 * 1024) {} + /*pre_buffer_limit=*/256 * 1024 * 1024) {} class ReadAheadCache::Impl { public: @@ -67,18 +106,37 @@ class ReadAheadCache::Impl { ~Impl(); Status Init(std::vector&& ranges); - Result Read(const ByteRange& range); + Result Read(const ByteRange& range, char* dest); void Reset(); + void ReleaseBuffers(); + void Warmup(); + void CollectMetrics(std::shared_ptr* metrics) const; private: - std::vector MakeCacheEntries(const std::vector& ranges) const; + /// Dispatch the prefetch IOs for entries that have already been published + /// into entries_. + void DispatchFetches(const std::vector& fetches); + /// Find the entries fully covering the given range under the read lock. + /// Returns an empty vector on miss. Entries are copied (shared buffers) + /// so the caller may use them after releasing the lock. + std::vector FindCoveringEntries(const ByteRange& range); void PreBuffer(uint64_t offset); + void CountHit(uint64_t size) { + hits_.fetch_add(1, std::memory_order_relaxed); + hit_bytes_.fetch_add(size, std::memory_order_relaxed); + } + void CountMiss(uint64_t size) { + misses_.fetch_add(1, std::memory_order_relaxed); + miss_bytes_.fetch_add(size, std::memory_order_relaxed); + } - /// Cache the given ranges in the background. + /// Mark, publish and fetch the pending ranges at the given indices. /// - /// The caller must ensure that the ranges do not overlap with each other, - /// nor with previously cached ranges. Otherwise, behaviour will be undefined. - void Cache(std::vector ranges); + /// Marking is_cached_ and publishing the promise-backed entries happen + /// atomically under the write lock, before any IO is dispatched, so a + /// reader racing the prefetch waits on the in-flight entries instead of + /// re-fetching the same bytes. + void Cache(std::vector pending_indices); std::shared_ptr stream_; CacheConfig config_; @@ -89,38 +147,52 @@ class ReadAheadCache::Impl { std::vector> is_cached_; std::vector pending_ranges_; bool is_initialized_ = false; + // Statistics of the Read() requests issued to the cache, aggregated over + // all streams sharing this cache. + std::atomic read_count_{0}; + std::atomic read_bytes_{0}; + std::atomic hits_{0}; + std::atomic hit_bytes_{0}; + std::atomic misses_{0}; + std::atomic miss_bytes_{0}; + // Prefetch IO statistics: how many requests and bytes were actually issued + // to the underlying stream. + std::atomic io_count_{0}; + std::atomic io_bytes_{0}; }; -void ReadAheadCache::Impl::Cache(std::vector ranges) { - std::sort(ranges.begin(), ranges.end(), - [](const ByteRange& a, const ByteRange& b) { return a.offset < b.offset; }); - std::vector new_entries = MakeCacheEntries(ranges); - // Add new entries, themselves ordered by offset - std::unique_lock lock(rw_mutex_); - if (entries_.size() > 0) { - size_t new_entries_size = 0; - for (const auto& e : new_entries) { - new_entries_size += e.range.length; - } - - size_t total_size = 0; - for (const auto& e : entries_) { - total_size += e.range.length; +void ReadAheadCache::Impl::Cache(std::vector pending_indices) { + std::vector new_entries; + std::vector fetches; + // Mark is_cached_, publish the promise-backed entries and only then + // dispatch the IOs. The mark and the publication happen atomically under + // the write lock: a reader racing the prefetch observes is_cached_=true + // only once the covering entries are already visible, so it waits on + // their futures instead of issuing a duplicate underlying read. + { + std::unique_lock lock(rw_mutex_); + for (size_t idx : pending_indices) { + if (is_cached_[idx].exchange(true)) { + continue; + } + const ByteRange& range = pending_ranges_[idx]; + auto promise = std::make_shared>(); + auto future = promise->get_future(); + auto buffer = std::make_shared(range.length, memory_pool_.get()); + fetches.push_back({range, buffer, promise}); + new_entries.emplace_back(range, std::move(buffer), std::move(future)); } - size_t limit = config_.GetBufferSizeLimit(); - while (!entries_.empty() && total_size + new_entries_size > limit) { - auto iter = entries_.begin(); - total_size -= entries_.front().range.length; - entries_.erase(iter); + if (!new_entries.empty()) { + // Entries are never evicted: the cache holds every published + // range until ReleaseBuffers()/Reset(), so an in-flight fetch + // always keeps its entry and thus its future reachable. + std::vector merged(entries_.size() + new_entries.size()); + std::merge(entries_.begin(), entries_.end(), new_entries.begin(), new_entries.end(), + merged.begin()); + entries_ = std::move(merged); } - - std::vector merged(entries_.size() + new_entries.size()); - std::merge(entries_.begin(), entries_.end(), new_entries.begin(), new_entries.end(), - merged.begin()); - entries_ = std::move(merged); - } else { - entries_ = std::move(new_entries); } + DispatchFetches(fetches); } Status ReadAheadCache::Impl::Init(std::vector&& ranges) { @@ -154,22 +226,18 @@ void ReadAheadCache::Impl::PreBuffer(uint64_t offset) { } size_t start_idx = std::distance(pending_ranges_.begin(), it); - std::vector ranges; + std::vector pending_indices; size_t total_bytes = 0; for (size_t i = start_idx; i < pending_ranges_.size(); ++i) { - size_t range_size = pending_ranges_[i].length; - total_bytes += range_size; + total_bytes += pending_ranges_[i].length; if (total_bytes > config_.GetPreBufferLimit()) { break; } - if (is_cached_[i].exchange(true)) { - continue; - } - ranges.emplace_back(pending_ranges_[i]); + pending_indices.push_back(i); } - if (!ranges.empty()) { - Cache(std::move(ranges)); + if (!pending_indices.empty()) { + Cache(std::move(pending_indices)); } } @@ -185,7 +253,22 @@ ReadAheadCache::Impl::~Impl() { } void ReadAheadCache::Impl::Reset() { + ReleaseBuffers(); + read_count_.store(0, std::memory_order_relaxed); + read_bytes_.store(0, std::memory_order_relaxed); + hits_.store(0, std::memory_order_relaxed); + hit_bytes_.store(0, std::memory_order_relaxed); + misses_.store(0, std::memory_order_relaxed); + miss_bytes_.store(0, std::memory_order_relaxed); + io_count_.store(0, std::memory_order_relaxed); + io_bytes_.store(0, std::memory_order_relaxed); +} + +void ReadAheadCache::Impl::ReleaseBuffers() { std::unique_lock lock(rw_mutex_); + // Entries are never evicted, so waiting on entries_ covers every + // dispatched fetch: no async callback can outlive the stream or the + // memory pool its buffer belongs to. for (auto& entry : entries_) { entry.future.wait(); } @@ -193,45 +276,106 @@ void ReadAheadCache::Impl::Reset() { is_cached_.clear(); pending_ranges_.clear(); is_initialized_ = false; + // The read/io counters are deliberately kept: a reader closed at EOF must + // still be able to report them through CollectMetrics(). } -Result ReadAheadCache::Impl::Read(const ByteRange& range) { +void ReadAheadCache::Impl::CollectMetrics(std::shared_ptr* metrics) const { + if (metrics == nullptr || !*metrics) { + return; + } + auto& m = *metrics; + m->SetCounter(ReadAheadCacheMetrics::READ_COUNT, read_count_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::READ_BYTES, read_bytes_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::READ_HITS, hits_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES, + hit_bytes_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::READ_MISSES, misses_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES, + miss_bytes_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::IO_COUNT, io_count_.load(std::memory_order_relaxed)); + m->SetCounter(ReadAheadCacheMetrics::IO_BYTES, io_bytes_.load(std::memory_order_relaxed)); +} + +void ReadAheadCache::Impl::Warmup() { + // Init() only registers the pending ranges; without this the first fetch + // starts when the first Read() arrives, racing the reader's own miss fetch. + if (!pending_ranges_.empty()) { + PreBuffer(pending_ranges_.front().offset); + } +} + +std::vector ReadAheadCache::Impl::FindCoveringEntries(const ByteRange& range) { + std::vector covering; + std::shared_lock lock(rw_mutex_); + // Find the entry holding the start of the range: the first entry whose + // end is beyond range.offset (entries are disjoint and sorted by offset). + auto it = std::lower_bound(entries_.begin(), entries_.end(), range.offset, + [](const RangeCacheEntry& e, uint64_t offset) { + return e.range.offset + e.range.length <= offset; + }); + if (it == entries_.end() || it->range.offset > range.offset) { + return covering; + } + if (it->range.Contains(range)) { + covering.push_back(*it); + return covering; + } + // The request spans several adjacent entries (a column chunk larger than + // one coalesced range): collect the contiguous run and check it covers + // the whole request. Entries are published before their fetch is + // dispatched, so a reader racing the prefetch waits for the in-flight + // fetch instead of issuing a second one for the same bytes. + uint64_t covered_end = it->range.offset + it->range.length; + covering.push_back(*it); + auto next = std::next(it); + while (covered_end < range.offset + range.length && next != entries_.end() && + next->range.offset == covered_end) { + covered_end = next->range.offset + next->range.length; + covering.push_back(*next); + ++next; + } + if (covered_end < range.offset + range.length) { + covering.clear(); + } + return covering; +} + +Result ReadAheadCache::Impl::Read(const ByteRange& range, char* dest) { if (range.length == 0) { - return ByteSlice{std::make_shared(0, memory_pool_.get()), 0, 0}; + return true; } + read_count_.fetch_add(1, std::memory_order_relaxed); + read_bytes_.fetch_add(range.length, std::memory_order_relaxed); PreBuffer(range.offset); - ByteSlice result{}; - { - std::shared_lock lock(rw_mutex_); - auto it = std::lower_bound(entries_.begin(), entries_.end(), range.offset, - [](const RangeCacheEntry& e, uint64_t offset) { - return e.range.offset + e.range.length <= offset; - }); - if (it != entries_.end() && it->range.Contains(range)) { - PAIMON_RETURN_NOT_OK(it->future.get()); - result = ByteSlice{it->buffer, range.offset - it->range.offset, range.length}; - return result; - } + std::vector covering = FindCoveringEntries(range); + if (covering.empty()) { + CountMiss(range.length); + return false; + } + // Wait OUTSIDE the lock: the futures resolve when the prefetch stream's + // async reads complete, and holding rw_mutex_ would block Cache(). + for (const auto& entry : covering) { + PAIMON_RETURN_NOT_OK(entry.future.get()); } - return result; + // The data copy runs OUTSIDE the lock for the same reason. + CopyRangeFromEntries(covering, range, dest); + CountHit(range.length); + return true; } -std::vector ReadAheadCache::Impl::MakeCacheEntries( - const std::vector& ranges) const { - std::vector new_entries; - new_entries.reserve(ranges.size()); - for (const auto& range : ranges) { - auto promise = std::make_shared>(); - auto future = promise->get_future(); - auto buffer = std::make_shared(range.length, memory_pool_.get()); +void ReadAheadCache::Impl::DispatchFetches(const std::vector& fetches) { + for (const auto& fetch : fetches) { + auto promise = fetch.promise; + auto buffer = fetch.buffer; auto read_size = static_cast(buffer->size()); - auto read_offset = static_cast(range.offset); + auto read_offset = static_cast(fetch.range.offset); stream_->ReadAsync( buffer->data(), read_size, read_offset, [promise, buffer](Status status) mutable { promise->set_value(status); }); - new_entries.emplace_back(range, std::move(buffer), std::move(future)); + io_count_.fetch_add(1, std::memory_order_relaxed); + io_bytes_.fetch_add(fetch.range.length, std::memory_order_relaxed); } - return new_entries; } ReadAheadCache::ReadAheadCache(const std::shared_ptr& stream, @@ -245,12 +389,24 @@ Status ReadAheadCache::Init(std::vector&& ranges) { return impl_->Init(std::move(ranges)); } -Result ReadAheadCache::Read(const ByteRange& range) { - return impl_->Read(range); +Result ReadAheadCache::Read(const ByteRange& range, char* dest) { + return impl_->Read(range, dest); } void ReadAheadCache::Reset() { return impl_->Reset(); } +void ReadAheadCache::ReleaseBuffers() { + return impl_->ReleaseBuffers(); +} + +void ReadAheadCache::Warmup() { + impl_->Warmup(); +} + +void ReadAheadCache::CollectMetrics(std::shared_ptr* metrics) const { + impl_->CollectMetrics(metrics); +} + } // namespace paimon diff --git a/include/paimon/utils/read_ahead_cache.h b/src/paimon/common/utils/read_ahead_cache.h similarity index 50% rename from include/paimon/utils/read_ahead_cache.h rename to src/paimon/common/utils/read_ahead_cache.h index 196045b7..f039c22e 100644 --- a/include/paimon/utils/read_ahead_cache.h +++ b/src/paimon/common/utils/read_ahead_cache.h @@ -27,87 +27,31 @@ #include #include "paimon/fs/file_system.h" -#include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/result.h" #include "paimon/status.h" +#include "paimon/utils/prefetch_cache_config.h" #include "paimon/visibility.h" namespace paimon { -/// PrefetchCacheMode -/// Cache prefetch switch modes. -/// Controls whether to enable cache prefetching under different circumstances, such as queries with -/// predicates or bitmap indexes. -/// -/// - ALWAYS: Enable cache in all scenarios. -/// - EXCLUDE_PREDICATE: Disable cache when query has predicates. -/// - EXCLUDE_BITMAP: Disable cache when using bitmap index. -/// - EXCLUDE_BITMAP_OR_PREDICATE: Disable cache if query has predicates or bitmap index. -/// - NEVER: Always disable cache. -enum class PAIMON_EXPORT PrefetchCacheMode { - ALWAYS = 1, - EXCLUDE_PREDICATE = 2, - EXCLUDE_BITMAP = 3, - EXCLUDE_BITMAP_OR_PREDICATE = 4, - NEVER = 5 -}; +class Metrics; -/// Configuration parameters for the read-ahead cache behavior. -/// -/// This struct controls various limits and prefetching strategies used by -/// ReadAheadCache to balance memory usage, I/O efficiency, and latency hiding. -class PAIMON_EXPORT CacheConfig { +/// Metric names for the read-ahead cache. +class PAIMON_EXPORT ReadAheadCacheMetrics { public: - CacheConfig(); - CacheConfig(uint64_t buffer_size_limit, uint64_t range_size_limit, uint64_t hole_size_limit, - uint64_t pre_buffer_limit); - - /// Returns the maximum total size (in bytes) of cached data. - uint64_t GetBufferSizeLimit() const { - return buffer_size_limit_; - } - - /// Sets the maximum total size (in bytes) of cached data. - void SetBufferSizeLimit(uint64_t buffer_size_limit) { - buffer_size_limit_ = buffer_size_limit; - } - - /// Returns the maximum allowed size (in bytes) for a single cached range. - uint64_t GetRangeSizeLimit() const { - return range_size_limit_; - } - - /// Sets the maximum allowed size (in bytes) for a single cached range. - void SetRangeSizeLimit(uint64_t range_size_limit) { - range_size_limit_ = range_size_limit; - } - - /// Returns the maximum gap size (in bytes) considered mergeable between adjacent ranges. - uint64_t GetHoleSizeLimit() const { - return hole_size_limit_; - } - - /// Sets the maximum gap size (in bytes) considered mergeable between adjacent ranges. - void SetHoleSizeLimit(uint64_t hole_size_limit) { - hole_size_limit_ = hole_size_limit; - } - - /// Returns the maximum size to pre-buffer ahead of the current read position. - uint64_t GetPreBufferLimit() const { - return pre_buffer_limit_; - } - - /// Sets the maximum size to pre-buffer ahead of the current read position. - void SetPreBufferLimit(uint64_t pre_buffer_limit) { - pre_buffer_limit_ = pre_buffer_limit; - } - - private: - uint64_t buffer_size_limit_; - uint64_t range_size_limit_; - uint64_t hole_size_limit_; - uint64_t pre_buffer_limit_; + /// Number of non-zero-sized Read() requests issued to the cache. + static inline const char READ_COUNT[] = "read-ahead-cache.read.count"; + /// Total bytes requested by the Read() requests issued to the cache. + static inline const char READ_BYTES[] = "read-ahead-cache.read.bytes"; + static inline const char READ_HITS[] = "read-ahead-cache.read.hits"; + static inline const char READ_HIT_BYTES[] = "read-ahead-cache.read.hit-bytes"; + static inline const char READ_MISSES[] = "read-ahead-cache.read.misses"; + static inline const char READ_MISS_BYTES[] = "read-ahead-cache.read.miss-bytes"; + /// Number of prefetch IO requests actually issued to the underlying stream. + static inline const char IO_COUNT[] = "read-ahead-cache.io.count"; + /// Total bytes requested by the prefetch IOs issued to the underlying stream. + static inline const char IO_BYTES[] = "read-ahead-cache.io.bytes"; }; /// A byte range with offset and length. @@ -132,22 +76,15 @@ struct PAIMON_EXPORT ByteRange { } }; -/// A byte slice with buffer, offset and length. -struct PAIMON_EXPORT ByteSlice { - std::shared_ptr buffer = nullptr; - uint64_t offset = 0; - uint64_t length = 0; -}; - /// A read cache designed to hide IO latencies when reading. /// Prefetching strategy: When a range is read, the cache will prefetch up to /// `pre_buffer_range_count` additional adjacent ranges ahead of the requested offset. This helps /// hide I/O latency for sequential access. Example: If you read range [0, 100), and /// pre_buffer_range_count=2, the next two configured ranges will also be prefetched. /// -/// Eviction policy: The cache uses a simple FIFO eviction policy based on total cached byte size. -/// When adding new ranges would exceed `buffer_size_limit`, the oldest cached ranges are evicted -/// first until there is enough space for the new data. +/// The cache never evicts: every published range stays cached until +/// ReleaseBuffers() or Reset(). It is meant to hold the prefetched ranges of +/// a single data file, whose size is bounded by the reader's scan scope. class PAIMON_EXPORT ReadAheadCache { public: /// Construct a read cache with given options @@ -162,11 +99,30 @@ class PAIMON_EXPORT ReadAheadCache { /// on the cache configuration. Status Init(std::vector&& ranges); - /// Read a range previously provided to Init(). + /// Read a range previously provided to Init(), copying the cached data + /// directly into the given destination buffer. + /// + /// Multi-segment hits are copied into `dest` segment by segment, without + /// an intermediate assembled buffer. /// @param range The byte range to read. - /// @return The byte slice containing the requested data. If the data is not yet cached - /// (cache miss), the returned `ByteSlice` will have a null buffer (`buffer == nullptr`) - Result Read(const ByteRange& range); + /// @param dest Destination buffer with at least `range.length` bytes. + /// @return true if the range was served from the cache and `dest` was + /// filled; false on cache miss (`dest` is left untouched). + Result Read(const ByteRange& range, char* dest); + + /// Start fetching the first batch of pending ranges immediately. + /// Init() only registers the ranges; without Warmup() the first fetch starts + /// when the first Read() arrives, racing the caller's own miss fetch. + void Warmup(); + + /// Collect hit/miss counters of Read() calls and the prefetch IO + /// counters into the given metrics as counters named after + /// `ReadAheadCacheMetrics`. Only reads issued through Read() are counted + /// as hits/misses; prefetch fetches dispatched by the cache itself are + /// counted in the fetch counters instead. + /// @param metrics The metrics to write the counters into. A null + /// pointer or a null shared pointer is a no-op. + void CollectMetrics(std::shared_ptr* metrics) const; /// Reset the cache to its initial state, clearing all cached data and configuration. /// @@ -175,6 +131,14 @@ class PAIMON_EXPORT ReadAheadCache { /// After calling Reset, the cache can be safely re-initialized with new ranges. void Reset(); + /// Release all cached buffers and pending ranges while keeping the hit/miss + /// counters intact. + /// + /// Unlike Reset(), the counters recorded by Read() remain readable through + /// CollectMetrics() afterwards, so this is safe to call when the owning reader + /// is closed while its metrics are still being aggregated. + void ReleaseBuffers(); + private: class Impl; std::unique_ptr impl_; diff --git a/src/paimon/common/utils/read_ahead_cache_test.cpp b/src/paimon/common/utils/read_ahead_cache_test.cpp index 67625422..e7900b7b 100644 --- a/src/paimon/common/utils/read_ahead_cache_test.cpp +++ b/src/paimon/common/utils/read_ahead_cache_test.cpp @@ -17,12 +17,18 @@ * under the License. */ -#include "paimon/utils/read_ahead_cache.h" +#include "paimon/common/utils/read_ahead_cache.h" +#include #include +#include +#include #include #include "gtest/gtest.h" +#include "paimon/common/factories/io_hook.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/fs/file_system.h" #include "paimon/fs/file_system_factory.h" #include "paimon/testing/utils/testharness.h" @@ -61,8 +67,94 @@ TestCacheEnv CreateTestFileAndCache(const std::string& filename, const std::stri return {path, cache, pool}; } +// Assert that reading the range is a cache hit filling the destination with +// the expected content. +void AssertReadEquals(const ByteRange& range, const std::string& expected, ReadAheadCache* cache) { + std::string dest(std::max(range.length, 1), 'X'); + bool hit = false; + ASSERT_OK_AND_ASSIGN(hit, cache->Read(range, dest.data())); + ASSERT_TRUE(hit) << expected; + EXPECT_EQ(expected, std::string_view(dest.data(), range.length)); +} + +// Assert that reading the range misses and leaves the destination untouched. +void AssertReadMiss(const ByteRange& range, ReadAheadCache* cache) { + std::string dest(std::max(range.length, 1), 'X'); + bool hit = true; + ASSERT_OK_AND_ASSIGN(hit, cache->Read(range, dest.data())); + ASSERT_FALSE(hit); + EXPECT_EQ(std::string(dest.size(), 'X'), dest); +} + +// An InputStream wrapper that holds ReadAsync callbacks until ReleaseAll() is +// called, letting tests observe the cache while prefetch IOs are in flight. +class GatedAsyncInputStream : public InputStream { + public: + explicit GatedAsyncInputStream(std::shared_ptr inner) : inner_(std::move(inner)) {} + + Status Close() override { + return inner_->Close(); + } + Status Seek(int64_t offset, SeekOrigin origin) override { + return inner_->Seek(offset, origin); + } + Result GetPos() const override { + return inner_->GetPos(); + } + Result Read(char* buffer, int64_t size) override { + return inner_->Read(buffer, size); + } + Result Read(char* buffer, int64_t size, int64_t offset) override { + return inner_->Read(buffer, size, offset); + } + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + std::lock_guard lock(mutex_); + async_read_count_++; + pending_.push_back({buffer, size, offset, std::move(callback)}); + } + Result GetUri() const override { + return inner_->GetUri(); + } + Result Length() const override { + return inner_->Length(); + } + + int AsyncReadCount() { + std::lock_guard lock(mutex_); + return async_read_count_; + } + + /// Complete all held fetches against the underlying stream. + void ReleaseAll() { + std::vector taken; + { + std::lock_guard lock(mutex_); + taken = std::move(pending_); + pending_.clear(); + } + for (auto& read : taken) { + Result res = inner_->Read(read.buffer, read.size, read.offset); + read.callback(res.ok() ? Status::OK() : res.status()); + } + } + + private: + struct PendingRead { + char* buffer; + int64_t size; + int64_t offset; + std::function callback; + }; + + std::shared_ptr inner_; + std::mutex mutex_; + std::vector pending_; + int async_read_count_ = 0; +}; + TEST(TestReadAheadCache, TestBasics) { - CacheConfig config(/*buffer_size_limit=*/256 * 1024 * 1024, /*range_size_limit=*/10, + CacheConfig config(/*range_size_limit=*/10, /*hole_size_limit=*/2, /*pre_buffer_limit=*/128 * 1024 * 1024); std::string content = "abcdefghijklmnopqrstuvwxyz"; auto env = CreateTestFileAndCache( @@ -70,90 +162,380 @@ TEST(TestReadAheadCache, TestBasics) { {{1, 2}, {3, 2}, {8, 2}, {10, 4}, {14, 0}, {15, 4}, {20, 2}, {25, 0}}); auto& cache = *env.cache; - auto assert_slice_equal = [](const ByteSlice& slice, const std::string& expected) { - ASSERT_TRUE(slice.buffer) << expected; - EXPECT_EQ(expected, std::string_view(slice.buffer->data() + slice.offset, slice.length)); - }; - - ByteSlice slice; - - ASSERT_OK_AND_ASSIGN(slice, cache.Read({20, 2})); - assert_slice_equal(slice, "uv"); + AssertReadEquals({20, 2}, "uv", &cache); + AssertReadEquals({1, 2}, "bc", &cache); + AssertReadEquals({3, 2}, "de", &cache); + AssertReadEquals({8, 2}, "ij", &cache); + AssertReadEquals({10, 4}, "klmn", &cache); + AssertReadEquals({15, 4}, "pqrs", &cache); + AssertReadEquals({19, 3}, "tuv", &cache); - ASSERT_OK_AND_ASSIGN(slice, cache.Read({1, 2})); - assert_slice_equal(slice, "bc"); + // Zero-sized reads are immediate hits touching nothing. + AssertReadEquals({14, 0}, "", &cache); + AssertReadEquals({25, 0}, "", &cache); - ASSERT_OK_AND_ASSIGN(slice, cache.Read({3, 2})); - assert_slice_equal(slice, "de"); - - ASSERT_OK_AND_ASSIGN(slice, cache.Read({8, 2})); - assert_slice_equal(slice, "ij"); - - ASSERT_OK_AND_ASSIGN(slice, cache.Read({10, 4})); - assert_slice_equal(slice, "klmn"); + // Non-cached ranges miss and leave the destination untouched. + AssertReadMiss({20, 3}, &cache); + AssertReadMiss({0, 3}, &cache); + AssertReadMiss({25, 2}, &cache); +} - ASSERT_OK_AND_ASSIGN(slice, cache.Read({15, 4})); - assert_slice_equal(slice, "pqrs"); +// Test that a read spanning several adjacent cache entries is served from the +// contiguous run of entries and counted as a single hit. +TEST(TestReadAheadCache, TestMultiSegmentContiguousHit) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + // A single 25-byte range exceeds range_size_limit, so Init() coalesces it + // into three adjacent entries: {0,10}, {10,10} and {20,5}. + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 25}}); + auto& cache = *env.cache; - ASSERT_OK_AND_ASSIGN(slice, cache.Read({19, 3})); - assert_slice_equal(slice, "tuv"); + // Spans all three entries. + AssertReadEquals({5, 20}, "fghijklmnopqrstuvwxy", &cache); - // Zero-sized - ASSERT_OK_AND_ASSIGN(slice, cache.Read({14, 0})); - assert_slice_equal(slice, ""); - ASSERT_OK_AND_ASSIGN(slice, cache.Read({25, 0})); - assert_slice_equal(slice, ""); + // Spans the first two entries only, trimming both ends of the run. + AssertReadEquals({5, 10}, "fghijklmno", &cache); - // Non-cached ranges + // Runs past the end of the last entry: no contiguous cover, a miss. + AssertReadMiss({5, 21}, &cache); - ASSERT_FALSE(cache.Read({20, 3}).value().buffer); - ASSERT_FALSE(cache.Read({0, 3}).value().buffer); - ASSERT_FALSE(cache.Read({25, 2}).value().buffer); + // A multi-segment hit counts once with the full requested length. + std::shared_ptr metrics = std::make_shared(); + cache.CollectMetrics(&metrics); + ASSERT_OK_AND_ASSIGN(uint64_t hits, metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_EQ(hits, 2u); + ASSERT_OK_AND_ASSIGN(uint64_t hit_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES)); + ASSERT_EQ(hit_bytes, 30u); + ASSERT_OK_AND_ASSIGN(uint64_t misses, metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, 1u); + ASSERT_OK_AND_ASSIGN(uint64_t miss_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); + ASSERT_EQ(miss_bytes, 21u); } // Test repeated reads to the same range to ensure cache reuse. TEST(TestReadAheadCache, TestRepeatedReadCacheReuse) { - CacheConfig config(/*buffer_size_limit=*/64, /*range_size_limit=*/10, + CacheConfig config(/*range_size_limit=*/10, /*hole_size_limit=*/2, /*pre_buffer_limit=*/64); std::string content = "abcdefghijklmnopqrstuvwxyz"; auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}, {7, 5}}); auto& cache = *env.cache; - ByteSlice slice; - ASSERT_OK_AND_ASSIGN(slice, cache.Read({0, 5})); - ASSERT_TRUE(slice.buffer); - std::string first_read(slice.buffer->data() + slice.offset, slice.length); - ASSERT_EQ(first_read, "abcde"); - - ASSERT_OK_AND_ASSIGN(slice, cache.Read({0, 5})); - ASSERT_TRUE(slice.buffer); - std::string second_read(slice.buffer->data() + slice.offset, slice.length); - ASSERT_EQ(second_read, "abcde"); + AssertReadEquals({0, 5}, "abcde", &cache); + AssertReadEquals({0, 5}, "abcde", &cache); } -// Test cache eviction when buffer size is limited. -TEST(TestReadAheadCache, TestCacheEviction) { - CacheConfig config(/*buffer_size_limit=*/10, /*range_size_limit=*/5, - /*hole_size_limit=*/2, /*pre_buffer_limit=*/10); +// The cache never evicts: every prefetched range stays cached until +// ReleaseBuffers()/Reset(), regardless of how much data accumulates. +TEST(TestReadAheadCache, TestNoEvictionKeepsAllRanges) { + CacheConfig config(/*range_size_limit=*/5, /*hole_size_limit=*/2, + /*pre_buffer_limit=*/10); std::string content = "abcdefghijklmnopqrstuvwxyz"; auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}, {8, 5}, {16, 5}}); auto& cache = *env.cache; - ByteSlice slice; - ASSERT_OK_AND_ASSIGN(slice, cache.Read({0, 5})); - ASSERT_TRUE(slice.buffer); - std::string first_read(slice.buffer->data() + slice.offset, slice.length); - ASSERT_EQ(first_read, "abcde"); - - // Reading another range should evict the first one due to buffer size limit - ASSERT_OK_AND_ASSIGN(slice, cache.Read({8, 5})); - ASSERT_TRUE(slice.buffer); - std::string second_read(slice.buffer->data() + slice.offset, slice.length); - ASSERT_EQ(second_read, "ijklm"); - - // The first range should now be a cache miss (buffer is nullptr) - auto miss = cache.Read({0, 5}); - ASSERT_FALSE(miss.value().buffer); + AssertReadEquals({0, 5}, "abcde", &cache); + + // Reading further ranges keeps the earlier ones cached. + AssertReadEquals({8, 5}, "ijklm", &cache); + AssertReadEquals({16, 5}, "qrstu", &cache); + AssertReadEquals({0, 5}, "abcde", &cache); +} + +// Test that Read() hits and misses are recorded in the cache metrics. +TEST(TestReadAheadCache, TestMetrics) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/128 * 1024 * 1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}, {8, 5}}); + auto& cache = *env.cache; + + AssertReadEquals({0, 5}, "abcde", &cache); + // Out of any cached range: a miss. + AssertReadMiss({20, 3}, &cache); + + std::shared_ptr metrics = std::make_shared(); + cache.CollectMetrics(&metrics); + // Both Read() requests are counted, regardless of hit or miss. + ASSERT_OK_AND_ASSIGN(uint64_t read_count, + metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); + ASSERT_EQ(read_count, 2u); + ASSERT_OK_AND_ASSIGN(uint64_t read_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_BYTES)); + ASSERT_EQ(read_bytes, 8u); + ASSERT_OK_AND_ASSIGN(uint64_t hits, metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_EQ(hits, 1u); + ASSERT_OK_AND_ASSIGN(uint64_t hit_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES)); + ASSERT_EQ(hit_bytes, 5u); + ASSERT_OK_AND_ASSIGN(uint64_t misses, metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, 1u); + ASSERT_OK_AND_ASSIGN(uint64_t miss_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); + ASSERT_EQ(miss_bytes, 3u); + // The hit prefetches both pending ranges in one window: two IO requests + // for 10 bytes in total; the miss issues no further fetch. + ASSERT_OK_AND_ASSIGN(uint64_t io_count, metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 2u); + ASSERT_OK_AND_ASSIGN(uint64_t io_bytes, metrics->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(io_bytes, 10u); +} + +// Test that ReleaseBuffers() drops the cached data but keeps the hit/miss counters +// readable, while Reset() zeroes them as well. +TEST(TestReadAheadCache, TestReleaseBuffersKeepsMetrics) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/128 * 1024 * 1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); + auto& cache = *env.cache; + + AssertReadEquals({0, 5}, "abcde", &cache); + + cache.ReleaseBuffers(); + + // The previously cached range is gone: the read now misses. + AssertReadMiss({0, 5}, &cache); + + std::shared_ptr metrics = std::make_shared(); + cache.CollectMetrics(&metrics); + // The read counters survive ReleaseBuffers() as well. + ASSERT_OK_AND_ASSIGN(uint64_t read_count, + metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); + ASSERT_EQ(read_count, 2u); + ASSERT_OK_AND_ASSIGN(uint64_t read_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_BYTES)); + ASSERT_EQ(read_bytes, 10u); + ASSERT_OK_AND_ASSIGN(uint64_t hits, metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_EQ(hits, 1u); + ASSERT_OK_AND_ASSIGN(uint64_t hit_bytes, + metrics->GetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES)); + ASSERT_EQ(hit_bytes, 5u); + ASSERT_OK_AND_ASSIGN(uint64_t misses, metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, 1u); + // The io counters survive ReleaseBuffers() as well. + ASSERT_OK_AND_ASSIGN(uint64_t io_count, metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1u); + ASSERT_OK_AND_ASSIGN(uint64_t io_bytes, metrics->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(io_bytes, 5u); + + // Reset() clears the counters too. + cache.Reset(); + std::shared_ptr reset_metrics = std::make_shared(); + cache.CollectMetrics(&reset_metrics); + ASSERT_OK_AND_ASSIGN(read_count, reset_metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); + ASSERT_EQ(read_count, 0u); + ASSERT_OK_AND_ASSIGN(hits, reset_metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_EQ(hits, 0u); + ASSERT_OK_AND_ASSIGN(misses, reset_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, 0u); + ASSERT_OK_AND_ASSIGN(io_count, reset_metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 0u); + ASSERT_OK_AND_ASSIGN(io_bytes, reset_metrics->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_EQ(io_bytes, 0u); +} + +// Test that a failed prefetch surfaces as an error Status from Read(), not as +// a miss: the entry exists from the moment its fetch is submitted and its +// future carries the IO error. +TEST(TestReadAheadCache, TestPrefetchIOErrorPropagation) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto io_hook = paimon::IOHook::GetInstance(); + + // Single entry: the prefetch is the first IO after the hook is armed. + { + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 10}}); + paimon::ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); + io_hook->Reset(0, paimon::IOHook::Mode::RETURN_ERROR); + std::string dest(5, 'X'); + ASSERT_NOK_WITH_MSG(env.cache->Read({0, 5}, dest.data()), + "io hook triggered io error at position"); + } + + // Several adjacent entries: the error of any segment aborts the read. + { + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 25}}); + paimon::ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); + io_hook->Reset(1, paimon::IOHook::Mode::RETURN_ERROR); + std::string dest(20, 'X'); + ASSERT_NOK_WITH_MSG(env.cache->Read({0, 20}, dest.data()), + "io hook triggered io error at position"); + } +} + +// Test that Warmup() fetches the pending ranges up front so the first Read() +// issues no further IO, while without Warmup() the first Read() triggers the +// prefetch itself. +TEST(TestReadAheadCache, TestWarmupPrefetchesBeforeFirstRead) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env1 = CreateTestFileAndCache("data_file", content, config, {{0, 5}, {8, 5}}); + env1.cache->Warmup(); + auto env2 = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); + + auto io_hook = paimon::IOHook::GetInstance(); + paimon::ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); + // Any new IO fails: the warmed-up reads must be served without fetching. + io_hook->Reset(0, paimon::IOHook::Mode::RETURN_ERROR); + + AssertReadEquals({0, 5}, "abcde", env1.cache.get()); + AssertReadEquals({8, 5}, "ijklm", env1.cache.get()); + + // Without Warmup() the first Read() starts the prefetch and sees the error. + std::string dest(5, 'X'); + ASSERT_NOK(env2.cache->Read({0, 5}, dest.data())); +} + +// Warmup() without any pending ranges is a safe no-op. +TEST(TestReadAheadCache, TestWarmupWithEmptyRanges) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {}); + env.cache->Warmup(); + AssertReadMiss({0, 5}, env.cache.get()); +} + +// A reader racing an in-flight prefetch must find the published entry and wait +// on its future instead of missing and re-fetching the same bytes: entries are +// published under the lock before their fetch is dispatched. +TEST(TestReadAheadCache, TestInFlightEntryServesRacingReader) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::string path = dir->Str() + "/data_file"; + std::ofstream file(path, std::ios::binary); + ASSERT_TRUE(file.is_open()); + file.write(content.data(), content.size()); + ASSERT_FALSE(file.fail()); + file.close(); + ASSERT_OK_AND_ASSIGN(auto fs, FileSystemFactory::Get("local", path, {})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs->Open(path)); + auto gated = std::make_shared(std::move(in)); + + ReadAheadCache cache(gated, config, GetDefaultPool()); + ASSERT_OK(cache.Init({{0, 5}})); + cache.Warmup(); + + // The prefetch entry is published, but its fetch is still held. + ASSERT_EQ(gated->AsyncReadCount(), 1); + + // A racing reader blocks on the in-flight entry's future and is served + // from it once the fetch completes, without triggering a second fetch. + std::thread reader([&cache, &gated]() { + std::string dest(5, 'X'); + Result res = cache.Read({0, 5}, dest.data()); + EXPECT_TRUE(res.ok()); + if (res.ok()) { + EXPECT_TRUE(res.value()); + } + EXPECT_EQ("abcde", std::string_view(dest.data(), 5)); + EXPECT_EQ(gated->AsyncReadCount(), 1); + }); + // Give the reader time to block on the in-flight entry's future before + // completing the fetch. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + gated->ReleaseAll(); + reader.join(); + + std::shared_ptr metrics = std::make_shared(); + cache.CollectMetrics(&metrics); + ASSERT_OK_AND_ASSIGN(uint64_t misses, metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t io_count, metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_EQ(io_count, 1u); +} + +// Test that pre_buffer_limit truncates the prefetch window: only ranges within +// the window are fetched at once, later reads fetch the remaining batches. +TEST(TestReadAheadCache, TestPreBufferWindowLimit) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/0, /*pre_buffer_limit=*/10); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 10}, {16, 10}}); + auto& cache = *env.cache; + + auto io_hook = paimon::IOHook::GetInstance(); + paimon::ScopeGuard guard([&io_hook]() { io_hook->Clear(); }); + // IOCount() only counts while armed; INT64_MAX never triggers the error mode. + io_hook->Reset(INT64_MAX, paimon::IOHook::Mode::RETURN_ERROR); + + AssertReadEquals({0, 10}, "abcdefghij", &cache); + // The second range did not fit into the window: only one prefetch IO. + ASSERT_EQ(io_hook->IOCount(), 1); + + AssertReadEquals({16, 10}, "qrstuvwxyz", &cache); + // The second read triggered exactly one more prefetch IO. + ASSERT_EQ(io_hook->IOCount(), 2); + + // The range is cached now: re-reading it issues no IO at all. + io_hook->Reset(INT64_MAX, paimon::IOHook::Mode::RETURN_ERROR); + AssertReadEquals({16, 10}, "qrstuvwxyz", &cache); + ASSERT_EQ(io_hook->IOCount(), 0); +} + +// Test that Init() rejects a second call until the cache is reset. +TEST(TestReadAheadCache, TestDoubleInit) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); + auto& cache = *env.cache; + + Status status = cache.Init({{8, 5}}); + ASSERT_FALSE(status.ok()); + + // The original ranges still work. + AssertReadEquals({0, 5}, "abcde", &cache); +} + +// Test that the cache can be re-initialized after Reset() and serves the new ranges. +TEST(TestReadAheadCache, TestReinitAfterReset) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); + auto& cache = *env.cache; + + AssertReadEquals({0, 5}, "abcde", &cache); + + cache.Reset(); + ASSERT_OK(cache.Init({{3, 4}})); + AssertReadEquals({3, 4}, "defg", &cache); + + // The old ranges are gone. + AssertReadMiss({20, 2}, &cache); +} + +// Test that Init() merges ranges separated by a small hole, so a read +// spanning the hole is served by the single coalesced entry. +TEST(TestReadAheadCache, TestInitCoalescesSmallHoles) { + CacheConfig config(/*range_size_limit=*/1024, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + // Byte 5 sits in a 1-byte hole, within hole_size_limit: one entry {0,11}. + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}, {6, 5}}); + auto& cache = *env.cache; + + AssertReadEquals({4, 3}, "efg", &cache); +} + +// CollectMetrics() with a null metrics output is a safe no-op. +TEST(TestReadAheadCache, TestCollectMetricsWithNullMetrics) { + CacheConfig config(/*range_size_limit=*/10, + /*hole_size_limit=*/2, /*pre_buffer_limit=*/1024); + std::string content = "abcdefghijklmnopqrstuvwxyz"; + auto env = CreateTestFileAndCache("data_file", content, config, {{0, 5}}); + env.cache->CollectMetrics(/*metrics=*/nullptr); + std::shared_ptr null_metrics; + env.cache->CollectMetrics(&null_metrics); } } // namespace paimon::test diff --git a/src/paimon/common/utils/saturating_cast.h b/src/paimon/common/utils/saturating_cast.h new file mode 100644 index 00000000..cb9c6039 --- /dev/null +++ b/src/paimon/common/utils/saturating_cast.h @@ -0,0 +1,53 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +namespace paimon { + +/// Converts a double to int32_t or int64_t with Java's float-to-int or float-to-long saturation +/// policy: NaN converts to 0 and an out-of-range value saturates at the bounds of TargetType. +/// Narrower Java integer conversions require a subsequent narrowing step and are not supported by +/// this helper. A bare static_cast of an unrepresentable double is undefined behavior and diverges +/// across architectures (x86 cvttsd2si yields the "integer indefinite" value, while aarch64 fcvtzs +/// saturates), so doubles that are not provably in range must go through this helper. +template +inline TargetType SaturatingDoubleToInteger(double value) { + static_assert(std::is_same_v || std::is_same_v, + "TargetType must be int32_t or int64_t"); + if (std::isnan(value)) { + return 0; + } + // Comparing against the bounds converted to double keeps the final truncation defined: + // (double)INT64_MAX rounds up to 2^63, so every value that reaches the truncation is + // representable in TargetType. + if (value >= static_cast(std::numeric_limits::max())) { + return std::numeric_limits::max(); + } + if (value <= static_cast(std::numeric_limits::lowest())) { + return std::numeric_limits::lowest(); + } + return static_cast(value); +} + +} // namespace paimon diff --git a/src/paimon/common/utils/saturating_cast_test.cpp b/src/paimon/common/utils/saturating_cast_test.cpp new file mode 100644 index 00000000..cc6b74cf --- /dev/null +++ b/src/paimon/common/utils/saturating_cast_test.cpp @@ -0,0 +1,74 @@ +/* + * 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 "paimon/common/utils/saturating_cast.h" + +#include +#include + +#include "gtest/gtest.h" + +namespace paimon::test { + +TEST(SaturatingCastTest, TestInt64InRangeTruncatesTowardZero) { + ASSERT_EQ(SaturatingDoubleToInteger(0.0), 0); + ASSERT_EQ(SaturatingDoubleToInteger(1.9), 1); + ASSERT_EQ(SaturatingDoubleToInteger(-1.9), -1); + // 2^63 - 1024 is the largest double below 2^63: it stays on the truncation path. + ASSERT_EQ(SaturatingDoubleToInteger(9223372036854774784.0), 9223372036854774784LL); +} + +TEST(SaturatingCastTest, TestInt64Saturation) { + // (double)INT64_MAX rounds up to 2^63, so the boundary double already saturates. + ASSERT_EQ(SaturatingDoubleToInteger( + static_cast(std::numeric_limits::max())), + std::numeric_limits::max()); + ASSERT_EQ(SaturatingDoubleToInteger(1e300), std::numeric_limits::max()); + ASSERT_EQ(SaturatingDoubleToInteger(-1e300), std::numeric_limits::lowest()); + ASSERT_EQ(SaturatingDoubleToInteger(std::numeric_limits::infinity()), + std::numeric_limits::max()); + ASSERT_EQ(SaturatingDoubleToInteger(-std::numeric_limits::infinity()), + std::numeric_limits::lowest()); + // The lowest bound is exactly representable and must survive as a value. + ASSERT_EQ(SaturatingDoubleToInteger( + static_cast(std::numeric_limits::lowest())), + std::numeric_limits::lowest()); +} + +TEST(SaturatingCastTest, TestInt64NaNBecomesZero) { + // Java's (long)Double.NaN == 0. + ASSERT_EQ(SaturatingDoubleToInteger(std::numeric_limits::quiet_NaN()), 0); +} + +TEST(SaturatingCastTest, TestInt32Path) { + // SstFileWriter converts through the int32_t instantiation. + ASSERT_EQ(SaturatingDoubleToInteger(42.7), 42); + ASSERT_EQ(SaturatingDoubleToInteger(-42.7), -42); + // The int32_t bounds are exactly representable as doubles and saturate inclusively. + ASSERT_EQ(SaturatingDoubleToInteger(2147483647.0), + std::numeric_limits::max()); + ASSERT_EQ(SaturatingDoubleToInteger(2147483648.0), + std::numeric_limits::max()); + ASSERT_EQ(SaturatingDoubleToInteger(-2147483648.0), + std::numeric_limits::lowest()); + ASSERT_EQ(SaturatingDoubleToInteger(-2147483649.0), + std::numeric_limits::lowest()); + ASSERT_EQ(SaturatingDoubleToInteger(std::numeric_limits::quiet_NaN()), 0); +} + +} // namespace paimon::test diff --git a/src/paimon/common/utils/serialization_utils.h b/src/paimon/common/utils/serialization_utils.h index c1e97d03..449766ca 100644 --- a/src/paimon/common/utils/serialization_utils.h +++ b/src/paimon/common/utils/serialization_utils.h @@ -78,7 +78,9 @@ class SerializationUtils { if (PAIMON_UNLIKELY(bytes->size() < 4)) { return Status::Invalid(fmt::format("bytes size {} is less than 4", bytes->size())); } - int32_t arity = *(reinterpret_cast(bytes->data())); + // The buffer is byte-filled, so memcpy avoids the strict-aliasing UB of reinterpret_cast. + int32_t arity; + memcpy(&arity, bytes->data(), sizeof(int32_t)); if (SystemByteOrder() == ByteOrder::PAIMON_LITTLE_ENDIAN) { arity = EndianSwapValue(arity); } diff --git a/src/paimon/common/utils/serialization_utils_test.cpp b/src/paimon/common/utils/serialization_utils_test.cpp index 5e612ff2..56a39a29 100644 --- a/src/paimon/common/utils/serialization_utils_test.cpp +++ b/src/paimon/common/utils/serialization_utils_test.cpp @@ -19,7 +19,20 @@ #include "paimon/common/utils/serialization_utils.h" +#include +#include +#include + #include "gtest/gtest.h" +#include "paimon/common/data/binary_row_writer.h" +#include "paimon/common/data/binary_string.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/io/data_input_stream.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -37,4 +50,39 @@ TEST_F(SerializationUtilsTest, TestSerializeBinaryRow) { ASSERT_TRUE(bytes); } +TEST_F(SerializationUtilsTest, TestDeserializeBinaryRowFromStream) { + std::shared_ptr memory_pool = GetDefaultPool(); + // a row with mixed field types, including negative integers and a string + BinaryRow row(3); + BinaryRowWriter writer(&row, 0, memory_pool.get()); + writer.WriteInt(0, -123456); + writer.WriteLong(1, static_cast(-9000000000)); + writer.WriteString(2, BinaryString::FromString("hello paimon!", memory_pool.get())); + writer.Complete(); + + // the first 4 bytes on the wire are the big-endian arity (Java-compatible format) + std::shared_ptr bytes = SerializationUtils::SerializeBinaryRow(row, memory_pool.get()); + ASSERT_TRUE(bytes); + ASSERT_GE(bytes->size(), 4); + ASSERT_EQ(static_cast(bytes->data()[0]), 0x00); + ASSERT_EQ(static_cast(bytes->data()[1]), 0x00); + ASSERT_EQ(static_cast(bytes->data()[2]), 0x00); + ASSERT_EQ(static_cast(bytes->data()[3]), 0x03); + + // round-trip through the stream overloads, which fill a fresh byte buffer on deserialize + MemorySegmentOutputStream out(MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, memory_pool); + ASSERT_OK(SerializationUtils::SerializeBinaryRow(row, &out)); + auto stream_bytes = + MemorySegmentUtils::CopyToBytes(out.Segments(), 0, out.CurrentSize(), memory_pool.get()); + auto input_stream = + std::make_shared(stream_bytes->data(), stream_bytes->size()); + DataInputStream data_input_stream(input_stream); + ASSERT_OK_AND_ASSIGN(BinaryRow de_row, SerializationUtils::DeserializeBinaryRow( + &data_input_stream, memory_pool.get())); + ASSERT_EQ(de_row.GetFieldCount(), 3); + ASSERT_EQ(de_row.GetInt(0), -123456); + ASSERT_EQ(de_row.GetLong(1), static_cast(-9000000000)); + ASSERT_EQ(de_row.GetString(2).ToString(), "hello paimon!"); +} + } // namespace paimon::test diff --git a/src/paimon/core/append/append_only_writer_test.cpp b/src/paimon/core/append/append_only_writer_test.cpp index 8ef6e135..11401609 100644 --- a/src/paimon/core/append/append_only_writer_test.cpp +++ b/src/paimon/core/append/append_only_writer_test.cpp @@ -56,9 +56,11 @@ #include "paimon/core/stats/simple_stats.h" #include "paimon/core/utils/commit_increment.h" #include "paimon/defs.h" +#include "paimon/file_index/file_index_format.h" #include "paimon/format/file_format_factory.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/io/byte_array_input_stream.h" #include "paimon/memory/memory_pool.h" #include "paimon/record_batch.h" #include "paimon/testing/utils/binary_row_generator.h" @@ -404,6 +406,41 @@ TEST_F(AppendOnlyWriterTest, TestWriteAndPrepareCommit) { ASSERT_OK(writer->Close()); } +TEST_F(AppendOnlyWriterTest, TestWritePublishesEmbeddedBitmapIndex) { + CoreOptions options = CreateOptions( + {{"file-index.bitmap.columns", "f0"}, {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}}); + auto schema = + arrow::schema({arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32())}); + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = CreatePathFactory(dir->Str(), "mock_format", options); + ASSERT_OK_AND_ASSIGN( + auto writer, CreateAppendOnlyWriter( + options, /*schema_id=*/0, schema, /*write_cols=*/std::nullopt, + /*max_sequence_number=*/-1, path_factory, compact_manager_, memory_pool_)); + + ASSERT_OK(writer->Write(CreateBatch(schema, R"([{"f0": 1, "f1": 10}, + {"f0": 2, "f1": 20}, + {"f0": 1, "f1": 30}])"))); + ASSERT_OK_AND_ASSIGN(CommitIncrement increment, + writer->PrepareCommit(/*wait_compaction=*/true)); + const auto& files = increment.GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + ASSERT_TRUE(files[0]->embedded_index); + ASSERT_TRUE(files[0]->extra_files.empty()); + + auto input = std::make_shared(files[0]->embedded_index->data(), + files[0]->embedded_index->size()); + ASSERT_OK_AND_ASSIGN(auto index_reader, FileIndexFormat::CreateReader(input, memory_pool_)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(auto column_readers, index_reader->ReadColumnIndex("f0", &c_schema)); + ASSERT_EQ(1, column_readers.size()); + ASSERT_OK_AND_ASSIGN(auto result, column_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0,2}", result->ToString()); + ASSERT_OK(writer->Close()); +} + TEST_F(AppendOnlyWriterTest, TestWriteAndClose) { std::map raw_options; raw_options[Options::FILE_FORMAT] = "orc"; diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index 578950fb..6320ec57 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -38,6 +38,7 @@ #include "paimon/defs.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" +#include "paimon/statistics_mode.h" #include "paimon/status.h" namespace paimon { @@ -386,11 +387,14 @@ struct CoreOptions::Impl { int64_t manifest_target_file_size = 8 * 1024 * 1024; int64_t deletion_vector_target_file_size = 2 * 1024 * 1024; int64_t manifest_full_compaction_file_size = 16 * 1024 * 1024; + int64_t file_index_in_manifest_threshold = 500; int64_t write_buffer_size = 256 * 1024 * 1024; int64_t commit_timeout = std::numeric_limits::max(); int64_t commit_min_retry_wait = 10; int64_t commit_max_retry_wait = 10 * 1000; + bool realtime_enabled = false; int64_t realtime_read_view_ttl_millis = 5 * 60 * 1000; + StatisticsMode realtime_store_statistics_mode = StatisticsMode::NONE; std::shared_ptr file_format; std::shared_ptr file_system; @@ -425,6 +429,7 @@ struct CoreOptions::Impl { int32_t manifest_merge_min_count = 30; int32_t scan_manifest_entry_cache_max_snapshots = 0; + bool scan_manifest_entry_lazy_decode_enabled = true; int32_t read_batch_size = 1024; int32_t write_batch_size = 1024; int32_t local_sort_max_num_file_handles = 128; @@ -803,12 +808,6 @@ struct CoreOptions::Impl { std::string scan_timestamp_str; PAIMON_RETURN_NOT_OK(parser.Parse(Options::SCAN_TIMESTAMP, &scan_timestamp_str)); PAIMON_RETURN_NOT_OK(parser.Parse(Options::SCAN_TIMESTAMP_MILLIS, &scan_timestamp_millis)); - PAIMON_RETURN_NOT_OK(parser.ParseTimeDuration(Options::REALTIME_READ_VIEW_TTL, - &realtime_read_view_ttl_millis)); - if (realtime_read_view_ttl_millis <= 0) { - return Status::Invalid( - fmt::format("{} must be positive", Options::REALTIME_READ_VIEW_TTL)); - } if (scan_timestamp_millis != std::nullopt && !scan_timestamp_str.empty()) { return Status::Invalid( "scan.timestamp-millis and scan.timestamp cannot be set at the same time"); @@ -827,6 +826,8 @@ struct CoreOptions::Impl { return Status::Invalid(fmt::format("{} must be non-negative", Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS)); } + PAIMON_RETURN_NOT_OK(parser.Parse(Options::SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED, + &scan_manifest_entry_lazy_decode_enabled)); // Parse scan.fallback-branch - fallback branch when partition not found PAIMON_RETURN_NOT_OK(parser.Parse(Options::SCAN_FALLBACK_BRANCH, &scan_fallback_branch)); // Parse branch - branch name, default "main" @@ -836,8 +837,39 @@ struct CoreOptions::Impl { return Status::OK(); } + // Parse statistics collected by real-time stores. + Status ParseRealtimeStoreStatisticsMode(const ConfigParser& parser) { + std::string statistics_mode = "none"; + PAIMON_RETURN_NOT_OK(parser.Parse(Options::REALTIME_STORE_STATS_MODE, &statistics_mode)); + statistics_mode = StringUtils::ToLowerCase(statistics_mode); + if (statistics_mode == "none") { + realtime_store_statistics_mode = StatisticsMode::NONE; + } else if (statistics_mode == "full") { + realtime_store_statistics_mode = StatisticsMode::FULL; + } else { + return Status::Invalid( + fmt::format("{} must be 'none' or 'full'", Options::REALTIME_STORE_STATS_MODE)); + } + return Status::OK(); + } + + // Parse real-time write and read configurations. + Status ParseRealtimeOptions(const ConfigParser& parser) { + PAIMON_RETURN_NOT_OK(parser.Parse(Options::REALTIME_ENABLED, &realtime_enabled)); + PAIMON_RETURN_NOT_OK(parser.ParseTimeDuration(Options::REALTIME_READ_VIEW_TTL, + &realtime_read_view_ttl_millis)); + if (realtime_read_view_ttl_millis <= 0) { + return Status::Invalid( + fmt::format("{} must be positive", Options::REALTIME_READ_VIEW_TTL)); + } + return ParseRealtimeStoreStatisticsMode(parser); + } + // Parse index-related configurations: file index, global index. Status ParseIndexOptions(const ConfigParser& parser) { + // Parse file-index.in-manifest-threshold - max inline file index size, default 500B + PAIMON_RETURN_NOT_OK(parser.ParseMemorySize(Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, + &file_index_in_manifest_threshold)); // Parse file-index.read.enabled - whether to enable reading file index, default true PAIMON_RETURN_NOT_OK( parser.Parse(Options::FILE_INDEX_READ_ENABLED, &file_index_read_enabled)); @@ -1046,6 +1078,7 @@ Result CoreOptions::FromMap( PAIMON_RETURN_NOT_OK(impl->ParseCommitOptions(parser)); PAIMON_RETURN_NOT_OK(impl->ParseMergeAndSequenceOptions(parser)); PAIMON_RETURN_NOT_OK(impl->ParseDeletionVectorOptions(parser)); + PAIMON_RETURN_NOT_OK(impl->ParseRealtimeOptions(parser)); PAIMON_RETURN_NOT_OK(impl->ParseScanAndBranchOptions(parser)); PAIMON_RETURN_NOT_OK(impl->ParseIndexOptions(parser)); PAIMON_RETURN_NOT_OK(impl->ParseCompactionOptions(parser)); @@ -1158,14 +1191,26 @@ std::optional CoreOptions::GetScanTimestampMillis() const { return impl_->scan_timestamp_millis; } +bool CoreOptions::RealtimeEnabled() const { + return impl_->realtime_enabled; +} + int64_t CoreOptions::GetRealtimeReadViewTtlMillis() const { return impl_->realtime_read_view_ttl_millis; } +StatisticsMode CoreOptions::GetRealtimeStoreStatisticsMode() const { + return impl_->realtime_store_statistics_mode; +} + int32_t CoreOptions::GetScanManifestEntryCacheMaxSnapshots() const { return impl_->scan_manifest_entry_cache_max_snapshots; } +bool CoreOptions::ScanManifestEntryLazyDecodeEnabled() const { + return impl_->scan_manifest_entry_lazy_decode_enabled; +} + int64_t CoreOptions::GetManifestTargetFileSize() const { return impl_->manifest_target_file_size; } @@ -1654,6 +1699,10 @@ bool CoreOptions::FileIndexReadEnabled() const { return impl_->file_index_read_enabled; } +int64_t CoreOptions::FileIndexInManifestThreshold() const { + return impl_->file_index_in_manifest_threshold; +} + std::optional CoreOptions::GetDataFileExternalPaths() const { return impl_->data_file_external_paths; } diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h index 9eb28988..85a4a7fd 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h @@ -40,6 +40,7 @@ #include "paimon/format/file_format.h" #include "paimon/fs/file_system.h" #include "paimon/result.h" +#include "paimon/statistics_mode.h" #include "paimon/table/source/startup_mode.h" #include "paimon/type_fwd.h" #include "paimon/visibility.h" @@ -106,8 +107,12 @@ class PAIMON_EXPORT CoreOptions { int64_t GetSourceSplitOpenFileCost() const; std::optional GetScanSnapshotId() const; std::optional GetScanTimestampMillis() const; + bool RealtimeEnabled() const; int64_t GetRealtimeReadViewTtlMillis() const; + /// Returns the statistics mode used by the real-time store. + StatisticsMode GetRealtimeStoreStatisticsMode() const; int32_t GetScanManifestEntryCacheMaxSnapshots() const; + bool ScanManifestEntryLazyDecodeEnabled() const; int64_t GetManifestTargetFileSize() const; std::shared_ptr GetCache() const; @@ -207,6 +212,7 @@ class PAIMON_EXPORT CoreOptions { bool NeedLookup() const; bool PrepareCommitWaitCompaction() const; bool FileIndexReadEnabled() const; + int64_t FileIndexInManifestThreshold() const; std::map GetFieldsSequenceGroups() const; bool PartialUpdateRemoveRecordOnDelete() const; diff --git a/src/paimon/core/core_options_test.cpp b/src/paimon/core/core_options_test.cpp index ebc127ed..c110623d 100644 --- a/src/paimon/core/core_options_test.cpp +++ b/src/paimon/core/core_options_test.cpp @@ -29,6 +29,7 @@ #include "paimon/core/options/expire_config.h" #include "paimon/defs.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/statistics_mode.h" #include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/utils/testharness.h" #include "paimon/testing/utils/timezone_guard.h" @@ -54,6 +55,8 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_EQ("__DEFAULT_PARTITION__", core_options.GetPartitionDefaultName()); ASSERT_EQ(std::nullopt, core_options.GetScanSnapshotId()); ASSERT_EQ(5 * 60 * 1000, core_options.GetRealtimeReadViewTtlMillis()); + ASSERT_FALSE(core_options.RealtimeEnabled()); + ASSERT_EQ(StatisticsMode::NONE, core_options.GetRealtimeStoreStatisticsMode()); ASSERT_EQ("zstd", core_options.GetFileCompression()); ASSERT_EQ("zstd", core_options.GetWriteFileCompression(0)); ASSERT_EQ("zstd", core_options.GetWriteFileCompression(3)); @@ -65,6 +68,7 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_EQ(30, core_options.GetManifestMergeMinCount()); ASSERT_FALSE(core_options.ManifestDeleteFileDropStats()); ASSERT_EQ(0, core_options.GetScanManifestEntryCacheMaxSnapshots()); + ASSERT_TRUE(core_options.ScanManifestEntryLazyDecodeEnabled()); ASSERT_EQ(nullptr, core_options.GetCache()); ASSERT_EQ(128 * 1024 * 1024L, core_options.GetSourceSplitTargetSize()); ASSERT_EQ(4 * 1024 * 1024L, core_options.GetSourceSplitOpenFileCost()); @@ -134,6 +138,7 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_EQ(std::nullopt, core_options.GetScanFallbackBranch()); ASSERT_EQ("main", core_options.GetBranch()); ASSERT_TRUE(core_options.FileIndexReadEnabled()); + ASSERT_EQ(500, core_options.FileIndexInManifestThreshold()); ASSERT_EQ(std::nullopt, core_options.GetDataFileExternalPaths()); ASSERT_EQ(ExternalPathStrategy::NONE, core_options.GetExternalPathStrategy()); ASSERT_TRUE(core_options.EnableAdaptivePrefetchStrategy()); @@ -217,6 +222,7 @@ TEST(CoreOptionsTest, TestFromMap) { {Options::SCAN_SNAPSHOT_ID, "5"}, {Options::SCAN_MODE, "from-snapshot-full"}, {Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS, "7"}, + {Options::SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED, "false"}, {Options::SNAPSHOT_NUM_RETAINED_MIN, "15"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "30"}, {Options::SNAPSHOT_EXPIRE_LIMIT, "20"}, @@ -248,6 +254,7 @@ TEST(CoreOptionsTest, TestFromMap) { {Options::SCAN_FALLBACK_BRANCH, "fallback"}, {Options::BRANCH, "rt"}, {Options::FILE_INDEX_READ_ENABLED, "false"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "2KB"}, {Options::DATA_FILE_EXTERNAL_PATHS, "FILE:///tmp/index"}, {Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY, "round-robin"}, {Options::FILE_COMPRESSION, "snappy"}, @@ -304,6 +311,7 @@ TEST(CoreOptionsTest, TestFromMap) { {Options::LOOKUP_REMOTE_LEVEL_THRESHOLD, "2"}, {Options::TABLE_READ_SEQUENCE_NUMBER_ENABLED, "true"}, {Options::KEY_VALUE_SEQUENCE_NUMBER_ENABLED, "true"}, + {Options::REALTIME_ENABLED, "true"}, {Options::BUCKET_FUNCTION_TYPE, "mod"}, {"fields.metrics.map.storage-layout", "shared-shredding"}, {"fields.metrics.map.shared-shredding.max-columns", "128"}, @@ -353,6 +361,7 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_TRUE(core_options.CommitDiscardDuplicateFiles()); ASSERT_EQ(5, core_options.GetScanSnapshotId().value_or(-1)); ASSERT_EQ(7, core_options.GetScanManifestEntryCacheMaxSnapshots()); + ASSERT_FALSE(core_options.ScanManifestEntryLazyDecodeEnabled()); ExpireConfig expire_config = core_options.GetExpireConfig(); ASSERT_EQ(15, expire_config.GetSnapshotRetainMin()); ASSERT_EQ(30, expire_config.GetSnapshotRetainMax()); @@ -398,6 +407,7 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_EQ(core_options.GetScanFallbackBranch(), std::optional("fallback")); ASSERT_EQ(core_options.GetBranch(), "rt"); ASSERT_FALSE(core_options.FileIndexReadEnabled()); + ASSERT_EQ(2 * 1024, core_options.FileIndexInManifestThreshold()); ASSERT_EQ(core_options.GetDataFileExternalPaths(), std::optional("FILE:///tmp/index")); ASSERT_EQ(core_options.GetExternalPathStrategy(), ExternalPathStrategy::ROUND_ROBIN); @@ -463,6 +473,7 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_EQ(10L * 1024 * 1024 * 1024, core_options.GetLookupCacheMaxDiskSize()); ASSERT_TRUE(core_options.TableReadSequenceNumberEnabled()); ASSERT_TRUE(core_options.KeyValueSequenceNumberEnabled()); + ASSERT_TRUE(core_options.RealtimeEnabled()); ASSERT_TRUE(core_options.LookupRemoteFileEnabled()); ASSERT_EQ(core_options.GetLookupRemoteLevelThreshold(), 2); ASSERT_EQ(BucketFunctionType::MOD, core_options.GetBucketFunctionType()); @@ -492,6 +503,7 @@ TEST(CoreOptionsTest, TestInvalidCase) { "invalid lookup mode: invalid"); ASSERT_NOK_WITH_MSG(CoreOptions::FromMap({{Options::LOOKUP_COMPACT_MAX_INTERVAL, "invalid"}}), "Invalid Config [lookup-compact.max-interval: invalid]"); + ASSERT_NOK(CoreOptions::FromMap({{Options::REALTIME_ENABLED, "invalid"}})); ASSERT_NOK_WITH_MSG( CoreOptions::FromMap({{Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS, "-1"}}), "scan.manifest-entry-cache.max-snapshots must be non-negative"); @@ -804,6 +816,19 @@ TEST(CoreOptionsTest, TestRealtimeReadViewTtlMillis) { "realtime.read-view-ttl must be positive"); } +TEST(CoreOptionsTest, TestRealtimeStoreStatisticsMode) { + ASSERT_OK_AND_ASSIGN(CoreOptions full_options, + CoreOptions::FromMap({{Options::REALTIME_STORE_STATS_MODE, "full"}})); + ASSERT_EQ(StatisticsMode::FULL, full_options.GetRealtimeStoreStatisticsMode()); + + ASSERT_OK_AND_ASSIGN(CoreOptions none_options, + CoreOptions::FromMap({{Options::REALTIME_STORE_STATS_MODE, "none"}})); + ASSERT_EQ(StatisticsMode::NONE, none_options.GetRealtimeStoreStatisticsMode()); + + ASSERT_NOK_WITH_MSG(CoreOptions::FromMap({{Options::REALTIME_STORE_STATS_MODE, "invalid"}}), + "realtime.store.stats-mode must be 'none' or 'full'"); +} + TEST(CoreOptionsTest, TestScanTimestampMillisExplicitMode) { ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::SCAN_MODE, "from-timestamp"}, diff --git a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp index 1fca5f37..5451c9d2 100644 --- a/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp +++ b/src/paimon/core/deletionvectors/apply_deletion_vector_batch_reader_test.cpp @@ -26,13 +26,13 @@ #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/executor.h" #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/mock/mock_format_reader_builder.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" -#include "paimon/utils/read_ahead_cache.h" namespace arrow { class Array; @@ -87,7 +87,7 @@ class ApplyDeletionVectorBatchReaderTest : public ::testing::Test, prefetch_batch_count, batch_size, prefetch_batch_count * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_, /*initialize_read_ranges=*/true, - /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, CacheConfig(), pool_)); + /*read_ahead_cache_enabled=*/true, CacheConfig(), pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/core/io/append_data_file_writer_factory.cpp b/src/paimon/core/io/append_data_file_writer_factory.cpp index e12374dd..d0b677a7 100644 --- a/src/paimon/core/io/append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/append_data_file_writer_factory.cpp @@ -24,6 +24,7 @@ #include "arrow/c/abi.h" #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/fs/file_system.h" @@ -57,6 +58,11 @@ AppendDataFileWriterFactory::CreateWriter() const { options_.GetFileCompression(), std::function(), schema_id_, seq_num_counter, file_source_, resources.stats_extractor, path_factory_->IsExternalPath(), write_cols_, pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); return std::unique_ptr>>( diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp new file mode 100644 index 00000000..5c97bb3d --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -0,0 +1,179 @@ +/* + * 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 "paimon/core/io/data_file_index_writer.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/file_index/file_index_writer.h" +#include "paimon/file_index/file_indexer.h" +#include "paimon/file_index/file_indexer_factory.h" +#include "paimon/fs/file_system.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/status.h" + +namespace paimon { +Result> DataFileIndexWriter::Create( + const std::shared_ptr& logical_schema, const FileIndexOptions& options, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool) { + assert(logical_schema); + assert(file_system); + assert(path_factory); + assert(pool); + std::vector writers; + writers.reserve(options.Definitions().size()); + for (const FileIndexDefinition& definition : options.Definitions()) { + if (SpecialFields::IsSystemField(definition.column_name)) { + return Status::Invalid( + fmt::format("File index column '{}' is a system field", definition.column_name)); + } + int32_t field_index = logical_schema->GetFieldIndex(definition.column_name); + if (field_index < 0) { + return Status::Invalid( + fmt::format("File index column '{}' does not exist in the write schema", + definition.column_name)); + } + std::shared_ptr field = logical_schema->field(field_index); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, + FileIndexerFactory::Get(definition.index_type, definition.options)); + if (!indexer) { + return Status::Invalid( + fmt::format("File index type '{}' is not registered", definition.index_type)); + } + ::ArrowSchema c_schema; + ArrowSchemaMarkReleased(&c_schema); + ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow::schema({field}), &c_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + indexer->CreateWriter(&c_schema, pool)); + writers.push_back( + {definition.column_name, definition.index_type, field_index, field, std::move(writer)}); + } + return std::unique_ptr(new DataFileIndexWriter( + std::move(writers), options.InManifestThreshold(), file_system, path_factory, pool)); +} + +DataFileIndexWriter::DataFileIndexWriter(std::vector&& writers, + int64_t in_manifest_threshold, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool) + : writers_(std::move(writers)), + in_manifest_threshold_(in_manifest_threshold), + file_system_(file_system), + path_factory_(path_factory), + pool_(pool) {} + +Status DataFileIndexWriter::AddBatch(const std::shared_ptr& logical_batch) { + if (finished_) { + return Status::Invalid("Data file index writer has already finished"); + } + for (const IndexWriterEntry& entry : writers_) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr projected, + arrow::StructArray::Make({logical_batch->field(entry.field_index)}, {entry.field})); + ::ArrowArray c_array; + ArrowArrayMarkReleased(&c_array); + ScopeGuard array_guard([&c_array]() { ArrowArrayRelease(&c_array); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*projected, &c_array)); + PAIMON_RETURN_NOT_OK(entry.writer->AddBatch(&c_array)); + } + return Status::OK(); +} + +Result> DataFileIndexWriter::SerializeContainer() { + FileIndexFormat::ColumnIndexes column_indexes; + for (const IndexWriterEntry& entry : writers_) { + PAIMON_ASSIGN_OR_RAISE(column_indexes[entry.column_name][entry.index_type], + entry.writer->SerializedBytes()); + } + + auto segment_output = std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + std::shared_ptr output = + std::make_shared(std::move(segment_output)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_writer, + FileIndexFormat::CreateWriter(output, pool_)); + PAIMON_RETURN_NOT_OK(format_writer->WriteColumnIndexes(column_indexes)); + PAIMON_RETURN_NOT_OK(format_writer->Close()); + return output->Finish(pool_.get()); +} + +Result DataFileIndexWriter::Finish(const std::string& data_file_path) { + if (finished_) { + return Status::Invalid("Data file index writer has already finished"); + } + finished_ = true; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, SerializeContainer()); + if (static_cast(bytes->size()) <= in_manifest_threshold_) { + return FileIndexWriteResult{bytes, {}}; + } + + external_index_path_ = path_factory_->ToFileIndexPath(data_file_path); + PAIMON_RETURN_NOT_OK(WriteExternal(external_index_path_.value(), bytes)); + return FileIndexWriteResult{nullptr, {PathUtil::GetName(external_index_path_.value())}}; +} + +Status DataFileIndexWriter::WriteExternal(const std::string& path, + const std::shared_ptr& bytes) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, + file_system_->Create(path, /*overwrite=*/false)); + ScopeGuard guard([this, &output]() { + if (output) { + [[maybe_unused]] Status _ = output->Close(); + } + Abort(); + }); + PAIMON_ASSIGN_OR_RAISE(int64_t written, + output->Write(bytes->data(), static_cast(bytes->size()))); + if (written != static_cast(bytes->size())) { + return Status::IOError(fmt::format("Short write for file index {}: expected {}, wrote {}", + path, bytes->size(), written)); + } + PAIMON_RETURN_NOT_OK(output->Flush()); + Status close_status = output->Close(); + output.reset(); + PAIMON_RETURN_NOT_OK(close_status); + guard.Release(); + return Status::OK(); +} + +void DataFileIndexWriter::Abort() { + if (external_index_path_) { + [[maybe_unused]] Status _ = file_system_->Delete(external_index_path_.value()); + } +} + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_index_writer.h b/src/paimon/core/io/data_file_index_writer.h new file mode 100644 index 00000000..883b719f --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer.h @@ -0,0 +1,100 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/core/io/file_index_options.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace arrow { +class Field; +class Schema; +class StructArray; +} // namespace arrow + +namespace paimon { + +class Bytes; +class DataFilePathFactory; +class FileIndexWriter; +class FileSystem; +class MemoryPool; + +struct FileIndexWriteResult { + std::shared_ptr embedded_index; + std::vector> extra_files; +}; + +/// Builds every configured column index for one data file. +class DataFileIndexWriter { + public: + static Result> Create( + const std::shared_ptr& logical_schema, const FileIndexOptions& options, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool); + + Status AddBatch(const std::shared_ptr& logical_batch); + + /// Finalizes and publishes all configured indexes. This is a terminal, one-shot operation. + /// + /// @param data_file_path Path of the data file associated with these indexes. + /// @return Embedded index bytes or the external index file name. + Result Finish(const std::string& data_file_path); + + void Abort(); + + const std::optional& ExternalIndexPath() const { + return external_index_path_; + } + + private: + struct IndexWriterEntry { + std::string column_name; + std::string index_type; + int32_t field_index; + std::shared_ptr field; + std::shared_ptr writer; + }; + + DataFileIndexWriter(std::vector&& writers, int64_t in_manifest_threshold, + const std::shared_ptr& file_system, + const std::shared_ptr& path_factory, + const std::shared_ptr& pool); + + Result> SerializeContainer(); + Status WriteExternal(const std::string& path, const std::shared_ptr& bytes); + + std::vector writers_; + int64_t in_manifest_threshold_; + std::shared_ptr file_system_; + std::shared_ptr path_factory_; + std::shared_ptr pool_; + std::optional external_index_path_; + bool finished_ = false; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp new file mode 100644 index 00000000..e9ab7940 --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -0,0 +1,253 @@ +/* + * 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 "paimon/core/io/data_file_index_writer.h" + +#include +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "arrow/type.h" +#include "gtest/gtest.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/file_index_options.h" +#include "paimon/defs.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/testing/mock/mock_file_system.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +struct CloseFailingState { + int32_t close_count = 0; + int32_t delete_count = 0; +}; + +class CloseFailingOutputStream : public MockOutputStream { + public: + explicit CloseFailingOutputStream(const std::shared_ptr& state) + : state_(state) {} + + Result Write(const char*, int64_t size) override { + return size; + } + + Status Close() override { + ++state_->close_count; + return Status::IOError("close failed"); + } + + private: + std::shared_ptr state_; +}; + +class CloseFailingFileSystem : public MockFileSystem { + public: + explicit CloseFailingFileSystem(const std::shared_ptr& state) + : state_(state) {} + + Result> Create(const std::string&, bool) const override { + return std::unique_ptr(new CloseFailingOutputStream(state_)); + } + + Status Delete(const std::string&, bool = true) const override { + ++state_->delete_count; + return Status::OK(); + } + + private: + std::shared_ptr state_; +}; + +} // namespace + +class DataFileIndexWriterTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + file_system_ = std::make_shared(); + directory_ = UniqueTestDirectory::Create(); + ASSERT_TRUE(directory_); + path_factory_ = std::make_shared(); + ASSERT_OK(path_factory_->Init(directory_->Str(), "orc", "data-", nullptr)); + schema_ = + arrow::schema({arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32())}); + } + + Result> CreateWriter( + const std::map& index_options) const { + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, + CoreOptions::FromMap(index_options, file_system_)); + PAIMON_ASSIGN_OR_RAISE(FileIndexOptions parsed, + FileIndexOptions::FromCoreOptions(core_options)); + return DataFileIndexWriter::Create(schema_, parsed, file_system_, path_factory_, pool_); + } + + std::shared_ptr CreateBatch(const std::string& json) const { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json) + .ValueOrDie(); + return checked_pointer_cast(array); + } + + Result> CreateReader( + const std::shared_ptr& bytes) const { + auto input = std::make_shared(bytes->data(), bytes->size()); + return FileIndexFormat::CreateReader(input, pool_); + } + + Result>> ReadColumn( + FileIndexFormat::Reader* reader, const std::string& column_name) const { + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, &c_schema)); + return reader->ReadColumnIndex(column_name, &c_schema); + } + + protected: + std::shared_ptr pool_; + std::shared_ptr file_system_; + std::unique_ptr directory_; + std::shared_ptr path_factory_; + std::shared_ptr schema_; +}; + +TEST_F(DataFileIndexWriterTest, TestBitmapAndRangeBitmapEmbeddedRoundTrip) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {"file-index.range-bitmap.columns", "f1"}, + {"file-index.range-bitmap.f1.chunk-size", "1KB"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}, + {"f0": 2, "f1": 20}])"))); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 30}, + {"f0": null, "f1": 40}])"))); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish("unused.orc")); + ASSERT_TRUE(result.embedded_index); + ASSERT_TRUE(result.extra_files.empty()); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(result.embedded_index)); + + ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bitmap_readers.size()); + ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0,2}", equal_result->ToString()); + ASSERT_OK_AND_ASSIGN(auto null_result, bitmap_readers[0]->VisitIsNull()); + ASSERT_EQ("{3}", null_result->ToString()); + + ASSERT_OK_AND_ASSIGN(auto range_readers, ReadColumn(reader.get(), "f1")); + ASSERT_EQ(1, range_readers.size()); + ASSERT_OK_AND_ASSIGN(auto greater_result, range_readers[0]->VisitGreaterThan(Literal(20))); + ASSERT_EQ("{2,3}", greater_result->ToString()); +} + +TEST_F(DataFileIndexWriterTest, TestExternalIndexAndAbortCleanup) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + std::string data_path = path_factory_->NewPath(); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish(data_path)); + ASSERT_FALSE(result.embedded_index); + ASSERT_EQ(1, result.extra_files.size()); + ASSERT_TRUE(result.extra_files[0]); + ASSERT_EQ(PathUtil::GetName(path_factory_->ToFileIndexPath(data_path)), + result.extra_files[0].value()); + std::string index_path = path_factory_->ToFileIndexPath(data_path); + ASSERT_OK_AND_ASSIGN(bool exists, file_system_->Exists(index_path)); + ASSERT_TRUE(exists); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(index_path)); + ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input, pool_)); + ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bitmap_readers.size()); + ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(1))); + ASSERT_EQ("{0}", equal_result->ToString()); + + writer->Abort(); + ASSERT_OK_AND_ASSIGN(exists, file_system_->Exists(index_path)); + ASSERT_FALSE(exists); +} + +TEST_F(DataFileIndexWriterTest, TestUnavailableWriterFailsCreation) { + ASSERT_NOK_WITH_MSG(CreateWriter({{"file-index.unknown.columns", "f0"}}), + "File index type 'unknown' is not registered"); + ASSERT_NOK_WITH_MSG(CreateWriter({{"file-index.bloom-filter.columns", "f0"}}), + "do not support index writer in bloom filter"); +} + +TEST_F(DataFileIndexWriterTest, TestRejectSystemFieldIndex) { + std::shared_ptr key_value_schema = + SpecialFields::CompleteSequenceAndValueKindField(schema_); + for (const std::string& field_name : + {SpecialFields::SequenceNumber().Name(), SpecialFields::ValueKind().Name()}) { + ASSERT_OK_AND_ASSIGN( + CoreOptions core_options, + CoreOptions::FromMap({{"file-index.bitmap.columns", field_name}}, file_system_)); + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, + FileIndexOptions::FromCoreOptions(core_options)); + ASSERT_NOK_WITH_MSG(DataFileIndexWriter::Create(key_value_schema, options, file_system_, + path_factory_, pool_), + "is a system field"); + } +} + +TEST_F(DataFileIndexWriterTest, TestFinishIsOneShot) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + ASSERT_OK(writer->Finish("unused.orc")); + + ASSERT_NOK_WITH_MSG(writer->Finish("unused.orc"), "already finished"); + ASSERT_NOK_WITH_MSG(writer->AddBatch(CreateBatch(R"([{"f0": 2, "f1": 20}])")), + "already finished"); +} + +TEST_F(DataFileIndexWriterTest, TestCloseFailureClosesExternalStreamOnce) { + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, + CoreOptions::FromMap({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}}, + file_system_)); + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, FileIndexOptions::FromCoreOptions(core_options)); + auto state = std::make_shared(); + auto close_failing_file_system = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(auto writer, + DataFileIndexWriter::Create(schema_, options, close_failing_file_system, + path_factory_, pool_)); + ASSERT_OK(writer->AddBatch(CreateBatch(R"([{"f0": 1, "f1": 10}])"))); + + ASSERT_NOK_WITH_MSG(writer->Finish(path_factory_->NewPath()), "close failed"); + ASSERT_EQ(1, state->close_count); + ASSERT_EQ(1, state->delete_count); +} + +} // namespace paimon::test diff --git a/src/paimon/core/io/data_file_writer.cpp b/src/paimon/core/io/data_file_writer.cpp index 4ed3e040..9275fcf2 100644 --- a/src/paimon/core/io/data_file_writer.cpp +++ b/src/paimon/core/io/data_file_writer.cpp @@ -19,6 +19,7 @@ #include "paimon/core/io/data_file_writer.h" #include +#include #include "arrow/c/abi.h" #include "paimon/common/utils/long_counter.h" @@ -36,7 +37,7 @@ DataFileWriter::DataFileWriter( const std::shared_ptr& stats_extractor, bool is_external_path, const std::optional>& write_cols, const std::shared_ptr& pool) - : SingleFileWriter(compression, converter), + : DataFileWriterBase(compression, std::move(converter)), pool_(pool), schema_id_(schema_id), is_external_path_(is_external_path), @@ -45,28 +46,13 @@ DataFileWriter::DataFileWriter( stats_extractor_(stats_extractor), write_cols_(write_cols) {} -void DataFileWriter::SetMetadataFinalizer(MetadataFinalizer finalizer) { - metadata_finalizer_ = std::move(finalizer); -} - Status DataFileWriter::Write(ArrowArray* batch) { int64_t record_count = batch->length; - PAIMON_RETURN_NOT_OK(SingleFileWriter::Write(batch)); + PAIMON_RETURN_NOT_OK(WriteRecordWithFileIndex(batch)); seq_num_counter_->Add(record_count); return Status::OK(); } -Status DataFileWriter::BeforeFinish() { - if (metadata_finalizer_) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, - metadata_finalizer_()); - if (updated_schema) { - PAIMON_RETURN_NOT_OK(UpdateSchema(updated_schema)); - } - } - return Status::OK(); -} - Result> DataFileWriter::GetResult() { PAIMON_ASSIGN_OR_RAISE(std::vector> field_stats, GetFieldStats()); PAIMON_ASSIGN_OR_RAISE(SimpleStats stats, @@ -77,11 +63,12 @@ Result> DataFileWriter::GetResult() { PAIMON_ASSIGN_OR_RAISE(Path external_path, PathUtil::ToPath(path_)); final_path = external_path.ToString(); } + const FileIndexWriteResult& file_index = GetFileIndexWriteResult(); return DataFileMeta::ForAppend( PathUtil::GetName(path_), output_bytes_, RecordCount(), stats, seq_num_counter_->GetValue() - RecordCount(), seq_num_counter_->GetValue() - 1, schema_id_, - {}, /*embedded_index=*/nullptr, file_source_, /*value_stats_cols=*/std::nullopt, final_path, - /*first_row_id=*/std::nullopt, write_cols_); + file_index.extra_files, file_index.embedded_index, file_source_, + /*value_stats_cols=*/std::nullopt, final_path, /*first_row_id=*/std::nullopt, write_cols_); } Result>> DataFileWriter::GetFieldStats() { diff --git a/src/paimon/core/io/data_file_writer.h b/src/paimon/core/io/data_file_writer.h index 60cc808a..f56f3495 100644 --- a/src/paimon/core/io/data_file_writer.h +++ b/src/paimon/core/io/data_file_writer.h @@ -28,7 +28,7 @@ #include "arrow/c/abi.h" #include "paimon/common/utils/long_counter.h" #include "paimon/core/io/data_file_meta.h" -#include "paimon/core/io/single_file_writer.h" +#include "paimon/core/io/data_file_writer_base.h" #include "paimon/core/manifest/file_source.h" #include "paimon/result.h" #include "paimon/status.h" @@ -44,13 +44,8 @@ class FormatStatsExtractor; class LongCounter; class MemoryPool; -class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr> { +class DataFileWriter : public DataFileWriterBase<::ArrowArray*> { public: - /// Callback invoked during BeforeFinish() to finalize file metadata. - /// Produces an updated schema with per-field metadata (e.g. shredding metadata) - /// and may perform other finalization work (e.g. reporting stats to cross-file context). - using MetadataFinalizer = std::function>()>; - DataFileWriter(const std::string& compression, std::function converter, int64_t schema_id, const std::shared_ptr& seq_num_counter, FileSource file_source, @@ -58,17 +53,10 @@ class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr>& write_cols, const std::shared_ptr& pool); - /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated - /// schema and perform finalization callbacks. Must be set before Close(). - void SetMetadataFinalizer(MetadataFinalizer finalizer); - Status Write(::ArrowArray* batch) override; Result> GetResult() override; - protected: - Status BeforeFinish() override; - private: Result>> GetFieldStats(); @@ -80,7 +68,6 @@ class DataFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr stats_extractor_; std::optional> write_cols_; - MetadataFinalizer metadata_finalizer_; }; } // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h new file mode 100644 index 00000000..ccea898a --- /dev/null +++ b/src/paimon/core/io/data_file_writer_base.h @@ -0,0 +1,145 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/io/data_file_index_writer.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/single_file_writer.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon { + +struct KeyValueBatch; + +/// Common lifecycle for data file writers which may finalize schema metadata and publish a +/// file-level index. Concrete writers remain responsible for their record-specific state and +/// DataFileMeta construction. +template +class DataFileWriterBase : public SingleFileWriter> { + public: + using Base = SingleFileWriter>; + using AbortExecutor = typename Base::AbortExecutor; + /// Callback invoked during BeforeFinish() to finalize file metadata. + /// Produces an updated schema with per-field metadata (e.g. shredding metadata) + /// and may perform other finalization work (e.g. reporting stats to cross-file context). + using MetadataFinalizer = std::function>()>; + + /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated + /// schema and perform finalization callbacks. Must be set before Close(). + void SetMetadataFinalizer(MetadataFinalizer finalizer) { + metadata_finalizer_ = std::move(finalizer); + } + + void SetFileIndexWriter(std::unique_ptr&& file_index_writer, + const std::shared_ptr& logical_schema) { + file_index_writer_ = std::move(file_index_writer); + logical_type_ = arrow::struct_(logical_schema->fields()); + } + + void Abort() override { + if (file_index_writer_) { + // The external index uses a path different from the data file path deleted by Base. + file_index_writer_->Abort(); + } + Base::Abort(); + } + + Result GetAbortExecutor() const override { + PAIMON_ASSIGN_OR_RAISE(AbortExecutor executor, Base::GetAbortExecutor()); + if (file_index_writer_ && file_index_writer_->ExternalIndexPath()) { + executor.Add(this->fs_, file_index_writer_->ExternalIndexPath().value()); + } + return executor; + } + + protected: + DataFileWriterBase(const std::string& compression, + std::function converter) + : Base(compression, std::move(converter)) {} + + /// Extracts the pre-conversion Arrow batch from record for file index construction, then + /// passes record to the underlying data file writer, which may convert it to a physical schema. + Status WriteRecordWithFileIndex(Record record) { + PAIMON_RETURN_NOT_OK(AddFileIndexBatch(GetFileIndexBatch(record))); + return Base::Write(std::move(record)); + } + + const FileIndexWriteResult& GetFileIndexWriteResult() const { + return file_index_result_; + } + + Status BeforeFinish() override { + if (metadata_finalizer_) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, + metadata_finalizer_()); + if (updated_schema) { + PAIMON_RETURN_NOT_OK(this->UpdateSchema(updated_schema)); + } + } + return Status::OK(); + } + + Status BeforeCompletion() override { + if (file_index_writer_) { + PAIMON_ASSIGN_OR_RAISE(file_index_result_, file_index_writer_->Finish(this->path_)); + } + return Status::OK(); + } + + private: + static ::ArrowArray* GetFileIndexBatch(Record& record) { + if constexpr (std::is_same_v) { + return record; + } else { + static_assert(std::is_same_v, + "Unsupported data file record type"); + return record.batch.get(); + } + } + + Status AddFileIndexBatch(::ArrowArray* batch) { + if (!file_index_writer_) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_array, + arrow::ImportArray(batch, logical_type_)); + std::shared_ptr logical_batch = + checked_pointer_cast(logical_array); + PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*logical_batch, batch)); + return Status::OK(); + } + + MetadataFinalizer metadata_finalizer_; + std::unique_ptr file_index_writer_; + std::shared_ptr logical_type_; + FileIndexWriteResult file_index_result_; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_factory.cpp b/src/paimon/core/io/data_file_writer_factory.cpp index b929dde8..07195ab7 100644 --- a/src/paimon/core/io/data_file_writer_factory.cpp +++ b/src/paimon/core/io/data_file_writer_factory.cpp @@ -24,6 +24,9 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/file_index_options.h" #include "paimon/format/file_format.h" #include "paimon/format/writer_builder.h" @@ -58,4 +61,16 @@ Result DataFileWriterFactory::CreateWrit return resources; } +Result> DataFileWriterFactory::CreateFileIndexWriter( + const std::shared_ptr& logical_schema, + const std::shared_ptr& path_factory) const { + PAIMON_ASSIGN_OR_RAISE(FileIndexOptions file_index_options, + FileIndexOptions::FromCoreOptions(options_)); + if (file_index_options.Empty()) { + return std::unique_ptr(); + } + return DataFileIndexWriter::Create(logical_schema, file_index_options, options_.GetFileSystem(), + path_factory, pool_); +} + } // namespace paimon diff --git a/src/paimon/core/io/data_file_writer_factory.h b/src/paimon/core/io/data_file_writer_factory.h index c727b47d..cab942f0 100644 --- a/src/paimon/core/io/data_file_writer_factory.h +++ b/src/paimon/core/io/data_file_writer_factory.h @@ -32,6 +32,8 @@ class Schema; namespace paimon { class FileFormat; +class DataFileIndexWriter; +class DataFilePathFactory; class FormatStatsExtractor; class MemoryPool; class WriterBuilder; @@ -52,6 +54,10 @@ class DataFileWriterFactory { const std::shared_ptr& file_schema, bool create_stats_extractor) const; + Result> CreateFileIndexWriter( + const std::shared_ptr& logical_schema, + const std::shared_ptr& path_factory) const; + CoreOptions options_; int64_t schema_id_; std::shared_ptr pool_; diff --git a/src/paimon/core/io/file_index_options.cpp b/src/paimon/core/io/file_index_options.cpp new file mode 100644 index 00000000..a9587363 --- /dev/null +++ b/src/paimon/core/io/file_index_options.cpp @@ -0,0 +1,109 @@ +/* + * 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 "paimon/core/io/file_index_options.h" + +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/defs.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +constexpr char kFileIndexPrefix[] = "file-index."; +constexpr char kColumnsSuffix[] = ".columns"; +constexpr size_t kFileIndexPrefixLength = sizeof(kFileIndexPrefix) - 1; +constexpr size_t kColumnsSuffixLength = sizeof(kColumnsSuffix) - 1; + +} // namespace + +Result FileIndexOptions::FromCoreOptions(const CoreOptions& options) { + FileIndexOptions result; + const std::map& raw_options = options.ToMap(); + result.in_manifest_threshold_ = options.FileIndexInManifestThreshold(); + + std::set> declared; + for (const auto& [key, value] : raw_options) { + if (!StringUtils::StartsWith(key, kFileIndexPrefix) || + !StringUtils::EndsWith(key, kColumnsSuffix)) { + continue; + } + if (key.size() < kFileIndexPrefixLength + kColumnsSuffixLength) { + return Status::Invalid(fmt::format("Invalid file index option {}", key)); + } + const size_t index_type_length = key.size() - kFileIndexPrefixLength - kColumnsSuffixLength; + const std::string index_type = key.substr(kFileIndexPrefixLength, index_type_length); + if (index_type.empty()) { + return Status::Invalid(fmt::format("Invalid file index option {}", key)); + } + // TODO(jinli.zjw): Align malformed list option parsing (for example, "f1,f2,,") with Java. + // Update this together with ConfigParser::ParseList to keep option parsing consistent. + for (std::string column_name : StringUtils::Split(value, ",", /*ignore_empty=*/false)) { + StringUtils::Trim(&column_name); + if (column_name.empty()) { + return Status::Invalid( + fmt::format("Wrong option in {}, should not have empty column", key)); + } + if (column_name.find('[') != std::string::npos && + StringUtils::EndsWith(column_name, "]")) { + return Status::NotImplemented( + "Writing file indexes for nested map columns is not supported"); + } + if (declared.emplace(column_name, index_type).second) { + result.definitions_.push_back({column_name, index_type, {}}); + } + } + } + + for (const auto& [key, value] : raw_options) { + if (!StringUtils::StartsWith(key, kFileIndexPrefix) || + StringUtils::EndsWith(key, kColumnsSuffix) || + key == Options::FILE_INDEX_IN_MANIFEST_THRESHOLD) { + continue; + } + std::vector parts = + StringUtils::Split(key.substr(kFileIndexPrefixLength), ".", /*ignore_empty=*/false); + if (parts.size() != 3) { + continue; + } + bool found = false; + for (FileIndexDefinition& definition : result.definitions_) { + if (definition.index_type == parts[0] && definition.column_name == parts[1]) { + definition.options[parts[2]] = value; + found = true; + break; + } + } + if (!found) { + return Status::Invalid( + fmt::format("Wrong file index option '{}': column '{}' is not declared in " + "'file-index.{}.columns'", + key, parts[1], parts[0])); + } + } + return result; +} + +} // namespace paimon diff --git a/src/paimon/core/io/file_index_options.h b/src/paimon/core/io/file_index_options.h new file mode 100644 index 00000000..7b7c019b --- /dev/null +++ b/src/paimon/core/io/file_index_options.h @@ -0,0 +1,63 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/result.h" + +namespace paimon { + +class CoreOptions; + +struct FileIndexDefinition { + std::string column_name; + std::string index_type; + std::map options; +}; + +/// Parsed write-side file index configuration. +class FileIndexOptions { + public: + static Result FromCoreOptions(const CoreOptions& options); + + const std::vector& Definitions() const { + return definitions_; + } + + int64_t InManifestThreshold() const { + return in_manifest_threshold_; + } + + bool Empty() const { + return definitions_.empty(); + } + + private: + FileIndexOptions() = default; + + std::vector definitions_; + int64_t in_manifest_threshold_ = 0; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/file_index_options_test.cpp b/src/paimon/core/io/file_index_options_test.cpp new file mode 100644 index 00000000..157203f9 --- /dev/null +++ b/src/paimon/core/io/file_index_options_test.cpp @@ -0,0 +1,58 @@ +/* + * 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 "paimon/core/io/file_index_options.h" + +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/core/core_options.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +Result ParseOptions(const std::map& index_options) { + std::shared_ptr file_system = std::make_shared(); + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, + CoreOptions::FromMap(index_options, file_system)); + return FileIndexOptions::FromCoreOptions(core_options); +} + +} // namespace + +TEST(FileIndexOptionsTest, TestRejectOverlappingPrefixAndSuffix) { + ASSERT_NOK_WITH_MSG(ParseOptions({{"file-index.columns", "f0"}}), + "Invalid file index option file-index.columns"); +} + +TEST(FileIndexOptionsTest, TestNestedMapColumnSyntax) { + ASSERT_OK_AND_ASSIGN(FileIndexOptions options, + ParseOptions({{"file-index.bitmap.columns", "col[key"}})); + ASSERT_EQ(1, options.Definitions().size()); + ASSERT_EQ("col[key", options.Definitions()[0].column_name); + + ASSERT_NOK_WITH_MSG(ParseOptions({{"file-index.bitmap.columns", "col[key]"}}), + "nested map columns is not supported"); +} + +} // namespace paimon::test diff --git a/src/paimon/core/io/key_value_data_file_writer.cpp b/src/paimon/core/io/key_value_data_file_writer.cpp index 9393c7c3..9c32e067 100644 --- a/src/paimon/core/io/key_value_data_file_writer.cpp +++ b/src/paimon/core/io/key_value_data_file_writer.cpp @@ -25,7 +25,6 @@ #include #include -#include "arrow/type.h" #include "fmt/format.h" #include "paimon/common/data/binary_array.h" #include "paimon/common/data/binary_array_writer.h" @@ -53,7 +52,7 @@ KeyValueDataFileWriter::KeyValueDataFileWriter( const std::shared_ptr& stats_extractor, const std::shared_ptr& write_schema, bool is_external_path, const std::shared_ptr& pool) - : SingleFileWriter(compression, converter), + : DataFileWriterBase(compression, std::move(converter)), pool_(pool), schema_id_(schema_id), level_(level), @@ -64,10 +63,6 @@ KeyValueDataFileWriter::KeyValueDataFileWriter( is_external_path_(is_external_path), disable_stats_(stats_extractor == nullptr) {} -void KeyValueDataFileWriter::SetMetadataFinalizer(MetadataFinalizer finalizer) { - metadata_finalizer_ = std::move(finalizer); -} - Status KeyValueDataFileWriter::Write(KeyValueBatch batch) { // update min and max key if (!min_key_) { @@ -80,19 +75,7 @@ Status KeyValueDataFileWriter::Write(KeyValueBatch batch) { // update delete row count delete_row_count_ += batch.delete_row_count; - PAIMON_RETURN_NOT_OK(SingleFileWriter::Write(std::move(batch))); - return Status::OK(); -} - -Status KeyValueDataFileWriter::BeforeFinish() { - if (metadata_finalizer_) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr updated_schema, - metadata_finalizer_()); - if (updated_schema) { - PAIMON_RETURN_NOT_OK(UpdateSchema(updated_schema)); - } - } - return Status::OK(); + return WriteRecordWithFileIndex(std::move(batch)); } Result> KeyValueDataFileWriter::GetResult() { @@ -120,14 +103,14 @@ Result> KeyValueDataFileWriter::GetResult() { final_path = external_path.ToString(); } PAIMON_ASSIGN_OR_RAISE(int64_t local_micro, DateTimeUtils::GetCurrentLocalTimeUs()); + const FileIndexWriteResult& file_index = GetFileIndexWriteResult(); return std::make_shared( PathUtil::GetName(path_), output_bytes_, RecordCount(), min_key, max_key, key_stats, value_stats, min_sequence_number_, max_sequence_number_, schema_id_, level_, - /*extra_files=*/std::vector>(), + file_index.extra_files, Timestamp(/*millisecond=*/local_micro / 1000, /*nano_of_millisecond=*/0), delete_row_count_, - /*embedded_index=*/nullptr, file_source_, - /*value_stats_cols=*/std::nullopt, final_path, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + file_index.embedded_index, file_source_, /*value_stats_cols=*/std::nullopt, final_path, + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); } Status KeyValueDataFileWriter::GenerateMinMaxKey(BinaryRow* min_key, BinaryRow* max_key) const { diff --git a/src/paimon/core/io/key_value_data_file_writer.h b/src/paimon/core/io/key_value_data_file_writer.h index e1e3fd92..eb7a2efc 100644 --- a/src/paimon/core/io/key_value_data_file_writer.h +++ b/src/paimon/core/io/key_value_data_file_writer.h @@ -17,6 +17,7 @@ */ #pragma once + #include #include #include @@ -25,7 +26,7 @@ #include #include "paimon/core/io/data_file_meta.h" -#include "paimon/core/io/single_file_writer.h" +#include "paimon/core/io/data_file_writer_base.h" #include "paimon/core/key_value.h" #include "paimon/core/manifest/file_source.h" #include "paimon/result.h" @@ -44,14 +45,8 @@ class InternalRow; class MemoryPool; class SimpleStats; -class KeyValueDataFileWriter - : public SingleFileWriter> { +class KeyValueDataFileWriter : public DataFileWriterBase { public: - /// Callback invoked during BeforeFinish() to finalize file metadata. - /// Produces an updated schema with per-field metadata (e.g. shredding metadata) - /// and may perform other finalization work (e.g. reporting stats to cross-file context). - using MetadataFinalizer = std::function>()>; - KeyValueDataFileWriter(const std::string& compression, std::function converter, int64_t schema_id, int32_t level, FileSource file_source, @@ -60,17 +55,10 @@ class KeyValueDataFileWriter const std::shared_ptr& write_schema, bool is_external_path, const std::shared_ptr& pool); - /// Sets the metadata finalizer. Called during BeforeFinish() to produce an updated - /// schema and perform finalization callbacks. Must be set before Close(). - void SetMetadataFinalizer(MetadataFinalizer finalizer); - Status Write(KeyValueBatch batch) override; Result> GetResult() override; - protected: - Status BeforeFinish() override; - private: Result>> GetFieldStats(); @@ -96,7 +84,6 @@ class KeyValueDataFileWriter int64_t max_sequence_number_ = std::numeric_limits::min(); std::shared_ptr min_key_; std::shared_ptr max_key_; - MetadataFinalizer metadata_finalizer_; }; } // namespace paimon diff --git a/src/paimon/core/io/key_value_data_file_writer_factory.cpp b/src/paimon/core/io/key_value_data_file_writer_factory.cpp index 07d50b98..8f388559 100644 --- a/src/paimon/core/io/key_value_data_file_writer_factory.cpp +++ b/src/paimon/core/io/key_value_data_file_writer_factory.cpp @@ -24,6 +24,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/key_value_data_file_writer.h" #include "paimon/format/file_format.h" @@ -60,6 +61,11 @@ KeyValueDataFileWriterFactory::CreateWriter() const { options_.GetWriteFileCompression(level_), std::move(converter), schema_id_, level_, file_source_, primary_keys_, resources.stats_extractor, write_schema_, path_factory_->IsExternalPath(), pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); return std::unique_ptr>>( diff --git a/src/paimon/core/io/merged_key_value_record_reader.cpp b/src/paimon/core/io/merged_key_value_record_reader.cpp index 70f2bcfb..8c395287 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader.cpp @@ -117,13 +117,21 @@ Status MergedKeyValueRecordReader::Iterator::LoadNextRawKeyValue() const { } Result> MergedKeyValueRecordReader::NextBatch() { + if (initialization_error_.has_value()) { + return initialization_error_.value(); + } if (visited_) { return std::unique_ptr(); } visited_ = true; auto iterator = std::make_unique(this); - PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); + Result has_next_result = iterator->HasNext(); + if (!has_next_result.ok()) { + initialization_error_ = has_next_result.status(); + return initialization_error_.value(); + } + bool has_next = std::move(has_next_result).value(); if (!has_next) { return std::unique_ptr(); } diff --git a/src/paimon/core/io/merged_key_value_record_reader.h b/src/paimon/core/io/merged_key_value_record_reader.h index a1b7aa5e..227a1593 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.h +++ b/src/paimon/core/io/merged_key_value_record_reader.h @@ -67,6 +67,7 @@ class MergedKeyValueRecordReader : public KeyValueRecordReader { private: bool visited_ = false; + std::optional initialization_error_; std::unique_ptr reader_; std::shared_ptr key_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 1b6b71c6..3ba81f03 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -18,8 +18,12 @@ #include "paimon/core/io/merged_key_value_record_reader.h" +#include #include +#include +#include #include +#include #include "arrow/api.h" #include "arrow/array/array_nested.h" @@ -27,10 +31,13 @@ #include "gtest/gtest.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/fields_comparator.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include "paimon/memory/memory_pool.h" +#include "paimon/realtime/offset_range.h" #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/key_value_checker.h" @@ -38,6 +45,61 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { + +namespace { + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))); +} + +std::shared_ptr MakePreparedSchema(const arrow::FieldVector& value_fields) { + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + prepared_fields.insert(prepared_fields.end(), value_fields.begin(), value_fields.end()); + return arrow::schema(prepared_fields); +} + +Result> AdaptPreparedBatchReaderForTest( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + return AdaptPreparedBatchReader(std::move(reader), prepared_schema, visible_offsets, key_schema, + value_schema, memory_pool); +} + +class TrackingBatchReader : public BatchReader { + public: + TrackingBatchReader(std::unique_ptr&& delegate, int32_t* close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + ++(*close_count_); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + int32_t* close_count_; +}; + +} // namespace + class MergedKeyValueRecordReaderTest : public testing::Test { public: void SetUp() override { @@ -143,4 +205,367 @@ TEST_F(MergedKeyValueRecordReaderTest, TestSkipMergedNulloptResultInHasNext) { } } +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderOffsetFilter) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 100, 0, 1, 10], + [0, 101, 1, 2, 20], + [0, 102, 4, 3, 30], + [0, 103, 2, 4, 40], + [0, 104, 5, 5, 50], + [0, 105, 3, 6, 60] + ])") + .ValueOrDie()); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 2); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), + key_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + std::vector row_kinds = {const_cast(RowKind::Insert()), + const_cast(RowKind::Insert())}; + std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; + std::vector expected = KeyValueChecker::GenerateKeyValues( + row_kinds, {103, 105}, levels, {{4}, {6}}, {{4, 40}, {6, 60}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderRejectsReversedVisibleOffsets) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + std::shared_ptr prepared_array = + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + auto batch_reader = + std::make_unique(prepared_array, prepared_type, /*batch_size=*/1); + + Result> result = + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(2, 1), + value_schema, value_schema, pool_); + ASSERT_TRUE(result.status().IsInvalid()); + ASSERT_NOK_WITH_MSG(result, "prepared visible offset range begin exceeds end"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr extra = MakeField("extra", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, extra}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1, 2]])") + .ValueOrDie()); + + auto query_batch_reader = + std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr query_reader, + AdaptPreparedBatchReaderForTest(std::move(query_batch_reader), prepared_schema, + OffsetRange(0, 1), value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector query_results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(query_reader.get()))); + ASSERT_EQ(query_results.size(), 1); + ASSERT_EQ(query_results[0].value->GetFieldCount(), 1); + ASSERT_EQ(query_results[0].value->GetInt(0), 1); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); + ASSERT_NOK_WITH_MSG( + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, std::nullopt, + value_schema, value_schema, pool_), + "exact"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestBadCommitBatch) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value = MakeField("value", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, std::nullopt, + arrow::schema({key}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "field count"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMissingCompositeKey) { + std::shared_ptr key0 = MakeField("key0", arrow::int32(), 0); + std::shared_ptr key1 = MakeField("key1", arrow::int32(), 1); + std::shared_ptr value = MakeField("value", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key0, key1, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key0, key1, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key0, value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key0, key1}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "cannot find field id 1"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestQueryAddRename) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr old_value = MakeField("old_value", arrow::int32(), 1); + std::shared_ptr renamed_value = MakeField("renamed_value", arrow::int32(), 1); + std::shared_ptr added = MakeField("added", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key, renamed_value, added}); + std::shared_ptr prepared_schema = + MakePreparedSchema({key, renamed_value, added}); + std::shared_ptr actual_schema = MakePreparedSchema({key, old_value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key}), value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr iterator, + reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(KeyValue key_value, iterator->Next()); + ASSERT_EQ(20, key_value.value->GetInt(1)); + ASSERT_TRUE(key_value.value->IsNullAt(2)); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMergedReaderErrorRetry) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + std::shared_ptr prepared_array = + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + auto failing_reader = std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("stable prepared error")); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(failing_reader), prepared_schema, + OffsetRange(0, 1), value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({DataField(0, key)}, true)); + MergedKeyValueRecordReader merged_reader(std::move(reader), key_comparator, + merge_function_wrapper_); + + Result> first = merged_reader.NextBatch(); + Result> retry = merged_reader.NextBatch(); + ASSERT_NOK(first); + ASSERT_NOK(retry); + ASSERT_EQ(first.status().ToString(), retry.status().ToString()); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderSafeDecode) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + + arrow::FieldVector invalid_fields = prepared_schema->fields(); + invalid_fields[0] = invalid_fields[0]->WithName("wrong_value_kind"); + invalid_fields[3] = MakeField("wrong_key", arrow::int32(), 99); + std::shared_ptr invalid_type = arrow::struct_(invalid_fields); + auto invalid_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[0, 10, 0, 1]])").ValueOrDie()); + + auto batch_reader = std::make_unique(invalid_array, invalid_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "prepared batch field"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { + std::shared_ptr id = MakeField("id", arrow::int32(), 0); + std::shared_ptr item_a = MakeField("a", arrow::int32(), 10); + std::shared_ptr item_b = MakeField("b", arrow::int32(), 11); + std::shared_ptr items = + MakeField("items", arrow::list(arrow::field("item", arrow::struct_({item_a, item_b}))), 2); + std::shared_ptr attr_x = MakeField("x", arrow::int32(), 20); + std::shared_ptr attr_y = MakeField("y", arrow::int32(), 21); + std::shared_ptr attrs = + MakeField("attrs", arrow::map(arrow::utf8(), arrow::struct_({attr_x, attr_y})), 3); + std::shared_ptr key_left = MakeField("left", arrow::int32(), 30); + std::shared_ptr key_right = MakeField("right", arrow::int32(), 31); + std::shared_ptr keyed_values = MakeField( + "keyed_values", arrow::map(arrow::struct_({key_left, key_right}), arrow::int32()), 4); + std::shared_ptr full_value_schema = + arrow::schema({id, items, attrs, keyed_values}); + std::shared_ptr key_schema = arrow::schema({id}); + std::shared_ptr prepared_schema = + MakePreparedSchema(full_value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 9, 9, 0, [[1, 2]], [["prefix", [3, 4]]], [[[5, 6], 7]]], + [0, 10, 0, 1, [[100, 200], [300, 400]], [["k1", [7, 8]], ["k2", [9, 10]]], [[[11, 12], 13], [[21, 22], 23]]], + [0, 11, 11, 2, [[8, 9]], [["suffix", [10, 11]]], [[[12, 13], 14]]] + ])") + .ValueOrDie()); + prepared_array = checked_pointer_cast(prepared_array->Slice(1, 1)); + + std::shared_ptr query_item_b = MakeField("renamed_b", arrow::int32(), 11); + std::shared_ptr query_item_a = MakeField("renamed_a", arrow::int32(), 10); + std::shared_ptr query_items = MakeField( + "items_renamed", + arrow::list(arrow::field("element", arrow::struct_({query_item_b, query_item_a}))), 2); + std::shared_ptr query_attr_y = MakeField("renamed_y", arrow::int32(), 21); + std::shared_ptr query_attr_x = MakeField("renamed_x", arrow::int32(), 20); + std::shared_ptr query_attrs = + MakeField("attrs_renamed", + arrow::map(arrow::utf8(), arrow::struct_({query_attr_y, query_attr_x})), 3); + std::shared_ptr query_key_right = MakeField("renamed_right", arrow::int32(), 31); + std::shared_ptr query_key_left = MakeField("renamed_left", arrow::int32(), 30); + std::shared_ptr query_keyed_values = + MakeField("keyed_values_renamed", + arrow::map(arrow::struct_({query_key_right, query_key_left}), arrow::int32()), 4); + std::shared_ptr query_value_schema = + arrow::schema({id, query_items, query_attrs, query_keyed_values}); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + key_schema, query_value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0].key->GetInt(0), 1); + ASSERT_EQ(results[0].value->GetFieldCount(), 4); + ASSERT_EQ(results[0].value->GetInt(0), 1); + + std::shared_ptr item_array = results[0].value->GetArray(1); + ASSERT_EQ(item_array->Size(), 2); + std::shared_ptr first_item = item_array->GetRow(0, 2); + ASSERT_EQ(first_item->GetInt(0), 200); + ASSERT_EQ(first_item->GetInt(1), 100); + std::shared_ptr second_item = item_array->GetRow(1, 2); + ASSERT_EQ(second_item->GetInt(0), 400); + ASSERT_EQ(second_item->GetInt(1), 300); + + std::shared_ptr attr_map = results[0].value->GetMap(2); + ASSERT_EQ(attr_map->Size(), 2); + std::shared_ptr key_array = attr_map->KeyArray(); + ASSERT_EQ(std::string(key_array->GetStringView(0)), "k1"); + ASSERT_EQ(std::string(key_array->GetStringView(1)), "k2"); + std::shared_ptr value_array = attr_map->ValueArray(); + std::shared_ptr first_attr = value_array->GetRow(0, 2); + ASSERT_EQ(first_attr->GetInt(0), 8); + ASSERT_EQ(first_attr->GetInt(1), 7); + std::shared_ptr second_attr = value_array->GetRow(1, 2); + ASSERT_EQ(second_attr->GetInt(0), 10); + ASSERT_EQ(second_attr->GetInt(1), 9); + + std::shared_ptr keyed_value_map = results[0].value->GetMap(3); + ASSERT_EQ(keyed_value_map->Size(), 2); + std::shared_ptr struct_keys = keyed_value_map->KeyArray(); + std::shared_ptr first_key = struct_keys->GetRow(0, 2); + ASSERT_EQ(first_key->GetInt(0), 12); + ASSERT_EQ(first_key->GetInt(1), 11); + std::shared_ptr second_key = struct_keys->GetRow(1, 2); + ASSERT_EQ(second_key->GetInt(0), 22); + ASSERT_EQ(second_key->GetInt(1), 21); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(0), 13); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(1), 23); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderLifecycle) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1, 100] + ])") + .ValueOrDie()); + + int32_t explicit_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &explicit_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(tracking_reader), prepared_schema, + OffsetRange(0, 1), key_schema, value_schema, pool_)); + reader->Close(); + } + ASSERT_EQ(explicit_close_count, 1); + + int32_t destructor_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &destructor_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(tracking_reader), prepared_schema, + OffsetRange(0, 1), key_schema, value_schema, pool_)); + } + ASSERT_EQ(destructor_close_count, 1); + + int32_t factory_failure_close_count = 0; + { + std::unique_ptr tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &factory_failure_close_count); + std::shared_ptr invalid_schema = arrow::schema(value_schema->fields()); + ASSERT_NOK(AdaptPreparedBatchReaderForTest(std::move(tracking_reader), invalid_schema, + OffsetRange(0, 1), key_schema, value_schema, + pool_)); + ASSERT_EQ(nullptr, tracking_reader); + } + ASSERT_EQ(factory_failure_close_count, 1); + + int32_t read_failure_close_count = 0; + { + auto failing_reader = + std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("prepared reader failure")); + auto tracking_reader = std::make_unique(std::move(failing_reader), + &read_failure_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReaderForTest(std::move(tracking_reader), prepared_schema, + OffsetRange(0, 1), key_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "prepared reader failure"); + ASSERT_EQ(read_failure_close_count, 1); + } + ASSERT_EQ(read_failure_close_count, 1); +} + } // namespace paimon::test diff --git a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp index 6e4843bb..0e4e8219 100644 --- a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp @@ -23,6 +23,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_file_writer.h" #include "paimon/core/io/infer_shredding_file_writer.h" @@ -89,6 +90,11 @@ ShreddingAppendDataFileWriterFactory::CreateShreddedWriter( options_.GetFileCompression(), std::move(batch_converter), schema_id_, seq_num_counter, file_source_, resources.stats_extractor, path_factory_->IsExternalPath(), write_cols_, pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer = diff --git a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp index 8ac583ee..30d4c9fc 100644 --- a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp @@ -23,6 +23,7 @@ #include "arrow/c/helpers.h" #include "paimon/core/core_options.h" +#include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/infer_shredding_file_writer.h" #include "paimon/core/io/key_value_data_file_writer.h" @@ -88,6 +89,11 @@ ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter( options_.GetWriteFileCompression(level_), std::move(batch_converter), schema_id_, level_, file_source_, primary_keys_, resources.stats_extractor, file_schema, path_factory_->IsExternalPath(), pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_index_writer, + CreateFileIndexWriter(write_schema_, path_factory_)); + if (file_index_writer) { + writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer = diff --git a/src/paimon/core/io/single_file_writer.h b/src/paimon/core/io/single_file_writer.h index 99507b57..6db3a699 100644 --- a/src/paimon/core/io/single_file_writer.h +++ b/src/paimon/core/io/single_file_writer.h @@ -25,6 +25,7 @@ #include #include #include +#include #include "arrow/c/abi.h" #include "arrow/c/bridge.h" @@ -64,21 +65,27 @@ class SingleFileWriter : public FileWriter { class AbortExecutor { public: AbortExecutor(const std::shared_ptr& fs, const std::string& path) - : fs_(fs), path_(path), logger_(Logger::GetLogger("AbortExecutor")) {} + : paths_({{fs, path}}), logger_(Logger::GetLogger("AbortExecutor")) {} + + void Add(const std::shared_ptr& fs, const std::string& path) { + paths_.emplace_back(fs, path); + } void Abort() { - if (fs_) { - auto status = fs_->Delete(path_); + for (const auto& [fs, path] : paths_) { + if (!fs) { + continue; + } + auto status = fs->Delete(path); if (!status.ok()) { - PAIMON_LOG_WARN(logger_, "Exception occurs when deleting %s: %s", path_.c_str(), + PAIMON_LOG_WARN(logger_, "Exception occurs when deleting %s: %s", path.c_str(), status.ToString().c_str()); } } } private: - std::shared_ptr fs_; - std::string path_; + std::vector, std::string>> paths_; std::shared_ptr logger_; }; @@ -132,6 +139,11 @@ class SingleFileWriter : public FileWriter { return Status::OK(); } + /// Hook called after the data file is closed and before its completion callback is published. + virtual Status BeforeCompletion() { + return Status::OK(); + } + /// Serializes schema and forwards it as file metadata to FormatWriter. Status UpdateSchema(const std::shared_ptr& schema); @@ -239,6 +251,7 @@ Status SingleFileWriter::Close() { // guard still removes the file on a callback error, while a repeated Close() does not publish // the same file again. closed_ = true; + PAIMON_RETURN_NOT_OK(BeforeCompletion()); if (completion_callback_) { PAIMON_RETURN_NOT_OK(completion_callback_()); } diff --git a/src/paimon/core/io/single_file_writer_test.cpp b/src/paimon/core/io/single_file_writer_test.cpp index 78fce54c..4136702e 100644 --- a/src/paimon/core/io/single_file_writer_test.cpp +++ b/src/paimon/core/io/single_file_writer_test.cpp @@ -18,8 +18,10 @@ #include "paimon/core/io/single_file_writer.h" +#include #include -#include +#include +#include #include "arrow/api.h" #include "arrow/c/abi.h" diff --git a/src/paimon/core/io/vector_file_batch_reader.cpp b/src/paimon/core/io/vector_file_batch_reader.cpp new file mode 100644 index 00000000..f16025c4 --- /dev/null +++ b/src/paimon/core/io/vector_file_batch_reader.cpp @@ -0,0 +1,285 @@ +/* + * 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 "paimon/core/io/vector_file_batch_reader.h" + +#include +#include +#include +#include + +#include "arrow/array.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/arrow/vector_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +std::shared_ptr FindField(const std::shared_ptr& type, + const std::string& name) { + for (const auto& field : type->fields()) { + if (field->name() == name) { + return field; + } + } + return nullptr; +} + +/// Rebuilds `map_type` with new key and item types, keeping the name and metadata of its +/// entries field. +std::shared_ptr MakeMapType(const arrow::MapType& map_type, + const std::shared_ptr& key_field, + const std::shared_ptr& item_field) { + return std::make_shared( + map_type.value_field()->WithType(arrow::struct_({key_field, item_field})), + map_type.keys_sorted()); +} + +/// Returns the type to request from the file format plugin. A VECTOR is only read back as a +/// LIST when the file itself stores it as one: writers such as Paimon Java expose VECTOR +/// columns as Arrow LIST, while Paimon Rust and Python expose them as FixedSizeList. +std::shared_ptr GetPhysicalReadType( + const std::shared_ptr& logical_type, + const std::shared_ptr& file_type) { + switch (logical_type->id()) { + case arrow::Type::FIXED_SIZE_LIST: { + if (!file_type || file_type->id() != arrow::Type::LIST) { + return logical_type; + } + const auto& vector_type = checked_cast(*logical_type); + const auto& list_type = checked_cast(*file_type); + return arrow::list(vector_type.value_field()->WithType( + GetPhysicalReadType(vector_type.value_type(), list_type.value_type()))); + } + case arrow::Type::STRUCT: { + if (!file_type || file_type->id() != arrow::Type::STRUCT) { + return logical_type; + } + arrow::FieldVector fields; + fields.reserve(logical_type->num_fields()); + for (const auto& field : logical_type->fields()) { + std::shared_ptr file_field = FindField(file_type, field->name()); + fields.push_back(field->WithType( + GetPhysicalReadType(field->type(), file_field ? file_field->type() : nullptr))); + } + return arrow::struct_(fields); + } + case arrow::Type::LIST: { + if (!file_type || file_type->id() != arrow::Type::LIST) { + return logical_type; + } + return arrow::list(logical_type->field(0)->WithType( + GetPhysicalReadType(logical_type->field(0)->type(), file_type->field(0)->type()))); + } + case arrow::Type::MAP: { + if (!file_type || file_type->id() != arrow::Type::MAP) { + return logical_type; + } + const auto& map_type = checked_cast(*logical_type); + const auto& file_map_type = checked_cast(*file_type); + return MakeMapType(map_type, + map_type.key_field()->WithType(GetPhysicalReadType( + map_type.key_type(), file_map_type.key_type())), + map_type.item_field()->WithType(GetPhysicalReadType( + map_type.item_type(), file_map_type.item_type()))); + } + default: + return logical_type; + } +} + +Result> CastListToVector( + const std::shared_ptr& array, + const std::shared_ptr& read_type, arrow::MemoryPool* pool) { + if (array->type_id() != arrow::Type::LIST) { + return Status::Invalid( + fmt::format("Cannot restore VECTOR from type {}", array->type()->ToString())); + } + PAIMON_RETURN_NOT_OK(VectorUtils::ValidateVectorElements(*array)); + if (array->null_count() == array->length()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::MakeArrayOfNull(read_type, array->length(), pool)); + return result; + } + arrow::compute::ExecContext exec_context(pool); + arrow::TypeHolder type_holder(read_type.get()); + arrow::compute::CastOptions options = arrow::compute::CastOptions::Safe(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr result, + arrow::compute::Cast(*array, type_holder, options, &exec_context)); + return result; +} + +std::shared_ptr RebuildNestedType( + const std::shared_ptr& read_type, + const std::vector>& children) { + if (read_type->id() == arrow::Type::STRUCT) { + arrow::FieldVector fields; + fields.reserve(children.size()); + for (int32_t i = 0; i < static_cast(children.size()); ++i) { + fields.push_back(read_type->field(i)->WithType(children[i]->type)); + } + return arrow::struct_(fields); + } + if (read_type->id() == arrow::Type::LIST) { + return arrow::list(read_type->field(0)->WithType(children[0]->type)); + } + + const auto& entries_type = checked_cast(*children[0]->type); + const auto& map_type = checked_cast(*read_type); + return MakeMapType(map_type, map_type.key_field()->WithType(entries_type.field(0)->type()), + map_type.item_field()->WithType(entries_type.field(1)->type())); +} + +Result> ConvertToReadType( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* pool) { + if (!VectorUtils::ContainsVectorType(read_type)) { + return array; + } + switch (read_type->id()) { + case arrow::Type::FIXED_SIZE_LIST: { + if (array->type_id() == arrow::Type::FIXED_SIZE_LIST) { + const auto& source_type = + checked_cast(*array->type()); + const auto& vector_type = checked_cast(*read_type); + if (source_type.list_size() != vector_type.list_size() || + !source_type.value_type()->Equals(vector_type.value_type())) { + return Status::Invalid(fmt::format("VECTOR type mismatch: data {} vs read {}", + source_type.ToString(), + vector_type.ToString())); + } + PAIMON_RETURN_NOT_OK(VectorUtils::ValidateVectorElements(*array)); + // Writers disagree on the element field, for example `element: float not null` + // for Paimon Rust against the `item: float` of a Paimon schema. Restore the + // requested type so that files storing VECTOR as LIST and files storing it as + // FixedSizeList produce batches of one type. + std::shared_ptr data = array->data()->Copy(); + data->type = read_type; + return arrow::MakeArray(data); + } + return CastListToVector( + array, checked_pointer_cast(read_type), pool); + } + case arrow::Type::STRUCT: + case arrow::Type::LIST: + case arrow::Type::MAP: { + if (array->type_id() != read_type->id()) { + return Status::Invalid(fmt::format("Cannot reconcile file type {} with {}", + array->type()->ToString(), + read_type->ToString())); + } + if (array->type()->num_fields() != read_type->num_fields() || + array->data()->child_data.size() != static_cast(read_type->num_fields())) { + return Status::Invalid( + fmt::format("Cannot reconcile file type {} with {}: nested field count differs", + array->type()->ToString(), read_type->ToString())); + } + std::vector> children; + children.reserve(read_type->num_fields()); + for (int32_t i = 0; i < read_type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr child, + ConvertToReadType(arrow::MakeArray(array->data()->child_data[i]), + read_type->field(i)->type(), pool)); + children.push_back(child->data()); + } + std::shared_ptr data = array->data()->Copy(); + data->child_data = std::move(children); + data->type = RebuildNestedType(read_type, data->child_data); + return arrow::MakeArray(data); + } + default: + return array; + } +} + +} // namespace + +VectorFileBatchReader::VectorFileBatchReader(std::unique_ptr&& reader, + const std::shared_ptr& pool) + : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)) {} + +bool VectorFileBatchReader::ContainsVector(const std::shared_ptr& schema) { + return VectorUtils::ContainsVector(schema); +} + +Status VectorFileBatchReader::SetReadSchema( + ::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + if (!read_schema) { + return Status::Invalid("SetReadSchema failed: read schema cannot be nullptr"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_schema, + arrow::ImportSchema(read_schema)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> c_file_schema, reader_->GetFileSchema()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_schema, + arrow::ImportSchema(c_file_schema.get())); + arrow::FieldVector physical_fields; + physical_fields.reserve(logical_schema->num_fields()); + for (const auto& field : logical_schema->fields()) { + std::shared_ptr file_field = file_schema->GetFieldByName(field->name()); + physical_fields.push_back(field->WithType( + GetPhysicalReadType(field->type(), file_field ? file_field->type() : nullptr))); + } + std::shared_ptr physical_schema = + arrow::schema(physical_fields, logical_schema->metadata()); + ArrowSchema c_physical_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*physical_schema, &c_physical_schema)); + PAIMON_RETURN_NOT_OK(reader_->SetReadSchema(&c_physical_schema, predicate, selection_bitmap)); + read_type_ = arrow::struct_(logical_schema->fields()); + return Status::OK(); +} + +Result VectorFileBatchReader::ConvertBatch(ReadBatch&& batch) const { + if (BatchReader::IsEofBatch(batch) || !read_type_) { + return std::move(batch); + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + PAIMON_ASSIGN_OR_RAISE(array, ConvertToReadType(array, read_type_, arrow_pool_.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(array->Validate()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); + return std::move(batch); +} + +Result VectorFileBatchReader::NextBatch() { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, reader_->NextBatch()); + return ConvertBatch(std::move(batch)); +} + +Result VectorFileBatchReader::NextBatchWithBitmap() { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch_with_bitmap, reader_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + return std::move(batch_with_bitmap); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, ConvertBatch(std::move(batch_with_bitmap.first))); + batch_with_bitmap.first = std::move(batch); + return std::move(batch_with_bitmap); +} + +} // namespace paimon diff --git a/src/paimon/core/io/vector_file_batch_reader.h b/src/paimon/core/io/vector_file_batch_reader.h new file mode 100644 index 00000000..b7eab263 --- /dev/null +++ b/src/paimon/core/io/vector_file_batch_reader.h @@ -0,0 +1,85 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +#include "arrow/c/abi.h" +#include "paimon/reader/file_batch_reader.h" + +namespace arrow { +class DataType; +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { +class MemoryPool; + +/// Reconciles logical VECTOR values with the variable-length LIST representation exposed to file +/// format plugins. +class VectorFileBatchReader : public FileBatchReader { + public: + VectorFileBatchReader(std::unique_ptr&& reader, + const std::shared_ptr& pool); + + static bool ContainsVector(const std::shared_ptr& schema); + + Result> GetFileSchema() const override { + return reader_->GetFileSchema(); + } + + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + + Result NextBatch() override; + + Result NextBatchWithBitmap() override; + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override { + return reader_->GetPreviousBatchFileRowId(batch_row_id); + } + + Result GetNumberOfRows() const override { + return reader_->GetNumberOfRows(); + } + + bool SupportPreciseBitmapSelection() const override { + return reader_->SupportPreciseBitmapSelection(); + } + + private: + Result ConvertBatch(ReadBatch&& batch) const; + + std::shared_ptr arrow_pool_; + std::shared_ptr read_type_; + std::unique_ptr reader_; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/vector_file_batch_reader_test.cpp b/src/paimon/core/io/vector_file_batch_reader_test.cpp new file mode 100644 index 00000000..e334e62a --- /dev/null +++ b/src/paimon/core/io/vector_file_batch_reader_test.cpp @@ -0,0 +1,232 @@ +/* + * 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 "paimon/core/io/vector_file_batch_reader.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr AsStructType(const std::shared_ptr& type) { + return checked_pointer_cast(type); +} + +} // namespace + +TEST(VectorFileBatchReaderTest, ConvertSchemaAndNextBatch) { + auto physical_type = AsStructType(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::list(arrow::float32())), + })); + auto logical_type = AsStructType(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)), + })); + const std::string json = R"([ + [1, [1.0, 2.0, 3.0]], + [2, null], + [3, [4.0, 5.0, 6.0]] + ])"; + auto physical_array = + arrow::ipc::internal::json::ArrayFromJSON(physical_type, json).ValueOrDie(); + auto mock_reader = + std::make_unique(physical_array, physical_type, /*batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + MockFileBatchReader* inner_reader = mock_reader.get(); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + + ASSERT_TRUE(VectorFileBatchReader::ContainsVector(arrow::schema(logical_type->fields()))); + ASSERT_FALSE(VectorFileBatchReader::ContainsVector(arrow::schema(physical_type->fields()))); + ArrowSchema c_read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_EQ(inner_reader->read_schema_->field(1)->type()->id(), arrow::Type::LIST); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader.NextBatch()); + arrow::Result> actual_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + std::shared_ptr actual = std::move(actual_result).ValueOrDie(); + auto expected = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json).ValueOrDie(); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); + ASSERT_OK_AND_ASSIGN(batch, reader.NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(VectorFileBatchReaderTest, KeepFixedSizeListFileSchema) { + auto logical_type = AsStructType(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)), + })); + const std::string json = R"([ + [1, [1.0, 2.0, 3.0]], + [2, null], + [3, [4.0, 5.0, 6.0]] + ])"; + auto logical_array = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json).ValueOrDie(); + auto mock_reader = + std::make_unique(logical_array, logical_type, /*batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + MockFileBatchReader* inner_reader = mock_reader.get(); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + + ArrowSchema c_read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_EQ(inner_reader->read_schema_->field(1)->type()->id(), arrow::Type::FIXED_SIZE_LIST); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader.NextBatch()); + arrow::Result> actual_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + ASSERT_TRUE(logical_array->Equals(std::move(actual_result).ValueOrDie())); +} + +// Paimon Rust names the element field of a VECTOR column `element` and marks it non-nullable, +// while a Paimon schema names it `item`. Batches must carry the requested type either way, +// otherwise they cannot be combined with batches read from a file storing VECTOR as LIST. +TEST(VectorFileBatchReaderTest, NormalizeFixedSizeListElementField) { + auto file_vector = + arrow::fixed_size_list(arrow::field("element", arrow::float32(), /*nullable=*/false), 3); + auto logical_vector = arrow::fixed_size_list(arrow::float32(), 3); + auto file_type = AsStructType(arrow::struct_({ + arrow::field("embedding", file_vector), + arrow::field("history", arrow::list(file_vector)), + })); + auto logical_type = AsStructType(arrow::struct_({ + arrow::field("embedding", logical_vector), + arrow::field("history", arrow::list(logical_vector)), + })); + const std::string json = R"([ + [[1.0, 2.0, 3.0], [[4.0, 5.0, 6.0]]], + [null, []] + ])"; + auto file_array = arrow::ipc::internal::json::ArrayFromJSON(file_type, json).ValueOrDie(); + auto mock_reader = + std::make_unique(file_array, file_type, /*batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + + ArrowSchema c_read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader.NextBatch()); + arrow::Result> actual_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + std::shared_ptr actual = std::move(actual_result).ValueOrDie(); + ASSERT_TRUE(actual->type()->Equals(logical_type)) << actual->type()->ToString(); + auto expected = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json).ValueOrDie(); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); +} + +TEST(VectorFileBatchReaderTest, ConvertNestedVectorsWithBitmap) { + auto logical_vector = + arrow::fixed_size_list(arrow::field("item", arrow::float64(), /*nullable=*/false), 2); + auto physical_vector = arrow::list(arrow::field("item", arrow::float64(), /*nullable=*/false)); + auto logical_type = AsStructType(arrow::struct_({ + arrow::field("vectors", arrow::list(logical_vector)), + arrow::field("by_name", arrow::map(arrow::utf8(), logical_vector)), + })); + auto physical_type = AsStructType(arrow::struct_({ + arrow::field("vectors", arrow::list(physical_vector)), + arrow::field("by_name", arrow::map(arrow::utf8(), physical_vector)), + })); + const std::string json = R"([[[[1.0, 2.0], null], [["a", [3.0, 4.0]]]], + [null, [["b", null]]]])"; + auto physical_array = + arrow::ipc::internal::json::ArrayFromJSON(physical_type, json).ValueOrDie(); + RoaringBitmap32 bitmap; + bitmap.Add(1); + auto mock_reader = std::make_unique(physical_array, physical_type, bitmap, + /*read_batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + ArrowSchema c_read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader.NextBatchWithBitmap()); + ASSERT_FALSE(batch_with_bitmap.second.Contains(0)); + ASSERT_TRUE(batch_with_bitmap.second.Contains(1)); + arrow::Result> actual_result = arrow::ImportArray( + batch_with_bitmap.first.first.get(), batch_with_bitmap.first.second.get()); + ASSERT_TRUE(actual_result.ok()) << actual_result.status().ToString(); + std::shared_ptr actual = std::move(actual_result).ValueOrDie(); + auto expected = arrow::ipc::internal::json::ArrayFromJSON(logical_type, json).ValueOrDie(); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); +} + +TEST(VectorFileBatchReaderTest, RejectInvalidVectorValues) { + auto physical_type = + AsStructType(arrow::struct_({arrow::field("embedding", arrow::list(arrow::float32()))})); + auto logical_type = AsStructType( + arrow::struct_({arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))})); + for (const char* json : {R"([[[1.0, 2.0]]])", R"([[[1.0, null, 3.0]]])"}) { + auto physical_array = + arrow::ipc::internal::json::ArrayFromJSON(physical_type, json).ValueOrDie(); + auto mock_reader = std::make_unique(physical_array, physical_type, + /*read_batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + ArrowSchema c_read_schema; + ASSERT_TRUE( + arrow::ExportSchema(*arrow::schema(logical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_NOK(reader.NextBatch()); + } +} + +TEST(VectorFileBatchReaderTest, RejectInvalidFixedSizeListVectorValues) { + auto physical_type = AsStructType( + arrow::struct_({arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))})); + auto physical_array = + arrow::ipc::internal::json::ArrayFromJSON(physical_type, R"([[[1.0, null, 3.0]]])") + .ValueOrDie(); + auto mock_reader = std::make_unique(physical_array, physical_type, + /*read_batch_size=*/10); + mock_reader->EnableRandomizeBatchSize(false); + VectorFileBatchReader reader(std::move(mock_reader), GetDefaultPool()); + ArrowSchema c_read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(physical_type->fields()), &c_read_schema).ok()); + ASSERT_OK(reader.SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_NOK(reader.NextBatch()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/manifest/manifest_entry_serializer.cpp b/src/paimon/core/manifest/manifest_entry_serializer.cpp index 053405b8..2389cd8d 100644 --- a/src/paimon/core/manifest/manifest_entry_serializer.cpp +++ b/src/paimon/core/manifest/manifest_entry_serializer.cpp @@ -31,17 +31,26 @@ namespace paimon { class MemoryPool; struct DataFileMeta; +Status ManifestEntrySerializer::ValidateVersion(int32_t version) { + if (version == VERSION_2) { + return Status::OK(); + } + if (version == VERSION_1) { + return Status::Invalid( + fmt::format("The current version {} is not compatible with the version {}, " + "please recreate the table.", + VERSION_2, version)); + } + return Status::Invalid(fmt::format("Unsupported version: {}", version)); +} + +int32_t ManifestEntrySerializer::GetBucket(const InternalRow& row) { + return row.GetInt(3); +} + Result ManifestEntrySerializer::ConvertFrom(int32_t version, const InternalRow& row) const { - if (version != VERSION_2) { - if (version == VERSION_1) { - return Status::Invalid( - fmt::format("The current version {} is not compatible with the version {}, " - "please recreate the table.", - GetVersion(), version)); - } - return Status::Invalid("Unsupported version", std::to_string(version)); - } + PAIMON_RETURN_NOT_OK(ValidateVersion(version)); auto kind = row.GetByte(0); PAIMON_ASSIGN_OR_RAISE(FileKind file_kind, FileKind::FromByteValue(kind)); auto partition_bytes = row.GetBinary(1); diff --git a/src/paimon/core/manifest/manifest_entry_serializer.h b/src/paimon/core/manifest/manifest_entry_serializer.h index 7438895f..4a71a1b6 100644 --- a/src/paimon/core/manifest/manifest_entry_serializer.h +++ b/src/paimon/core/manifest/manifest_entry_serializer.h @@ -50,6 +50,12 @@ class ManifestEntrySerializer : public VersionedObjectSerializer return VERSION_2; } + /// Validate the serialization version before reading fields that may vary by version. + static Status ValidateVersion(int32_t version); + + /// Get the bucket from a versioned manifest entry row without fully deserializing it. + static int32_t GetBucket(const InternalRow& row); + Result ConvertFrom(int32_t version, const InternalRow& row) const override; Result ToRow(const ManifestEntry& record) const override; diff --git a/src/paimon/core/manifest/manifest_entry_serializer_test.cpp b/src/paimon/core/manifest/manifest_entry_serializer_test.cpp index 2aa2db52..2d8cffc3 100644 --- a/src/paimon/core/manifest/manifest_entry_serializer_test.cpp +++ b/src/paimon/core/manifest/manifest_entry_serializer_test.cpp @@ -55,12 +55,22 @@ TEST_F(ManifestEntrySerializerTest, TestToFromRow) { ManifestEntrySerializer serializer(pool); for (const auto& entry : entries) { ASSERT_OK_AND_ASSIGN(auto row, serializer.ToRow(entry)); + ASSERT_EQ(entry.Bucket(), ManifestEntrySerializer::GetBucket(row)); ASSERT_OK_AND_ASSIGN(auto result_entry, serializer.FromRow(row)); ASSERT_EQ(entry, result_entry); ASSERT_EQ(entry.ToString(), result_entry.ToString()); } } +TEST_F(ManifestEntrySerializerTest, TestValidateVersion) { + ASSERT_OK(ManifestEntrySerializer::ValidateVersion(/*version=*/2)); + ASSERT_NOK_WITH_MSG(ManifestEntrySerializer::ValidateVersion(/*version=*/1), + "The current version 2 is not compatible with the version 1, please " + "recreate the table."); + ASSERT_NOK_WITH_MSG(ManifestEntrySerializer::ValidateVersion(/*version=*/3), + "Unsupported version: 3"); +} + TEST_F(ManifestEntrySerializerTest, TestNullableRecordCount) { std::vector empty_entries; ASSERT_FALSE(ManifestEntry::NullableRecordCount(empty_entries).has_value()); diff --git a/src/paimon/core/manifest/manifest_file.cpp b/src/paimon/core/manifest/manifest_file.cpp index 22f2681f..9f9c8aee 100644 --- a/src/paimon/core/manifest/manifest_file.cpp +++ b/src/paimon/core/manifest/manifest_file.cpp @@ -24,6 +24,7 @@ #include "arrow/c/abi.h" #include "arrow/c/bridge.h" +#include "paimon/common/data/columnar/columnar_row.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/io/rolling_file_writer.h" #include "paimon/core/manifest/manifest_entry.h" @@ -86,6 +87,25 @@ Result> ManifestFile::Create( manifest_file_factory, target_file_size, pool, options, partition_type)); } +Status ManifestFile::ReadBucketEntries(const std::string& file_name, int32_t bucket, + std::vector* entries) const { + return ReadArrowBatches( + file_name, + [this, bucket, entries](const std::shared_ptr& batch) -> Status { + const arrow::ArrayVector& fields = batch->fields(); + for (int64_t i = 0; i < batch->length(); i++) { + ColumnarRow row(fields, pool_, i); + PAIMON_RETURN_NOT_OK(ManifestEntrySerializer::ValidateVersion(row.GetInt(0))); + if (ManifestEntrySerializer::GetBucket(row) != bucket) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(ManifestEntry entry, serializer_->FromRow(row)); + entries->push_back(std::move(entry)); + } + return Status::OK(); + }); +} + Result> ManifestFile::Write( const std::vector& entries) { if (entries.empty()) { diff --git a/src/paimon/core/manifest/manifest_file.h b/src/paimon/core/manifest/manifest_file.h index d34764b5..0211e14d 100644 --- a/src/paimon/core/manifest/manifest_file.h +++ b/src/paimon/core/manifest/manifest_file.h @@ -62,6 +62,10 @@ class ManifestFile : public ObjectsFile { /// @note This method is atomic. Result> Write(const std::vector& entries); + /// Read a manifest file and deserialize only entries for the specified bucket. + Status ReadBucketEntries(const std::string& file_name, int32_t bucket, + std::vector* entries) const; + private: ManifestFile(const std::shared_ptr& file_system, const std::shared_ptr& reader_builder, diff --git a/src/paimon/core/manifest/manifest_file_test.cpp b/src/paimon/core/manifest/manifest_file_test.cpp index 8f6e0b2e..a34f4152 100644 --- a/src/paimon/core/manifest/manifest_file_test.cpp +++ b/src/paimon/core/manifest/manifest_file_test.cpp @@ -23,7 +23,6 @@ #include #include #include -#include #include "arrow/api.h" #include "gtest/gtest.h" @@ -100,10 +99,10 @@ class CountingFileSystem : public FileSystem { class ManifestFileTest : public testing::Test { public: - std::vector ReadManifestEntry(const std::string& file_format_str, - const std::string& root_path, - const std::string& file_name, - const std::shared_ptr& pool) const { + std::vector ReadManifestEntry( + const std::string& file_format_str, const std::string& root_path, + const std::string& file_name, const std::shared_ptr& pool, + const std::optional& bucket = std::nullopt) const { std::shared_ptr file_system = std::make_shared(); EXPECT_OK_AND_ASSIGN(std::shared_ptr file_format, FileFormatFactory::Get(file_format_str, {})); @@ -124,7 +123,12 @@ class ManifestFileTest : public testing::Test { ManifestFile::Create(file_system, file_format, "zstd", path_factory, /*target_file_size=*/1024, pool, options, unused_schema)); std::vector manifest_entries; - EXPECT_OK(manifest_file->Read(file_name, /*filter=*/nullptr, &manifest_entries)); + if (bucket) { + EXPECT_OK( + manifest_file->ReadBucketEntries(file_name, bucket.value(), &manifest_entries)); + } else { + EXPECT_OK(manifest_file->Read(file_name, /*filter=*/nullptr, &manifest_entries)); + } return manifest_entries; } @@ -316,6 +320,104 @@ TEST_F(ManifestFileTest, TestManifestCacheReusesCachedBytes) { ASSERT_EQ(1, manifest_cache->Size()); } +TEST_F(ManifestFileTest, TestReadBucketEntriesMaterializesOnlySelectedBucket) { + auto pool = GetDefaultPool(); + auto counting_file_system = std::make_shared(); + auto manifest_cache = + std::make_shared(CacheKind::MANIFEST, 64 * 1024 * 1024); + ASSERT_OK_AND_ASSIGN(std::shared_ptr file_format, + FileFormatFactory::Get("orc", {})); + std::string root_path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; + auto unused_schema = arrow::schema(arrow::FieldVector({arrow::field("f0", arrow::utf8())})); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr path_factory, + FileStorePathFactory::Create(root_path, unused_schema, /*partition_keys=*/{}, + /*default_part_value=*/"", file_format->Identifier(), + /*data_file_prefix=*/"data-", + /*legacy_partition_name_enabled=*/true, /*external_paths=*/{}, + /*global_index_external_path=*/std::nullopt, + /*index_file_in_data_file_dir=*/false, pool)); + ASSERT_OK_AND_ASSIGN( + CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}, {Options::MANIFEST_FORMAT, "orc"}})); + options.WithCache(manifest_cache); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr manifest_file, + ManifestFile::Create(counting_file_system, file_format, "zstd", path_factory, + /*target_file_size=*/1024, pool, options, unused_schema)); + + const std::string manifest_name = "manifest-3a44a0da-1008-463c-914e-28d271375e24-0"; + std::vector all_entries; + ASSERT_OK(manifest_file->Read(manifest_name, /*filter=*/nullptr, &all_entries)); + ASSERT_EQ(2, all_entries.size()); + + std::vector bucket_one_entries; + ASSERT_OK(manifest_file->ReadBucketEntries(manifest_name, /*bucket=*/1, &bucket_one_entries)); + ASSERT_EQ(std::vector({all_entries[0]}), bucket_one_entries); + + std::vector bucket_zero_entries; + ASSERT_OK(manifest_file->ReadBucketEntries(manifest_name, /*bucket=*/0, &bucket_zero_entries)); + ASSERT_EQ(std::vector({all_entries[1]}), bucket_zero_entries); + + std::vector missing_bucket_entries; + ASSERT_OK( + manifest_file->ReadBucketEntries(manifest_name, /*bucket=*/2, &missing_bucket_entries)); + ASSERT_TRUE(missing_bucket_entries.empty()); + + ASSERT_EQ(1, counting_file_system->open_count); + ASSERT_EQ(4, manifest_cache->GetCount()); + ASSERT_EQ(1, manifest_cache->SupplierCallCount()); +} + +TEST_F(ManifestFileTest, TestReadBucketEntriesSkipsDeserializingOtherBuckets) { + auto pool = GetDefaultPool(); + std::vector source_entries = + ReadManifestEntry("orc", paimon::test::GetDataDir() + "/orc/append_09.db/append_09", + "manifest-3a44a0da-1008-463c-914e-28d271375e24-0", pool); + ASSERT_EQ(2, source_entries.size()); + + auto test_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(test_dir); + std::shared_ptr file_system = test_dir->GetFileSystem(); + ASSERT_OK(file_system->Mkdirs(FileStorePathFactory::ManifestPath(test_dir->Str()))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr file_format, + FileFormatFactory::Get("orc", {})); + auto unused_schema = arrow::schema(arrow::FieldVector({arrow::field("f0", arrow::utf8())})); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr path_factory, + FileStorePathFactory::Create(test_dir->Str(), unused_schema, /*partition_keys=*/{}, + /*default_part_value=*/"", file_format->Identifier(), + /*data_file_prefix=*/"data-", + /*legacy_partition_name_enabled=*/true, /*external_paths=*/{}, + /*global_index_external_path=*/std::nullopt, + /*index_file_in_data_file_dir=*/false, pool)); + ASSERT_OK_AND_ASSIGN( + CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}, {Options::MANIFEST_FORMAT, "orc"}})); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr manifest_file, + ManifestFile::Create(file_system, file_format, "zstd", path_factory, + /*target_file_size=*/1024, pool, options, unused_schema)); + + ManifestEntry invalid_other_bucket(FileKind(static_cast(2)), + source_entries[0].Partition(), /*bucket=*/1, + /*total_buckets=*/2, source_entries[0].File()); + ManifestEntry valid_target_bucket(FileKind::Add(), source_entries[1].Partition(), /*bucket=*/0, + /*total_buckets=*/2, source_entries[1].File()); + using WrittenFile = std::pair; + ASSERT_OK_AND_ASSIGN( + WrittenFile written_file, + manifest_file->WriteWithoutRolling({invalid_other_bucket, valid_target_bucket})); + + std::vector all_entries; + ASSERT_NOK_WITH_MSG(manifest_file->Read(written_file.first, /*filter=*/nullptr, &all_entries), + "Unsupported byte value 2 for file kind."); + + std::vector bucket_entries; + ASSERT_OK(manifest_file->ReadBucketEntries(written_file.first, /*bucket=*/0, &bucket_entries)); + ASSERT_EQ(std::vector({valid_target_bucket}), bucket_entries); +} + TEST_F(ManifestFileTest, TestWithNullCount) { auto pool = GetDefaultPool(); auto manifest_entries = @@ -406,6 +508,9 @@ TEST_F(ManifestFileTest, TestManifestFileCompatibleWithJavaPaimon09) { std::vector expected_manifest_entries; expected_manifest_entries.emplace_back(manifest_entry); ASSERT_EQ(expected_manifest_entries, manifest_entries); + ASSERT_EQ(expected_manifest_entries, + ReadManifestEntry("avro", paimon::test::GetDataDir() + "/avro", "avro_manifest_09", + pool, /*bucket=*/0)); } TEST_F(ManifestFileTest, TestManifestFileCompatibleWithJavaPaimon11) { @@ -442,6 +547,9 @@ TEST_F(ManifestFileTest, TestManifestFileCompatibleWithJavaPaimon11) { std::vector expected_manifest_entries; expected_manifest_entries.emplace_back(manifest_entry); ASSERT_EQ(expected_manifest_entries, manifest_entries); + ASSERT_EQ(expected_manifest_entries, + ReadManifestEntry("avro", paimon::test::GetDataDir() + "/avro", "avro_manifest_11", + pool, /*bucket=*/0)); } } // namespace paimon::test diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 3b6806c7..49961536 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -154,6 +154,59 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReaders( + std::vector>&& readers) { + auto raw_readers_guard = ScopeGuard([&]() -> void { + for (std::unique_ptr& reader : readers) { + if (reader != nullptr) { + reader->Close(); + } + } + }); + if (readers.empty()) { + return Status::Invalid("sorted readers must not be empty"); + } + for (const std::unique_ptr& reader : readers) { + if (reader == nullptr) { + return Status::Invalid("sorted readers must not contain null reader"); + } + } + + auto sort_merge_reader = std::make_unique( + std::move(readers), key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_); + auto create_consumer = [target_schema = write_schema_, pool = pool_]() + -> Result>> { + return KeyValueMetaProjectionConsumer::Create(target_schema, pool); + }; + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), 1, pool_); + raw_readers_guard.Release(); + ScopeGuard async_readers_guard([&]() -> void { async_key_value_producer_consumer->Close(); }); + std::unique_ptr>> rolling_writer; + PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); + ScopeGuard abort_writer_guard([&]() -> void { rolling_writer->Abort(); }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_key_value_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + } + PAIMON_RETURN_NOT_OK(rolling_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, + rolling_writer->GetResult()); + abort_writer_guard.Release(); + + for (const std::shared_ptr& flushed_file : flushed_files) { + new_files_.emplace_back(flushed_file); + PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); + } + metrics_->Merge(rolling_writer->GetMetrics()); + return Status::OK(); +} + Status MergeTreeWriter::Compact(bool full_compaction) { return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction); } @@ -256,49 +309,9 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, wait_for_latest_compaction = true; } auto cleanup_guard = ScopeGuard([&]() { write_buffer_->Clear(); }); - // 1. flush write buffer to get sorted readers PAIMON_ASSIGN_OR_RAISE(std::vector> readers, write_buffer_->CreateReaders()); - // 2. prepare loser tree sort merge reader - auto sort_merge_reader = std::make_unique( - std::move(readers), key_comparator_, user_defined_seq_comparator_, - merge_function_wrapper_); - // 3. project key value to arrow array - auto create_consumer = [target_schema = write_schema_, pool = pool_]() - -> Result>> { - return KeyValueMetaProjectionConsumer::Create(target_schema, pool); - }; - // consumer batch size is WriteBatchSize - auto async_key_value_producer_consumer = - std::make_unique>( - std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), - /*projection_thread_num=*/1, pool_); - std::unique_ptr>> - rolling_writer; - PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); - ScopeGuard write_guard([&]() -> void { - rolling_writer->Abort(); - async_key_value_producer_consumer->Close(); - }); - while (true) { - PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, - async_key_value_producer_consumer->NextBatch()); - if (key_value_batch.batch == nullptr) { - break; - } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); - } - PAIMON_RETURN_NOT_OK(rolling_writer->Close()); - PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, - rolling_writer->GetResult()); - async_key_value_producer_consumer->Close(); - write_guard.Release(); - - for (const auto& flushed_file : flushed_files) { - new_files_.emplace_back(flushed_file); - PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); - } - metrics_->Merge(rolling_writer->GetMetrics()); + PAIMON_RETURN_NOT_OK(WriteSortedReaders(std::move(readers))); } PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction)); PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction)); diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index febce2af..01efd975 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -51,6 +51,7 @@ class IOManager; class FieldsComparator; class MemoryPool; class Metrics; +class KeyValueRecordReader; template class MergeFunctionWrapper; @@ -69,6 +70,10 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; + /// Consumes readers whose complete streams are individually sorted by primary key and sequence + /// number. Readers are closed on success or failure. + Status WriteSortedReaders(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 2155647a..9ce5498c 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -42,6 +43,7 @@ #include "paimon/core/io/compact_increment.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_increment.h" +#include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -52,6 +54,8 @@ #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/read_result_collector.h" @@ -64,6 +68,36 @@ class MergeFunctionWrapper; } // namespace paimon namespace paimon::test { +namespace { + +class TrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + TrackingKeyValueRecordReader(std::unique_ptr&& inner_reader, + bool* closed_flag) + : inner_reader_(std::move(inner_reader)), closed_flag_(closed_flag) {} + + Result> NextBatch() override { + return inner_reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return inner_reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_flag_ != nullptr) { + *closed_flag_ = true; + } + inner_reader_->Close(); + } + + private: + std::unique_ptr inner_reader_; + bool* closed_flag_; +}; + +} // namespace + class MergeTreeWriterTest : public ::testing::TestWithParam { public: class FakeCompactManager : public paimon::CompactManager { @@ -211,6 +245,23 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { writer_compact_manager, io_manager, /*enable_multi_thread_spill=*/false, pool_); } + std::unique_ptr CreateSingleReader( + const std::shared_ptr& array, int32_t batch_size = 16, + const Status& next_batch_status = Status::OK()) const { + std::vector write_fields = {SpecialFields::SequenceNumber(), + SpecialFields::ValueKind()}; + write_fields.insert(write_fields.end(), value_fields_.begin(), value_fields_.end()); + std::shared_ptr write_schema = + DataField::ConvertDataFieldsToArrowSchema(write_fields); + std::shared_ptr key_schema = + arrow::schema(arrow::FieldVector({write_schema->field(2)})); + auto file_batch_reader = + std::make_unique(array, array->type(), batch_size); + file_batch_reader->SetNextBatchStatus(next_batch_status); + return std::make_unique( + std::move(file_batch_reader), key_schema, value_schema_, 0, pool_); + } + private: std::shared_ptr pool_; std::shared_ptr file_system_; @@ -293,6 +344,29 @@ TEST_P(MergeTreeWriterTest, TestSimple) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [2, 0, "Alice", 10, 0, 13.1], + [0, 0, "Lucy", 20, 1, 14.1], + [1, 0, "Paul", 20, 1, null] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(-1, dir->Str(), sorted_reader_path_factory, 1, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); } TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { @@ -374,6 +448,210 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [16, 0, "Alice", 10, 0, 113.1], + [14, 0, "Lucy", 20, 1, 114.1], + [13, 0, "Paul", 20, 1, 15.1], + [15, 0, "Skye", 10, 0, 118.1] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(9, dir->Str(), sorted_reader_path_factory, 0, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + std::string uuid = path_factory->uuid_; + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(4, dir->Str(), path_factory, 7, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])") + .ValueOrDie()); + + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + std::string expected_data_file_name = "data-" + uuid + "-0.orc"; + std::string expected_data_file_path = dir->Str() + "/" + expected_data_file_name; + ASSERT_OK_AND_ASSIGN(FileStatus data_file_status, + options.GetFileSystem()->GetFileStatus(expected_data_file_path)); + + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])"}, + &expected_array) + .ok()); + CheckFileContent(expected_data_file_path, expected_array); + + ASSERT_TRUE(commit_increment.GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(expected_data_file_name, new_file->file_name); + ASSERT_EQ(data_file_status.GetLen(), new_file->file_size); + ASSERT_EQ(3, new_file->row_count); + ASSERT_EQ(7, new_file->min_sequence_number); + ASSERT_EQ(9, new_file->max_sequence_number); + ASSERT_EQ(7, new_file->schema_id); + ASSERT_EQ(1, new_file->delete_row_count); +} + +TEST_P(MergeTreeWriterTest, TestMergeSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/4, dir->Str(), path_factory, + /*schema_id=*/7, options)); + + auto first_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [5, 0, "Alice", 10, 0, 15.1], + [7, 0, "Carol", 20, 1, 17.1], + [10, 0, "Eve", 30, 2, 20.1] + ])") + .ValueOrDie()); + auto second_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1] + ])") + .ValueOrDie()); + bool first_closed = false; + bool second_closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(first_array), &first_closed)); + sorted_readers.push_back(std::make_unique( + CreateSingleReader(second_array), &second_closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(first_closed); + ASSERT_TRUE(second_closed); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(5, new_file->row_count); + ASSERT_EQ(1, new_file->delete_row_count); + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [5, 0, "Alice", 10, 0, 15.1], + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1], + [10, 0, "Eve", 30, 2, 20.1] + ])"}, + &expected_array) + .ok()); + CheckFileContent(path_factory->ToPath(new_file), expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderOwnership) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + + bool closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(closed); + ASSERT_OK(merge_writer->Close()); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + std::vector> empty_readers; + Status empty_status = merge_writer->WriteSortedReaders(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReaders(std::move(null_readers)); + ASSERT_TRUE(null_status.IsInvalid()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + Status expected_status = Status::IOError("sorted reader failure"); + bool failing_reader_closed = false; + std::vector> failing_readers; + failing_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array, /*batch_size=*/16, expected_status), + &failing_reader_closed)); + Status failing_status = merge_writer->WriteSortedReaders(std::move(failing_readers)); + ASSERT_EQ(expected_status, failing_status); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); } TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 2a3d9e10..bb82c5d8 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -43,6 +43,7 @@ #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/field_mapping_reader.h" +#include "paimon/core/io/vector_file_batch_reader.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/partition/partition_info.h" #include "paimon/core/schema/table_schema.h" @@ -139,6 +140,13 @@ Result> AbstractSplitRead::PrepareReaderBuilder( file_format->CreateReaderBuilder(options_.GetReadBatchSize())); reader_builder->WithMemoryPool(pool_); reader_builder->WithCache(options_.GetCache()); + // Propagate the framework runtime read state so each format can adapt its own + // behavior (e.g. parquet disabling its pre-buffer when the shared read-ahead cache + // takes over prefetching), instead of mutating format options here. + ReadHints read_hints; + read_hints.prefetch_enabled = context_->EnablePrefetch(); + read_hints.read_ahead_cache_enabled = context_->ReadAheadCacheEnabled(); + reader_builder->WithReadHints(read_hints); return reader_builder; } @@ -154,7 +162,7 @@ Result> AbstractSplitRead::CreateFileBatchReade context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, - /*initialize_read_ranges=*/false, context_->GetPrefetchCacheMode(), + /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), context_->GetCacheConfig(), pool_)); return std::make_unique(std::move(prefetch_reader)); } else { @@ -208,6 +216,9 @@ Result> AbstractSplitRead::CreateFieldMappingRe PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_reader, CreateFileBatchReader(file_format_identifier, data_file_path, file_meta->file_size, reader_builder)); + if (VectorFileBatchReader::ContainsVector(read_schema)) { + file_reader = std::make_unique(std::move(file_reader), pool_); + } std::set skip_map_selected_keys_filter_field_ids; if (file_format_identifier != "blob") { std::pair, std::set> shared_shredding_result; diff --git a/src/paimon/core/operation/append_only_file_store_scan_test.cpp b/src/paimon/core/operation/append_only_file_store_scan_test.cpp index f319498a..e1fb5a43 100644 --- a/src/paimon/core/operation/append_only_file_store_scan_test.cpp +++ b/src/paimon/core/operation/append_only_file_store_scan_test.cpp @@ -29,6 +29,7 @@ #include "paimon/common/data/binary_row.h" #include "paimon/common/data/binary_row_writer.h" #include "paimon/common/io/cache/lru_cache.h" +#include "paimon/core/manifest/manifest_entry.h" #include "paimon/core/manifest/partition_entry.h" #include "paimon/core/operation/metrics/scan_metrics.h" #include "paimon/core/schema/schema_manager.h" @@ -186,11 +187,14 @@ namespace { std::shared_ptr BuildScan(const std::string& table_path, const std::shared_ptr& cache, const std::optional& bucket = std::nullopt, - const std::shared_ptr& predicate = nullptr) { + const std::shared_ptr& predicate = nullptr, + bool manifest_entry_lazy_decode_enabled = true) { ScanContextBuilder context_builder(table_path); context_builder.AddOption(Options::FILE_FORMAT, "orc") .AddOption(Options::MANIFEST_FORMAT, "orc") .AddOption(Options::SCAN_MANIFEST_ENTRY_CACHE_MAX_SNAPSHOTS, "8") + .AddOption(Options::SCAN_MANIFEST_ENTRY_LAZY_DECODE_ENABLED, + manifest_entry_lazy_decode_enabled ? "true" : "false") .WithCache(cache); if (bucket) { context_builder.SetBucketFilter(bucket.value()); @@ -205,6 +209,16 @@ std::shared_ptr BuildScan(const std::string& table_path, return typed_table_scan->snapshot_reader_->scan_; } +std::vector SortedFileNames(std::vector&& entries) { + std::vector file_names; + file_names.reserve(entries.size()); + for (const auto& entry : entries) { + file_names.push_back(entry.FileName()); + } + std::sort(file_names.begin(), file_names.end()); + return file_names; +} + } // namespace TEST(AppendOnlyFileStoreScanTest, TestDropStatsAfterFiltering) { @@ -253,13 +267,13 @@ TEST(AppendOnlyFileStoreScanTest, TestSnapshotLiveManifestCachePath) { scan_first->GetSnapshotManager()->LoadSnapshot(/*snapshot_id=*/5)); scan_first->WithSnapshot(snapshot_5); ASSERT_OK_AND_ASSIGN(auto plan_first, scan_first->CreatePlan()); - size_t first_size = plan_first->Files().size(); + std::vector first_file_names = SortedFileNames(plan_first->Files()); // Second scan on the same snapshot should read the same bucket live entries from cache. auto scan_second = BuildScan(table_path, cache, /*bucket=*/0); scan_second->WithSnapshot(snapshot_5); ASSERT_OK_AND_ASSIGN(auto plan_second, scan_second->CreatePlan()); - ASSERT_EQ(first_size, plan_second->Files().size()); + ASSERT_EQ(first_file_names, SortedFileNames(plan_second->Files())); } TEST(AppendOnlyFileStoreScanTest, TestSnapshotLiveManifestCacheRebuildOnMiss) { @@ -285,6 +299,24 @@ TEST(AppendOnlyFileStoreScanTest, TestSnapshotLiveManifestCacheRebuildOnMiss) { auto scan_expected = BuildScan(table_path, /*cache=*/nullptr, /*bucket=*/0); scan_expected->WithSnapshot(snapshot_5); ASSERT_OK_AND_ASSIGN(auto plan_expected, scan_expected->CreatePlan()); - ASSERT_EQ(plan_expected->Files().size(), plan_next->Files().size()); + ASSERT_EQ(SortedFileNames(plan_expected->Files()), SortedFileNames(plan_next->Files())); +} + +TEST(AppendOnlyFileStoreScanTest, TestSnapshotLiveManifestCacheFallbackWithoutLazyDecode) { + TimezoneGuard guard("Asia/Shanghai"); + std::string table_path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09/"; + auto cache = std::make_shared(/*max_weight=*/16 * 1024 * 1024); + + auto scan_fallback = BuildScan(table_path, cache, /*bucket=*/0, /*predicate=*/nullptr, + /*manifest_entry_lazy_decode_enabled=*/false); + ASSERT_OK_AND_ASSIGN(Snapshot snapshot_5, + scan_fallback->GetSnapshotManager()->LoadSnapshot(/*snapshot_id=*/5)); + scan_fallback->WithSnapshot(snapshot_5); + ASSERT_OK_AND_ASSIGN(auto plan_fallback, scan_fallback->CreatePlan()); + + auto scan_expected = BuildScan(table_path, /*cache=*/nullptr, /*bucket=*/0); + scan_expected->WithSnapshot(snapshot_5); + ASSERT_OK_AND_ASSIGN(auto plan_expected, scan_expected->CreatePlan()); + ASSERT_EQ(SortedFileNames(plan_expected->Files()), SortedFileNames(plan_fallback->Files())); } } // namespace paimon::test diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index b2946e83..5d6c2c93 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -255,9 +255,9 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( partition_values.end()); auto c_write_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*write_schema_, c_write_schema.get())); - return RealtimeAppendOnlyWriter::Create(partition_map, bucket, std::move(c_write_schema), - realtime_context_, writer, write_schema_, - options_.ToMap(), pool_); + return RealtimeAppendOnlyWriter::Create( + partition_map, bucket, std::move(c_write_schema), realtime_context_, writer, write_schema_, + options_.GetRealtimeStoreStatisticsMode(), options_.ToMap(), pool_); } Result AppendOnlyFileStoreWrite::GetDataFileWriterFactory( diff --git a/src/paimon/core/operation/append_only_file_store_write_test.cpp b/src/paimon/core/operation/append_only_file_store_write_test.cpp index 26411c19..4e1cef50 100644 --- a/src/paimon/core/operation/append_only_file_store_write_test.cpp +++ b/src/paimon/core/operation/append_only_file_store_write_test.cpp @@ -273,6 +273,7 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestRealtimeWriteTracksInternalOffsetRange) {"write-only", "true"}, {"bucket", "1"}, {"bucket-key", "id"}, + {Options::REALTIME_ENABLED, "true"}, }; auto logical_schema = arrow::schema({arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())}); diff --git a/src/paimon/core/operation/commit/commit_scanner.cpp b/src/paimon/core/operation/commit/commit_scanner.cpp index 5590dd02..448bd7e4 100644 --- a/src/paimon/core/operation/commit/commit_scanner.cpp +++ b/src/paimon/core/operation/commit/commit_scanner.cpp @@ -38,6 +38,7 @@ #include "paimon/core/operation/commit/overwrite_changes_provider.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/table/bucket_mode.h" +#include "paimon/core/utils/partition_utils.h" #include "paimon/scan_context.h" namespace paimon { @@ -183,14 +184,9 @@ Result> CommitScanner::ReadAllIndexEntriesFromPa } for (const auto& partition_spec : partitions) { - bool matched = true; - for (const auto& [key, value] : partition_spec) { - auto iter = partition.find(key); - if (iter == partition.end() || iter->second != value) { - matched = false; - break; - } - } + PAIMON_ASSIGN_OR_RAISE(bool matched, + PartitionUtils::MatchPartitionSpec(partition, partition_spec, + *partition_computer_)); if (matched) { return true; } diff --git a/src/paimon/core/operation/commit/realtime_commit_properties.cpp b/src/paimon/core/operation/commit/realtime_commit_properties.cpp index 262fbbef..c79fc968 100644 --- a/src/paimon/core/operation/commit/realtime_commit_properties.cpp +++ b/src/paimon/core/operation/commit/realtime_commit_properties.cpp @@ -31,6 +31,7 @@ #include "paimon/common/utils/rapidjson_util.h" #include "paimon/common/utils/uuid.h" #include "paimon/core/utils/branch_manager.h" +#include "paimon/core/utils/partition_utils.h" #include "paimon/fs/file_system.h" #include "paimon/macros.h" @@ -159,24 +160,71 @@ std::string RealtimeCommitProperties::OffsetsDirectory(const std::string& table_ return PathUtil::JoinPath(BranchManager::BranchPath(table_root, branch), "metadata"); } +std::optional RealtimeCommitProperties::GetOffsetsPath(const Snapshot& snapshot) { + if (!snapshot.Properties()) { + return std::nullopt; + } + const std::map& properties = snapshot.Properties().value(); + auto iter = properties.find(kOffsetsKey); + if (iter == properties.end()) { + return std::nullopt; + } + return iter->second; +} + Result RealtimeCommitProperties::ReadOffsets( const std::optional& snapshot, const std::shared_ptr& file_system) { - if (!snapshot || !snapshot->Properties()) { + if (!snapshot) { return RealtimeOffsetMap{}; } - const std::map& properties = snapshot->Properties().value(); - auto iter = properties.find(kOffsetsKey); - if (iter == properties.end()) { + std::optional offsets_path = GetOffsetsPath(snapshot.value()); + if (!offsets_path) { return RealtimeOffsetMap{}; } if (file_system == nullptr) { return Status::Invalid("file system is null when reading real-time offsets"); } std::string content; - PAIMON_RETURN_NOT_OK(file_system->ReadFile(iter->second, &content)); + PAIMON_RETURN_NOT_OK(file_system->ReadFile(offsets_path.value(), &content)); return ParseOffsets(content); } +Result RealtimeCommitProperties::AreRangesCommitted( + const RealtimeOffsetMap& committed_offsets, + const std::map& realtime_ranges) { + std::optional all_committed; + for (const auto& [partition_bucket, offset_range] : realtime_ranges) { + if (partition_bucket.bucket < 0) { + return Status::Invalid( + fmt::format("real-time commit bucket {} is invalid", partition_bucket.bucket)); + } + if (offset_range.begin < 0 || offset_range.begin >= offset_range.end) { + return Status::Invalid("real-time commit offset range is invalid"); + } + + auto offset_iter = committed_offsets.find(partition_bucket); + int64_t committed_end_offset = + offset_iter == committed_offsets.end() ? 0 : offset_iter->second; + bool range_committed = offset_range.end <= committed_end_offset; + if (!range_committed && offset_range.begin < committed_end_offset) { + return Status::Invalid(fmt::format( + "real-time commit offset range partially overlaps committed offset for bucket {}", + partition_bucket.bucket)); + } + if (!range_committed && offset_range.begin != committed_end_offset) { + return Status::Invalid( + fmt::format("real-time commit offsets for bucket {} are not contiguous", + partition_bucket.bucket)); + } + if (all_committed && all_committed.value() != range_committed) { + return Status::Invalid( + "real-time commit ranges are only partially covered by committed offsets"); + } + all_committed = range_committed; + } + return all_committed.value_or(false); +} + Result RealtimeCommitProperties::SerializeOffsets(const RealtimeOffsetMap& offsets) { std::string result; PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(OffsetsJson(offsets), &result)); @@ -187,11 +235,17 @@ Result> RealtimeCommitProperties::Build( const std::map& properties, const std::optional& latest_snapshot, const std::map& realtime_ranges, + bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + const BinaryRowPartitionComputer& partition_computer, const std::shared_ptr& file_system, const std::string& table_root, const std::string& branch) { std::map merged_properties = properties; - if (realtime_ranges.empty()) { - if (latest_snapshot && latest_snapshot->Properties()) { + if (reset_all_realtime_progress || !removed_realtime_partitions.empty()) { + merged_properties.erase(kOffsetsKey); + } + if (realtime_ranges.empty() && removed_realtime_partitions.empty()) { + if (!reset_all_realtime_progress && latest_snapshot && latest_snapshot->Properties()) { const std::map& latest_properties = latest_snapshot->Properties().value(); auto offsets_iter = latest_properties.find(kOffsetsKey); @@ -202,8 +256,25 @@ Result> RealtimeCommitProperties::Build( return merged_properties; } - PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap merged_offsets, - ReadOffsets(latest_snapshot, file_system)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap merged_offsets, + ReadOffsets(reset_all_realtime_progress ? std::nullopt : latest_snapshot, file_system)); + for (auto iter = merged_offsets.begin(); iter != merged_offsets.end();) { + bool removed = false; + for (const auto& partition_spec : removed_realtime_partitions) { + PAIMON_ASSIGN_OR_RAISE( + removed, PartitionUtils::MatchPartitionSpec(iter->first.partition, partition_spec, + partition_computer)); + if (removed) { + break; + } + } + if (removed) { + iter = merged_offsets.erase(iter); + } else { + ++iter; + } + } for (const auto& [partition_bucket, offset_range] : realtime_ranges) { if (partition_bucket.bucket < 0) { return Status::Invalid( @@ -221,6 +292,10 @@ Result> RealtimeCommitProperties::Build( } merged_offsets[partition_bucket] = offset_range.end; } + if (merged_offsets.empty()) { + merged_properties.erase(kOffsetsKey); + return merged_properties; + } PAIMON_ASSIGN_OR_RAISE( merged_properties[kOffsetsKey], WriteOffsets(merged_offsets, file_system, OffsetsDirectory(table_root, branch))); diff --git a/src/paimon/core/operation/commit/realtime_commit_properties.h b/src/paimon/core/operation/commit/realtime_commit_properties.h index 31a014b4..4c406920 100644 --- a/src/paimon/core/operation/commit/realtime_commit_properties.h +++ b/src/paimon/core/operation/commit/realtime_commit_properties.h @@ -33,6 +33,7 @@ namespace paimon { +class BinaryRowPartitionComputer; class FileSystem; class RealtimeCommitProperties { @@ -46,16 +47,35 @@ class RealtimeCommitProperties { static std::string OffsetsDirectory(const std::string& table_root, const std::string& branch); + /// Returns the offset file referenced by `snapshot`, if present. + static std::optional GetOffsetsPath(const Snapshot& snapshot); + static Result ReadOffsets(const std::optional& snapshot, const std::shared_ptr& file_system); + /// Returns whether all ranges are already covered by committed offsets. + /// + /// Ranges must either all immediately follow committed offsets or all be fully covered. + /// Mixed states, gaps, and partial overlaps are rejected. + static Result AreRangesCommitted( + const RealtimeOffsetMap& committed_offsets, + const std::map& realtime_ranges); + static Result SerializeOffsets(const RealtimeOffsetMap& offsets); /// Builds snapshot properties against `latest_snapshot` and applies real-time progress. + /// + /// A full-table replacement sets `reset_all_realtime_progress`. A partition overwrite or + /// drop lists only the affected partition specs in `removed_realtime_partitions`; offsets for + /// all other partitions are retained. The two reset forms are independent of the snapshot's + /// commit kind because an ordinary commit may also use `OVERWRITE` for conflict handling. static Result> Build( const std::map& properties, const std::optional& latest_snapshot, const std::map& realtime_ranges, + bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + const BinaryRowPartitionComputer& partition_computer, const std::shared_ptr& file_system, const std::string& table_root, const std::string& branch); diff --git a/src/paimon/core/operation/commit/realtime_commit_properties_test.cpp b/src/paimon/core/operation/commit/realtime_commit_properties_test.cpp index 68b15945..afe529e2 100644 --- a/src/paimon/core/operation/commit/realtime_commit_properties_test.cpp +++ b/src/paimon/core/operation/commit/realtime_commit_properties_test.cpp @@ -28,9 +28,12 @@ #include #include +#include "arrow/type.h" #include "gtest/gtest.h" +#include "paimon/common/utils/binary_row_partition_computer.h" #include "paimon/fs/file_system.h" #include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -71,6 +74,13 @@ class RealtimeCommitPropertiesTest : public testing::Test { ASSERT_NE(nullptr, directory_); file_system_ = directory_->GetFileSystem(); ASSERT_NE(nullptr, file_system_); + std::shared_ptr schema = arrow::schema( + {arrow::field("dt", arrow::utf8()), arrow::field("region", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(partition_computer_, + BinaryRowPartitionComputer::Create( + /*partition_keys=*/{"dt", "region"}, schema, + /*default_part_value=*/"__DEFAULT_PARTITION__", + /*legacy_partition_name_enabled=*/true, GetDefaultPool())); } Snapshot MakeSnapshot( @@ -120,6 +130,7 @@ class RealtimeCommitPropertiesTest : public testing::Test { std::unique_ptr directory_; std::shared_ptr file_system_; + std::unique_ptr partition_computer_; int32_t next_file_id_ = 0; }; @@ -256,6 +267,52 @@ TEST_F(RealtimeCommitPropertiesTest, SortProgress) { ASSERT_EQ(OffsetRange(5, 7), commits[2].offset_range); } +TEST_F(RealtimeCommitPropertiesTest, CheckRangesCommitted) { + RealtimePartitionBucket bucket0({{"dt", "2"}}, /*bucket=*/0); + RealtimePartitionBucket bucket1({{"dt", "2"}}, /*bucket=*/1); + RealtimeOffsetMap committed_offsets = {{bucket0, 4}, {bucket1, 9}}; + + std::map pending = {{bucket0, OffsetRange(4, 6)}, + {bucket1, OffsetRange(9, 11)}}; + ASSERT_OK_AND_ASSIGN(bool pending_committed, + RealtimeCommitProperties::AreRangesCommitted(committed_offsets, pending)); + ASSERT_FALSE(pending_committed); + + std::map covered = {{bucket0, OffsetRange(0, 4)}, + {bucket1, OffsetRange(5, 8)}}; + ASSERT_OK_AND_ASSIGN(bool covered_committed, + RealtimeCommitProperties::AreRangesCommitted(committed_offsets, covered)); + ASSERT_TRUE(covered_committed); + + std::map single_bucket_covered = { + {bucket0, OffsetRange(0, 4)}}; + ASSERT_OK_AND_ASSIGN( + bool single_bucket_covered_committed, + RealtimeCommitProperties::AreRangesCommitted(committed_offsets, single_bucket_covered)); + ASSERT_TRUE(single_bucket_covered_committed); + + std::map single_bucket_pending = { + {bucket1, OffsetRange(9, 11)}}; + ASSERT_OK_AND_ASSIGN( + bool single_bucket_pending_committed, + RealtimeCommitProperties::AreRangesCommitted(committed_offsets, single_bucket_pending)); + ASSERT_FALSE(single_bucket_pending_committed); + + std::map partial_overlap = {{bucket0, OffsetRange(3, 5)}}; + ASSERT_NOK_WITH_MSG( + RealtimeCommitProperties::AreRangesCommitted(committed_offsets, partial_overlap), + "partially overlaps"); + + std::map gap = {{bucket0, OffsetRange(5, 7)}}; + ASSERT_NOK_WITH_MSG(RealtimeCommitProperties::AreRangesCommitted(committed_offsets, gap), + "are not contiguous"); + + std::map mixed = {{bucket0, OffsetRange(0, 4)}, + {bucket1, OffsetRange(9, 11)}}; + ASSERT_NOK_WITH_MSG(RealtimeCommitProperties::AreRangesCommitted(committed_offsets, mixed), + "only partially covered"); +} + TEST_F(RealtimeCommitPropertiesTest, BuildRejectsInvalidProgress) { RealtimePartitionBucket bucket0({{"dt", "2"}}, /*bucket=*/0); RealtimeOffsetMap committed_offsets = {{bucket0, 1}}; @@ -265,20 +322,28 @@ TEST_F(RealtimeCommitPropertiesTest, BuildRejectsInvalidProgress) { {RealtimePartitionBucket({{"dt", "2"}}, /*bucket=*/-1), OffsetRange(0, 1)}}; ASSERT_NOK_WITH_MSG( RealtimeCommitProperties::Build(/*properties=*/{}, /*latest_snapshot=*/std::nullopt, - invalid_bucket, file_system_, directory_->Str(), "main"), + invalid_bucket, /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, + file_system_, directory_->Str(), "main"), "bucket -1 is invalid"); std::map gap = { {RealtimePartitionBucket({{"dt", "2"}}, /*bucket=*/0), OffsetRange(3, 5)}}; - ASSERT_NOK_WITH_MSG(RealtimeCommitProperties::Build(/*properties=*/{}, latest_snapshot, gap, - file_system_, directory_->Str(), "main"), - "are not contiguous"); + ASSERT_NOK_WITH_MSG( + RealtimeCommitProperties::Build(/*properties=*/{}, latest_snapshot, gap, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, + file_system_, directory_->Str(), "main"), + "are not contiguous"); std::map overlap = { {RealtimePartitionBucket({{"dt", "2"}}, /*bucket=*/0), OffsetRange(0, 2)}}; - ASSERT_NOK_WITH_MSG(RealtimeCommitProperties::Build(/*properties=*/{}, latest_snapshot, overlap, - file_system_, directory_->Str(), "main"), - "are not contiguous"); + ASSERT_NOK_WITH_MSG( + RealtimeCommitProperties::Build(/*properties=*/{}, latest_snapshot, overlap, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, + file_system_, directory_->Str(), "main"), + "are not contiguous"); RealtimeOffsetMap exhausted_offsets = {{bucket0, std::numeric_limits::max()}}; ASSERT_OK_AND_ASSIGN(Snapshot exhausted_snapshot, MakeSnapshotWithOffsets(exhausted_offsets)); @@ -287,6 +352,8 @@ TEST_F(RealtimeCommitPropertiesTest, BuildRejectsInvalidProgress) { OffsetRange(std::numeric_limits::max(), std::numeric_limits::max())}}; ASSERT_NOK_WITH_MSG( RealtimeCommitProperties::Build(/*properties=*/{}, exhausted_snapshot, after_max, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, file_system_, directory_->Str(), "main"), "offset range is invalid"); } @@ -300,16 +367,54 @@ TEST_F(RealtimeCommitPropertiesTest, BuildWithoutProgress) { ASSERT_OK_AND_ASSIGN(Properties inherited, RealtimeCommitProperties::Build( properties, std::optional(MakeSnapshot(latest_properties)), - /*realtime_ranges=*/{}, /*file_system=*/nullptr, + /*realtime_ranges=*/{}, /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, + /*file_system=*/nullptr, /*table_root=*/"", /*branch=*/"main")); ASSERT_EQ("value", inherited.at("custom")); ASSERT_EQ(latest_offsets_path, inherited.at(RealtimeCommitProperties::kOffsetsKey)); - ASSERT_OK_AND_ASSIGN(Properties unchanged, RealtimeCommitProperties::Build( - properties, /*latest_snapshot=*/std::nullopt, - /*realtime_ranges=*/{}, /*file_system=*/nullptr, - /*table_root=*/"", /*branch=*/"main")); + ASSERT_OK_AND_ASSIGN( + Properties unchanged, + RealtimeCommitProperties::Build(properties, /*latest_snapshot=*/std::nullopt, + /*realtime_ranges=*/{}, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, + /*file_system=*/nullptr, + /*table_root=*/"", /*branch=*/"main")); ASSERT_EQ(properties, unchanged); + + Properties properties_with_stale_offset = properties; + properties_with_stale_offset[RealtimeCommitProperties::kOffsetsKey] = "stale.offsets"; + ASSERT_OK_AND_ASSIGN( + Properties overwritten, + RealtimeCommitProperties::Build( + properties_with_stale_offset, std::optional(MakeSnapshot(latest_properties)), + /*realtime_ranges=*/{}, /*reset_all_realtime_progress=*/true, + /*removed_realtime_partitions=*/{}, *partition_computer_, + /*file_system=*/nullptr, /*table_root=*/"", /*branch=*/"main")); + ASSERT_EQ("value", overwritten.at("custom")); + ASSERT_EQ(0, overwritten.count(RealtimeCommitProperties::kOffsetsKey)); +} + +TEST_F(RealtimeCommitPropertiesTest, BuildRemovesOnlyOverwrittenPartitions) { + RealtimePartitionBucket dt2_bucket0({{"dt", "2"}}, /*bucket=*/0); + RealtimePartitionBucket dt2_bucket1({{"dt", "2"}}, /*bucket=*/1); + RealtimePartitionBucket dt3_bucket0({{"dt", "3"}}, /*bucket=*/0); + RealtimeOffsetMap committed_offsets = {{dt2_bucket0, 3}, {dt2_bucket1, 4}, {dt3_bucket0, 5}}; + ASSERT_OK_AND_ASSIGN(Snapshot latest_snapshot, MakeSnapshotWithOffsets(committed_offsets)); + std::vector> removed_partitions = {{{"dt", "2"}}}; + + ASSERT_OK_AND_ASSIGN(Properties properties, + RealtimeCommitProperties::Build( + /*properties=*/{}, latest_snapshot, /*realtime_ranges=*/{}, + /*reset_all_realtime_progress=*/false, removed_partitions, + *partition_computer_, file_system_, directory_->Str(), "main")); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap actual, + RealtimeCommitProperties::ReadOffsets( + std::optional(MakeSnapshot(properties)), file_system_)); + RealtimeOffsetMap expected = {{dt3_bucket0, 5}}; + ASSERT_EQ(expected, actual); } TEST_F(RealtimeCommitPropertiesTest, BuildWritesMergedProgress) { @@ -322,10 +427,13 @@ TEST_F(RealtimeCommitPropertiesTest, BuildWritesMergedProgress) { {RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0), OffsetRange(7, 9)}}; std::map properties = {{"custom", "value"}}; - ASSERT_OK_AND_ASSIGN(Properties merged, - RealtimeCommitProperties::Build( - properties, std::optional(MakeSnapshot(latest_properties)), - ranges, file_system_, directory_->Str(), "main")); + ASSERT_OK_AND_ASSIGN( + Properties merged, + RealtimeCommitProperties::Build( + properties, std::optional(MakeSnapshot(latest_properties)), ranges, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, file_system_, + directory_->Str(), "main")); ASSERT_EQ("value", merged.at("custom")); ASSERT_NE(latest_offsets_path, merged.at(RealtimeCommitProperties::kOffsetsKey)); @@ -344,6 +452,8 @@ TEST_F(RealtimeCommitPropertiesTest, BuildRequiresFileSystem) { {RealtimePartitionBucket({{"dt", "2"}}, /*bucket=*/0), OffsetRange(0, 2)}}; ASSERT_NOK_WITH_MSG(RealtimeCommitProperties::Build( /*properties=*/{}, /*latest_snapshot=*/std::nullopt, ranges, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, *partition_computer_, /*file_system=*/nullptr, directory_->Str(), "main"), "file system is null"); } diff --git a/src/paimon/core/operation/data_evolution_split_read.cpp b/src/paimon/core/operation/data_evolution_split_read.cpp index 52a1bb51..9e71fe47 100644 --- a/src/paimon/core/operation/data_evolution_split_read.cpp +++ b/src/paimon/core/operation/data_evolution_split_read.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -53,8 +54,13 @@ #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/apply_deletion_vector_batch_reader.h" #include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/file_index_evaluator.h" #include "paimon/core/utils/blob_view_lookup.h" #include "paimon/core/utils/data_evolution_utils.h" +#include "paimon/core/utils/field_mapping.h" +#include "paimon/file_index/bitmap_index_result.h" +#include "paimon/file_index/file_index_result.h" +#include "paimon/predicate/predicate_utils.h" namespace paimon { namespace { @@ -389,12 +395,22 @@ Result> DataEvolutionSplitRead::InnerCreateReader( path_factory_->CreateDataFilePathFactory(split_impl->Partition(), split_impl->Bucket())); auto metas = split_impl->DataFiles(); DeletionVector::Factory split_dv_factory = CreateSplitDvFactory(*split_impl); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr push_down_predicate, + CreatePushDownPredicate(context_->GetPredicate(), raw_read_schema_)); PAIMON_ASSIGN_OR_RAISE(std::vector>> split_by_row_id, MergeRangesAndSort(std::move(metas))); std::vector> sub_readers; for (const std::vector>& need_merge_files : split_by_row_id) { + if (need_merge_files.size() > 1) { + PAIMON_ASSIGN_OR_RAISE( + bool skip_group, + SkipByFileIndex(push_down_predicate, need_merge_files, data_file_path_factory)); + if (skip_group) { + continue; + } + } PAIMON_ASSIGN_OR_RAISE(std::optional group_dv, ReadGroupDeletionVector(need_merge_files, split_dv_factory)); PAIMON_ASSIGN_OR_RAISE(DeletionVector::Factory group_dv_factory, @@ -404,10 +420,15 @@ Result> DataEvolutionSplitRead::InnerCreateReader( PAIMON_ASSIGN_OR_RAISE( std::vector> raw_file_readers, CreateRawFileReaders(split_impl->Partition(), need_merge_files, raw_read_schema_, - /*predicate=*/nullptr, group_dv_factory, row_ranges, + push_down_predicate, group_dv_factory, row_ranges, data_file_path_factory, /*extra_format_options=*/{})); - assert(raw_file_readers.size() == 1); + if (raw_file_readers.empty()) { + continue; + } + if (raw_file_readers.size() != 1) { + return Status::Invalid("Single-file data evolution group created multiple readers"); + } sub_readers.push_back(std::move(raw_file_readers[0])); } else { PAIMON_ASSIGN_OR_RAISE( @@ -424,17 +445,110 @@ Result> DataEvolutionSplitRead::InnerCreateReader( return std::make_unique(std::move(batch_reader), pool_); } +Result> DataEvolutionSplitRead::CreatePushDownPredicate( + const std::shared_ptr& predicate, + const std::shared_ptr& read_schema) { + std::map picked_field_name_to_idx; + for (int32_t i = 0; i < read_schema->num_fields(); ++i) { + const std::string& field_name = read_schema->field(i)->name(); + if (!SpecialFields::IsSystemField(field_name)) { + picked_field_name_to_idx.emplace(field_name, i); + } + } + return PredicateUtils::CreatePickedFieldFilter(predicate, picked_field_name_to_idx); +} + +Result DataEvolutionSplitRead::SkipByFileIndex( + const std::shared_ptr& predicate, + const std::vector>& files, + const std::shared_ptr& data_file_path_factory) const { + if (!options_.FileIndexReadEnabled() || !predicate) { + return false; + } + + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr field_mapping_builder, + FieldMappingBuilder::Create(raw_read_schema_, context_->GetPartitionKeys(), predicate)); + std::set claimed_field_ids; + for (const auto& file : files) { + // Blob and vector-store files may cover only part of the row range, so their indexes + // cannot prove that the complete merged group misses the predicate. + if (!DataEvolutionUtils::IsNormalFile(file->file_name)) { + continue; + } + + std::shared_ptr data_schema = context_->GetTableSchema(); + if (file->schema_id != data_schema->Id()) { + PAIMON_ASSIGN_OR_RAISE(data_schema, schema_manager_->ReadSchema(file->schema_id)); + } + std::vector written_fields; + if (file->write_cols) { + std::vector data_write_cols; + data_write_cols.reserve(file->write_cols->size()); + for (const auto& write_col : file->write_cols.value()) { + if (!SpecialFields::IsSystemField(write_col)) { + data_write_cols.push_back(write_col); + } + } + PAIMON_ASSIGN_OR_RAISE(written_fields, data_schema->GetFields(data_write_cols)); + } else { + written_fields = data_schema->Fields(); + } + + std::set overwritten_field_names; + for (const auto& field : written_fields) { + if (!claimed_field_ids.insert(field.Id()).second) { + overwritten_field_names.insert(field.Name()); + } + } + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr field_mapping, + field_mapping_builder->CreateFieldMapping(written_fields)); + std::shared_ptr data_predicate = + field_mapping->non_partition_info.non_partition_filter; + if (!overwritten_field_names.empty()) { + PAIMON_ASSIGN_OR_RAISE(data_predicate, PredicateUtils::ExcludePredicateWithFields( + data_predicate, overwritten_field_names)); + } + if (!data_predicate) { + continue; + } + + auto written_schema = DataField::ConvertDataFieldsToArrowSchema(written_fields); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr index_result, + FileIndexEvaluator::Evaluate(written_schema, data_predicate, data_file_path_factory, + file, options_.GetFileSystem(), pool_)); + PAIMON_ASSIGN_OR_RAISE(bool is_remain, index_result->IsRemain()); + if (!is_remain) { + return true; + } + } + return false; +} + Result> DataEvolutionSplitRead::ApplyIndexAndDvReaderIfNeeded( std::unique_ptr&& file_reader, const std::shared_ptr& file, const std::shared_ptr& data_schema, const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const { - if (predicate) { - assert(false); - // as DataEvolutionSplitRead will skip predicate - return Status::Invalid("DataEvolutionSplitRead do not support predicate"); + std::shared_ptr file_index_result; + if (options_.FileIndexReadEnabled()) { + PAIMON_ASSIGN_OR_RAISE( + file_index_result, + FileIndexEvaluator::Evaluate(data_schema, predicate, data_file_path_factory, file, + options_.GetFileSystem(), pool_)); + PAIMON_ASSIGN_OR_RAISE(bool is_remain, file_index_result->IsRemain()); + if (!is_remain) { + return std::unique_ptr(); + } + } + const RoaringBitmap32* index_selection = nullptr; + if (auto* bitmap_index = dynamic_cast(file_index_result.get())) { + PAIMON_ASSIGN_OR_RAISE(index_selection, bitmap_index->GetBitmap()); } + // the factory is per row range group and already returns a view taking file-local positions. // Unlike RawFileSplitRead the vector is not folded into the format reader's selection: it is // no BitmapDeletionVector, and the blob fallback path's gap segments have no format reader. @@ -444,10 +558,19 @@ Result> DataEvolutionSplitRead::ApplyIndexAndDv } PAIMON_ASSIGN_OR_RAISE(std::optional selection_row_ids, file->ToFileSelection(row_ranges)); + if (index_selection) { + if (selection_row_ids) { + selection_row_ids.value() &= *index_selection; + } else { + selection_row_ids = *index_selection; + } + } + if (selection_row_ids && selection_row_ids->IsEmpty()) { + return std::unique_ptr(); + } ::ArrowSchema c_read_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); - PAIMON_RETURN_NOT_OK( - file_reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, selection_row_ids)); + PAIMON_RETURN_NOT_OK(file_reader->SetReadSchema(&c_read_schema, predicate, selection_row_ids)); std::unique_ptr reader; if (!file_reader->SupportPreciseBitmapSelection() && selection_row_ids) { diff --git a/src/paimon/core/operation/data_evolution_split_read.h b/src/paimon/core/operation/data_evolution_split_read.h index 983ca29d..b4f80adb 100644 --- a/src/paimon/core/operation/data_evolution_split_read.h +++ b/src/paimon/core/operation/data_evolution_split_read.h @@ -65,12 +65,13 @@ struct DeletionFile; /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ApplyBitmapIndexBatchReader) /// ->(CompleteRowTrackingFieldsBatchReader)->(ShreddingFileReader) /// ->(MapSharedShreddingFileReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(VectorFileBatchReader)->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader /// /// -/// A union `SplitRead` to read multiple inner files to merge columns, note that this class -/// does not support filtering push down: a predicate would have to be evaluated consistently -/// across the files being merged, which is not implemented here. +/// A union `SplitRead` to read multiple inner files to merge columns. A single-file row range +/// group gets both file-index and format-level predicate pushdown. A merged group only uses file +/// indexes to skip the whole group: filtering its child readers independently would break their +/// positional alignment. /// /// Deletion vectors are supported: a row range group's vector is maintained against the /// group's anchor file (DataEvolutionUtils::RetrieveAnchorFile), so its positions are @@ -172,6 +173,21 @@ class DataEvolutionSplitRead : public AbstractSplitRead { const std::shared_ptr& data_split, const std::optional>& row_ranges) const; + /// Keeps top-level conjuncts whose fields all belong to `read_schema`, excluding conjuncts + /// over system fields. The returned predicate is for pushdown only; the original predicate is + /// still evaluated as a residual filter when requested by the read context. + static Result> CreatePushDownPredicate( + const std::shared_ptr& predicate, + const std::shared_ptr& read_schema); + + /// Returns true when file indexes prove that no row in a merged row range group can match. + /// Only normal files are considered, and an older copy of a field is excluded after a newer + /// file has claimed the same field id. + Result SkipByFileIndex( + const std::shared_ptr& predicate, + const std::vector>& files, + const std::shared_ptr& data_file_path_factory) const; + /// Builds the deletion vector factory over the split's deletion files, keyed by data file /// name. Only anchor files carry one. Returns a null factory when the split has none. DeletionVector::Factory CreateSplitDvFactory(const DataSplitImpl& split_impl) const; diff --git a/src/paimon/core/operation/data_evolution_split_read_test.cpp b/src/paimon/core/operation/data_evolution_split_read_test.cpp index 809933e9..03bc95b6 100644 --- a/src/paimon/core/operation/data_evolution_split_read_test.cpp +++ b/src/paimon/core/operation/data_evolution_split_read_test.cpp @@ -25,6 +25,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/table/special_fields.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" @@ -37,6 +38,8 @@ #include "paimon/executor.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -102,6 +105,30 @@ class DataEvolutionSplitReadTest : public ::testing::Test { std::shared_ptr pool_ = GetDefaultPool(); }; +TEST_F(DataEvolutionSplitReadTest, TestCreatePushDownPredicate) { + auto f0_predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, Literal(1)); + auto f1_predicate = + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f1", FieldType::INT, Literal(2)); + auto row_id_predicate = PredicateBuilder::Equal( + /*field_index=*/2, SpecialFields::RowId().Name(), FieldType::BIGINT, Literal(3l)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate, + PredicateBuilder::And({f0_predicate, f1_predicate, row_id_predicate})); + + auto read_schema = DataField::ConvertDataFieldsToArrowSchema( + {DataField(0, arrow::field("f0", arrow::int32())), SpecialFields::RowId()}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr push_down, + DataEvolutionSplitRead::CreatePushDownPredicate(predicate, read_schema)); + ASSERT_TRUE(push_down); + ASSERT_EQ(*push_down, *f0_predicate); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({f0_predicate, f1_predicate})); + ASSERT_OK_AND_ASSIGN( + push_down, DataEvolutionSplitRead::CreatePushDownPredicate(or_predicate, read_schema)); + ASSERT_FALSE(push_down); +} + TEST_F(DataEvolutionSplitReadTest, TestAddSingleBlobEntry) { auto blob_entry = CreateBlobFile("blob1", /*first_row_id=*/0, /*row_count=*/100, diff --git a/src/paimon/core/operation/expire_snapshots.cpp b/src/paimon/core/operation/expire_snapshots.cpp index d6f4b121..a5bd2505 100644 --- a/src/paimon/core/operation/expire_snapshots.cpp +++ b/src/paimon/core/operation/expire_snapshots.cpp @@ -38,6 +38,7 @@ #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/manifest/manifest_list.h" +#include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/snapshot.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" @@ -50,13 +51,14 @@ ExpireSnapshots::ExpireSnapshots(const std::shared_ptr& snapsho const std::shared_ptr& manifest_list, const std::shared_ptr& manifest_file, const std::shared_ptr& fs, const ExpireConfig& config, - const std::shared_ptr& executor) + bool realtime_enabled, const std::shared_ptr& executor) : snapshot_manager_(snapshot_manager), path_factory_(path_factory), manifest_list_(manifest_list), manifest_file_(manifest_file), fs_(fs), config_(config), + realtime_enabled_(realtime_enabled), executor_(executor), logger_(Logger::GetLogger("ExpireSnapshots")) {} @@ -157,8 +159,24 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, std::vector retained_snapshots; PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(end_exclusive_id)); retained_snapshots.push_back(snapshot); + std::set retained_offset_files; + if (realtime_enabled_) { + PAIMON_ASSIGN_OR_RAISE(std::vector all_snapshots, + snapshot_manager_->GetAllSnapshots()); + for (const Snapshot& retained_snapshot : all_snapshots) { + if (retained_snapshot.Id() < end_exclusive_id) { + continue; + } + std::optional offsets_path = + RealtimeCommitProperties::GetOffsetsPath(retained_snapshot); + if (offsets_path) { + retained_offset_files.insert(offsets_path.value()); + } + } + } std::set skipping_sets; PAIMON_RETURN_NOT_OK(GetManifestSkippingSet(retained_snapshots, &skipping_sets)); + std::set expired_offset_files; for (int64_t id = begin_inclusive_id; id < end_exclusive_id; id++) { PAIMON_LOG_DEBUG(logger_, "Ready to delete manifests in snapshot #%ld", id); PAIMON_ASSIGN_OR_RAISE(bool exist, snapshot_manager_->SnapshotExists(id)); @@ -169,10 +187,24 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(id)); PAIMON_RETURN_NOT_OK(CleanUnusedManifests(snapshot.BaseManifestList(), skipping_sets)); PAIMON_RETURN_NOT_OK(CleanUnusedManifests(snapshot.DeltaManifestList(), skipping_sets)); + if (realtime_enabled_) { + std::optional offsets_path = + RealtimeCommitProperties::GetOffsetsPath(snapshot); + if (offsets_path) { + expired_offset_files.insert(offsets_path.value()); + } + } auto status = fs_->Delete(snapshot_manager_->SnapshotPath(id)); // delete quietly will ignore any status error (void)status; } + for (const std::string& offsets_path : expired_offset_files) { + if (retained_offset_files.count(offsets_path) == 0) { + auto status = fs_->Delete(offsets_path); + // Orphan cleanup can retry offset files that fail to delete here. + (void)status; + } + } PAIMON_RETURN_NOT_OK(snapshot_manager_->CommitEarliestHint(end_exclusive_id)); return end_exclusive_id - begin_inclusive_id; } diff --git a/src/paimon/core/operation/expire_snapshots.h b/src/paimon/core/operation/expire_snapshots.h index 75b5fff5..238f599e 100644 --- a/src/paimon/core/operation/expire_snapshots.h +++ b/src/paimon/core/operation/expire_snapshots.h @@ -51,7 +51,7 @@ class ExpireSnapshots { const std::shared_ptr& manifest_list, const std::shared_ptr& manifest_file, const std::shared_ptr& fs, const ExpireConfig& config, - const std::shared_ptr& executor); + bool realtime_enabled, const std::shared_ptr& executor); Result Expire(); @@ -74,6 +74,7 @@ class ExpireSnapshots { std::shared_ptr manifest_file_; std::shared_ptr fs_; ExpireConfig config_; + bool realtime_enabled_; std::shared_ptr executor_; std::unordered_map> deletion_buckets_; diff --git a/src/paimon/core/operation/expire_snapshots_test.cpp b/src/paimon/core/operation/expire_snapshots_test.cpp index d6b965ba..6a9e46ba 100644 --- a/src/paimon/core/operation/expire_snapshots_test.cpp +++ b/src/paimon/core/operation/expire_snapshots_test.cpp @@ -170,7 +170,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_OK_AND_ASSIGN(int32_t count, expire.Expire()); ASSERT_EQ(count, 0); } @@ -178,7 +178,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "0"}})); ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { @@ -186,7 +186,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "9"}})); ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { @@ -194,7 +194,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_OK_AND_ASSIGN(int32_t count, expire.Expire()); ASSERT_EQ(count, 0); } @@ -204,7 +204,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { {Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { @@ -212,7 +212,7 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); ExpireSnapshots expire(nullptr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } } @@ -222,7 +222,7 @@ TEST_F(ExpireSnapshotsTest, TestGetDataFileToDelete) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); { ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); std::map data_file_to_delete; std::vector data_file_entries; data_file_entries.push_back(CreateManifestEntry("file1", /*bucket=*/0, FileKind::Delete())); @@ -237,7 +237,7 @@ TEST_F(ExpireSnapshotsTest, TestGetDataFileToDelete) { } { ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), executor_); + options.GetExpireConfig(), options.RealtimeEnabled(), executor_); std::map data_file_to_delete; std::vector data_file_entries; data_file_entries.push_back(CreateManifestEntry("file1", /*bucket=*/0, FileKind::Add())); diff --git a/src/paimon/core/operation/file_store_commit.cpp b/src/paimon/core/operation/file_store_commit.cpp index b2ba17fa..ad0942ad 100644 --- a/src/paimon/core/operation/file_store_commit.cpp +++ b/src/paimon/core/operation/file_store_commit.cpp @@ -175,7 +175,7 @@ Result> FileStoreCommit::Create( auto expire_snapshots = std::make_shared( snapshot_manager, path_factory, manifest_list, manifest_file, options.GetFileSystem(), - options.GetExpireConfig(), ctx->GetExecutor()); + options.GetExpireConfig(), options.RealtimeEnabled(), ctx->GetExecutor()); CommitScanner::ScanSupplier scan_supplier; if (table_schema.value()->PrimaryKeys().empty()) { diff --git a/src/paimon/core/operation/file_store_commit_impl.cpp b/src/paimon/core/operation/file_store_commit_impl.cpp index 5388b9e8..8f2b3c73 100644 --- a/src/paimon/core/operation/file_store_commit_impl.cpp +++ b/src/paimon/core/operation/file_store_commit_impl.cpp @@ -80,6 +80,7 @@ #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/utils/duration.h" #include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/core/utils/partition_utils.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/file_store_write.h" #include "paimon/fs/file_system.h" @@ -96,17 +97,6 @@ constexpr const char* kCommitStrictModeLastSafeSnapshot = "commit.strict-mode.la constexpr const char* kSequenceSnapshotOrdering = "sequence.snapshot-ordering"; constexpr const char* kPkClusteringOverride = "pk-clustering-override"; -bool MatchPartitionSpec(const std::map& partition, - const std::map& partition_spec) { - for (const auto& [key, value] : partition_spec) { - auto iter = partition.find(key); - if (iter == partition.end() || iter->second != value) { - return false; - } - } - return true; -} - } // namespace Status FileStoreCommitImpl::ValidateCommitOptions(const CoreOptions& options) { @@ -337,7 +327,6 @@ Result FileStoreCommitImpl::RollbackToAsLatest(int64_t target_snapshot_id) // snapshots between the target and the previous latest, breaking the global uniqueness of // _ROW_ID. Keep the larger of the previous latest and the target nextRowId. std::optional next_row_id = std::max(latest.NextRowId(), target_snapshot.NextRowId()); - int64_t delta_record_count = ManifestEntry::RecordCountAdd(delta_files) - ManifestEntry::RecordCountDelete(delta_files); Snapshot new_snapshot( @@ -668,7 +657,10 @@ Status FileStoreCommitImpl::ExecuteOverwrite( PAIMON_ASSIGN_OR_RAISE(partition_map, PartitionToMap(entry.Partition())); bool belongs_to_overwrite_partition = false; for (const auto& partition_spec : partitions) { - if (MatchPartitionSpec(partition_map, partition_spec)) { + PAIMON_ASSIGN_OR_RAISE( + bool matched, PartitionUtils::MatchPartitionSpec(partition_map, partition_spec, + *partition_computer_)); + if (matched) { belongs_to_overwrite_partition = true; break; } @@ -708,6 +700,8 @@ Status FileStoreCommitImpl::ExecuteOverwrite( changes->compact_index_files, identifier, watermark, committable->Properties(), /*realtime_ranges=*/{}, Snapshot::CommitKind::Compact(), + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, /*detect_conflicts=*/true, /*retry_on_conflict=*/true)); *attempt += cnt; @@ -809,8 +803,12 @@ Result FileStoreCommitImpl::TryOverwrite( const std::map& properties) { std::shared_ptr changes_provider = commit_scanner_->OverwriteChangesProvider(partitions, changes, index_entries); + // ExecuteOverwrite has already resolved dynamic overwrite to the concrete affected + // partitions. Only an empty final partition list denotes a full-table replacement. + const bool reset_all_realtime_progress = partitions.empty(); return TryCommit(changes_provider, commit_identifier, watermark, properties, /*realtime_ranges=*/{}, Snapshot::CommitKind::Overwrite(), + reset_all_realtime_progress, partitions, /*detect_conflicts=*/true, /*retry_on_conflict=*/true); } @@ -859,7 +857,9 @@ Status FileStoreCommitImpl::Commit( TryCommit(changes.append_table_files, changes.append_changelog, changes.append_index_files, committable->Identifier(), committable->Watermark(), committable->Properties(), realtime_ranges, - commit_kind, check_append_files, retry_on_conflict)); + commit_kind, + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, check_append_files, retry_on_conflict)); attempt += cnt; generated_snapshot += 1; } @@ -870,6 +870,8 @@ Status FileStoreCommitImpl::Commit( changes.compact_index_files, committable->Identifier(), committable->Watermark(), committable->Properties(), /*realtime_ranges=*/{}, Snapshot::CommitKind::Compact(), + /*reset_all_realtime_progress=*/false, + /*removed_realtime_partitions=*/{}, /*detect_conflicts=*/true, retry_on_conflict)); attempt += cnt; generated_snapshot += 1; @@ -889,6 +891,9 @@ Status FileStoreCommitImpl::Commit( Result FileStoreCommitImpl::CommitWithProgress( const std::vector& realtime_commits, int64_t identifier, std::optional watermark) { + if (!options_.RealtimeEnabled()) { + return Status::Invalid("CommitWithProgress requires realtime.enabled=true"); + } if (realtime_commits.empty()) { return Status::Invalid("real-time commits must not be empty"); } @@ -931,6 +936,30 @@ Result FileStoreCommitImpl::CommitWithProgress( std::shared_ptr committable = CreateManifestCommittable(identifier, commit_messages, watermark, /*properties=*/{}); + PAIMON_ASSIGN_OR_RAISE(std::vector> pending_committables, + FilterCommitted({committable})); + const bool identifier_committed = pending_committables.empty(); + + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager_->LatestSnapshot()); + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(latest_snapshot, fs_)); + PAIMON_ASSIGN_OR_RAISE(bool ranges_committed, RealtimeCommitProperties::AreRangesCommitted( + committed_offsets, realtime_ranges)); + if (ranges_committed != identifier_committed) { + return Status::Invalid( + ranges_committed + ? "real-time offset ranges were committed by another commit user or identifier" + : "real-time commit identifier was committed without the requested offset ranges"); + } + if (ranges_committed) { + if (!latest_snapshot) { + return Status::Invalid("real-time commit ranges are covered without a snapshot"); + } + return latest_snapshot->Id(); + } + + PAIMON_RETURN_NOT_OK(CheckFilesExistence(pending_committables)); const int64_t previous_snapshot_id = last_committed_snapshot_id_; PAIMON_RETURN_NOT_OK(Commit(committable, /*check_append_files=*/false, /*retry_on_conflict=*/false, realtime_ranges)); @@ -946,18 +975,23 @@ Result FileStoreCommitImpl::TryCommit( const std::vector& index_entries, int64_t identifier, std::optional watermark, const std::map& properties, const std::map& realtime_ranges, - Snapshot::CommitKind commit_kind, bool detect_conflicts, bool retry_on_conflict) { + Snapshot::CommitKind commit_kind, bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + bool detect_conflicts, bool retry_on_conflict) { std::shared_ptr changes_provider = CommitChangesProvider::Provider(delta_files, changelog_files, index_entries); return TryCommit(changes_provider, identifier, watermark, properties, realtime_ranges, - commit_kind, detect_conflicts, retry_on_conflict); + commit_kind, reset_all_realtime_progress, removed_realtime_partitions, + detect_conflicts, retry_on_conflict); } Result FileStoreCommitImpl::TryCommit( const std::shared_ptr& changes_provider, int64_t identifier, std::optional watermark, const std::map& properties, const std::map& realtime_ranges, - Snapshot::CommitKind commit_kind, bool detect_conflicts, bool retry_on_conflict) { + Snapshot::CommitKind commit_kind, bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + bool detect_conflicts, bool retry_on_conflict) { int32_t retry_count = 0; int64_t start_millis = DateTimeUtils::GetCurrentUTCTimeUs() / 1000; while (true) { @@ -968,7 +1002,9 @@ Result FileStoreCommitImpl::TryCommit( using SnapshotProperties = std::map; PAIMON_ASSIGN_OR_RAISE( SnapshotProperties snapshot_properties, - RealtimeCommitProperties::Build(properties, latest_snapshot, realtime_ranges, fs_, + RealtimeCommitProperties::Build(properties, latest_snapshot, realtime_ranges, + reset_all_realtime_progress, + removed_realtime_partitions, *partition_computer_, fs_, root_path_, snapshot_manager_->Branch())); PAIMON_ASSIGN_OR_RAISE( bool commit_success, diff --git a/src/paimon/core/operation/file_store_commit_impl.h b/src/paimon/core/operation/file_store_commit_impl.h index 7c804414..ca4de22f 100644 --- a/src/paimon/core/operation/file_store_commit_impl.h +++ b/src/paimon/core/operation/file_store_commit_impl.h @@ -185,21 +185,23 @@ class FileStoreCommitImpl : public FileStoreCommit { void ReportCommit(const ManifestEntryChanges& changes, int64_t commit_duration, int32_t generated_snapshot, int32_t attempt); - Result TryCommit(const std::vector& delta_files, - const std::vector& changelog_files, - const std::vector& index_entries, - int64_t identifier, std::optional watermark, - const std::map& properties, - const std::map& realtime_ranges, - Snapshot::CommitKind commit_kind, bool detect_conflicts, - bool retry_on_conflict); - - Result TryCommit(const std::shared_ptr& changes_provider, - int64_t identifier, std::optional watermark, - const std::map& properties, - const std::map& realtime_ranges, - Snapshot::CommitKind commit_kind, bool detect_conflicts, - bool retry_on_conflict); + Result TryCommit( + const std::vector& delta_files, + const std::vector& changelog_files, + const std::vector& index_entries, int64_t identifier, + std::optional watermark, const std::map& properties, + const std::map& realtime_ranges, + Snapshot::CommitKind commit_kind, bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + bool detect_conflicts, bool retry_on_conflict); + + Result TryCommit( + const std::shared_ptr& changes_provider, int64_t identifier, + std::optional watermark, const std::map& properties, + const std::map& realtime_ranges, + Snapshot::CommitKind commit_kind, bool reset_all_realtime_progress, + const std::vector>& removed_realtime_partitions, + bool detect_conflicts, bool retry_on_conflict); Result TryCommitOnce(const std::vector& delta_files, const std::vector& changelog_files, diff --git a/src/paimon/core/operation/file_store_scan.cpp b/src/paimon/core/operation/file_store_scan.cpp index 865e006f..f21b0bb7 100644 --- a/src/paimon/core/operation/file_store_scan.cpp +++ b/src/paimon/core/operation/file_store_scan.cpp @@ -365,6 +365,33 @@ Status FileStoreScan::StoreSnapshotLiveManifestEntries( Status FileStoreScan::ReadAndMergeBucketFileEntries( const std::vector& manifest_metas, int32_t bucket, std::vector* merged_entries) const { + if (core_options_.ScanManifestEntryLazyDecodeEnabled()) { + std::vector>>> futures; + futures.reserve(manifest_metas.size()); + for (const auto& meta : manifest_metas) { + auto read_meta_task = [this, meta, bucket]() -> Result> { + std::vector bucket_entries; + PAIMON_RETURN_NOT_OK( + manifest_file_->ReadBucketEntries(meta.FileName(), bucket, &bucket_entries)); + return bucket_entries; + }; + futures.push_back(Via(executor_.get(), read_meta_task)); + } + + std::vector bucket_entries; + std::vector>> entry_lists = CollectAll(futures); + for (auto& entry_list : entry_lists) { + if (!entry_list.ok()) { + return entry_list.status(); + } + bucket_entries.reserve(bucket_entries.size() + entry_list.value().size()); + for (auto& entry : entry_list.value()) { + bucket_entries.emplace_back(std::move(entry)); + } + } + return MergeLiveEntries(bucket_entries, merged_entries); + } + std::vector unmerged_entries; std::vector entries; PAIMON_RETURN_NOT_OK(ReadFileEntries(manifest_metas, &entries, /*apply_scan_filter=*/false)); diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 940608c7..84a32476 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -36,6 +36,7 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" @@ -120,6 +121,9 @@ Result> FileStoreWrite::Create(std::unique_ptrIgnorePreviousFiles(); + if (ctx->GetRealtimeContext() && !options.RealtimeEnabled()) { + return Status::Invalid("real-time write requires realtime.enabled=true"); + } if (schema->PrimaryKeys().empty()) { // append table bool need_dv_maintainer_factory = options.DeletionVectorsEnabled(); @@ -194,7 +198,26 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - return Status::Invalid("real-time write currently supports append tables only"); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options, *schema)); + if (ignore_previous_files) { + return Status::NotImplemented( + "PK realtime v1 requires restore from the latest snapshot"); + } + if (!ctx->GetWriteSchema().empty()) { + return Status::NotImplemented( + "PK realtime v1 does not support a custom write schema"); + } + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager->LatestSnapshot()); + if (latest_snapshot) { + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap realtime_committed_offsets, + RealtimeCommitProperties::ReadOffsets( + latest_snapshot, options.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(ctx->GetRealtimeContext())); + PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( + latest_snapshot->Id(), realtime_committed_offsets)); + } } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( @@ -250,7 +273,8 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory, io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options, ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), - ctx->EnableMultiThreadSpill(), ctx->GetExecutor(), ctx->GetMemoryPool()); + ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), + ctx->GetMemoryPool()); } } diff --git a/src/paimon/core/operation/internal_read_context.h b/src/paimon/core/operation/internal_read_context.h index f33b7a35..8ef9f2d2 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -96,8 +96,8 @@ class InternalReadContext { return read_context_->GetRealtimeContext(); } - PrefetchCacheMode GetPrefetchCacheMode() const { - return read_context_->GetPrefetchCacheMode(); + bool ReadAheadCacheEnabled() const { + return read_context_->ReadAheadCacheEnabled(); } const CacheConfig& GetCacheConfig() const { diff --git a/src/paimon/core/operation/key_value_file_store_write.cpp b/src/paimon/core/operation/key_value_file_store_write.cpp index 08c5ea0c..d8e7f5d1 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -18,21 +18,30 @@ #include "paimon/core/operation/key_value_file_store_write.h" +#include #include +#include "arrow/c/bridge.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/mergetree/levels.h" #include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/realtime/realtime_context.h" namespace arrow { class Schema; @@ -60,6 +69,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, bool enable_multi_thread_spill, + const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool) : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, table_schema, schema, /*write_schema=*/schema, @@ -67,6 +77,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, executor, pool), enable_multi_thread_spill_(enable_multi_thread_spill), + realtime_context_(realtime_context), key_comparator_(key_comparator), user_defined_seq_comparator_(user_defined_seq_comparator), merge_function_wrapper_(merge_function_wrapper), @@ -74,7 +85,11 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( options_, key_comparator_, user_defined_seq_comparator_, compaction_metrics_, table_schema_, schema_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, root_path_, pool_)), - logger_(Logger::GetLogger("KeyValueFileStoreWrite")) {} + logger_(Logger::GetLogger("KeyValueFileStoreWrite")) { + if (realtime_context_) { + writer_memory_manager_ = std::make_unique(); + } +} Result> KeyValueFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { @@ -109,19 +124,73 @@ Result> KeyValueFileStoreWrite::CreateWriter( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr levels, Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); - auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr compact_manager, - compact_manager_factory_->CreateCompactManager(partition, bucket, compact_strategy, - compact_executor_, levels, dv_maintainer)); + std::map partition_map; + std::shared_ptr compact_manager; + std::shared_ptr realtime_context_impl; + std::optional realtime_store_state; + if (realtime_context_) { + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + file_store_path_factory_->GeneratePartitionVector(partition)); + partition_map = + std::map(partition_values.begin(), partition_values.end()); + PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); + if (schema_->GetFieldByName(SpecialFields::RealtimeOffset().Name())) { + return Status::Invalid("PK real-time write schema contains reserved transport field " + + SpecialFields::RealtimeOffset().Name()); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + prepared_fields.insert(prepared_fields.end(), schema_->fields().begin(), + schema_->fields().end()); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(std::move(prepared_fields)), c_write_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(c_write_schema), options_.ToMap(), pool_, + partition_map, bucket, RealtimeStoreMode::PRIMARY_KEY})); + realtime_store_state = std::move(store_state); + compact_manager = std::make_shared(); + } else { + auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); + PAIMON_ASSIGN_OR_RAISE(compact_manager, compact_manager_factory_->CreateCompactManager( + partition, bucket, compact_strategy, + compact_executor_, levels, dv_maintainer)); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, - options_, compact_manager, io_manager_, enable_multi_thread_spill_, pool_)); - return writer; + options_, compact_manager, realtime_context_ ? nullptr : io_manager_, + enable_multi_thread_spill_, pool_)); + if (!realtime_context_) { + return std::shared_ptr(std::move(writer)); + } + return RealtimePrimaryKeyWriter::Create( + partition_map, bucket, schema_, trimmed_primary_keys, key_comparator_, + realtime_context_impl, realtime_store_state.value(), restore_max_seq_number, writer, pool_); +} + +Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + if (!realtime_context_) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), + options_.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context_)); + return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets); } Status KeyValueFileStoreWrite::Close() { diff --git a/src/paimon/core/operation/key_value_file_store_write.h b/src/paimon/core/operation/key_value_file_store_write.h index 14457590..66c362f2 100644 --- a/src/paimon/core/operation/key_value_file_store_write.h +++ b/src/paimon/core/operation/key_value_file_store_write.h @@ -45,6 +45,7 @@ class SnapshotManager; class SchemaManager; class TableSchema; class IOManager; +class RealtimeContext; struct KeyValue; template class MergeFunctionWrapper; @@ -65,8 +66,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, bool enable_multi_thread_spill, + const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool); + Status RefreshCommittedSnapshot(int64_t snapshot_id) override; Status Close() override; private: @@ -79,8 +82,13 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { Result> CreateFileStoreScan( const std::shared_ptr& filter) const override; + bool IsRealtimeWrite() const override { + return realtime_context_ != nullptr; + } + private: bool enable_multi_thread_spill_; + std::shared_ptr realtime_context_; std::shared_ptr key_comparator_; std::shared_ptr user_defined_seq_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index 35d938af..a2344e80 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -19,9 +19,13 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include +#include #include #include +#include +#include #include +#include #include #include @@ -44,6 +48,7 @@ #include "paimon/common/utils/path_util.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/operation/restore_files.h" +#include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/file_store_commit.h" @@ -52,7 +57,9 @@ #include "paimon/format/file_format_factory.h" #include "paimon/format/reader_builder.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" #include "paimon/testing/utils/test_helper.h" @@ -60,6 +67,50 @@ #include "paimon/write_context.h" namespace paimon::test { +namespace { + +class TestingMemoryPool final : public MemoryPool { + public: + void* Malloc(uint64_t size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* pointer, size_t old_size, size_t new_size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Realloc(pointer, old_size, new_size, alignment); + } + + void Free(void* pointer, uint64_t size) override { + delegate_->Free(pointer, size); + } + + void Free(void* pointer, uint64_t size, uint64_t alignment) override { + delegate_->Free(pointer, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + bool reject_allocations = false; + int64_t allocation_count = 0; + + private: + std::unique_ptr delegate_ = GetMemoryPool(); +}; + +} class KeyValueFileStoreWriteTest : public ::testing::Test { protected: @@ -127,14 +178,15 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } - std::unique_ptr MakeBatch(const std::shared_ptr& schema, - const std::string& json) const { + std::unique_ptr MakeBatch( + const std::shared_ptr& schema, const std::string& json, + const std::vector& row_kinds = {}) const { auto struct_type = arrow::struct_(schema->fields()); auto array = arrow::ipc::internal::json::ArrayFromJSON(struct_type, json).ValueOrDie(); ::ArrowArray arrow_array; EXPECT_TRUE(arrow::ExportArray(*array, &arrow_array).ok()); RecordBatchBuilder batch_builder(&arrow_array); - return batch_builder.SetBucket(0).Finish().value(); + return batch_builder.SetRowKinds(row_kinds).SetBucket(0).Finish().value(); } std::vector> WriteAndPrepare( @@ -193,6 +245,58 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { EXPECT_NE(nullptr, metadata); return MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy()).value(); } + + Result>> + ReadPreparedRows(const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + context->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected exactly one real-time store"); + } + RealtimeQueryContext query_context{nullptr, nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + views[0].store->CreateQueryReaders( + views[0].read_view, 0, query_context)); + std::vector> rows; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(array); + if (!values || values->num_fields() != 5) { + return Status::Invalid("unexpected prepared real-time batch"); + } + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(values->field(0)); + std::shared_ptr sequences = + std::dynamic_pointer_cast(values->field(1)); + std::shared_ptr offsets = + std::dynamic_pointer_cast(values->field(2)); + std::shared_ptr ids = + std::dynamic_pointer_cast(values->field(3)); + std::shared_ptr payloads = + std::dynamic_pointer_cast(values->field(4)); + if (!row_kinds || !sequences || !offsets || !ids || !payloads) { + return Status::Invalid("unexpected prepared real-time column type"); + } + for (int64_t row = 0; row < values->length(); ++row) { + rows.emplace_back(row_kinds->Value(row), ids->Value(row), + payloads->GetString(row), sequences->Value(row), + offsets->Value(row)); + } + } + reader->Close(); + } + return rows; + } }; TEST_F(KeyValueFileStoreWriteTest, TestWriteWithInvalidBatch) { @@ -303,6 +407,216 @@ TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenLookupEnabl ASSERT_EQ(commit_messages.size(), 1); } +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { + const std::map options = { + {Options::BUCKET, "1"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + {Options::REALTIME_ENABLED, "true"}, + }; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithTempDirectory(dir->Str()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + std::unique_ptr batch = + MakeBatch(schema, R"([ + [1, "old"], + [2, "two"], + [1, "new"] + ])", + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER}); + ASSERT_OK(writer->Write(std::move(batch))); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ( + (std::vector{{0, 1, "old", 0, 0}, {2, 1, "new", 2, 2}, {3, 2, "two", 1, 1}}), + prepared_rows); + ASSERT_OK_AND_ASSIGN(std::vector progresses, + writer->PrepareCommitWithProgress(0)); + ASSERT_EQ(1, progresses.size()); + ASSERT_EQ(OffsetRange(0, 3), progresses[0].offset_range); + std::shared_ptr commit_message = + std::dynamic_pointer_cast(progresses[0].commit_message); + ASSERT_NE(nullptr, commit_message); + int64_t row_count = 0; + for (const std::shared_ptr& file : + commit_message->GetNewFilesIncrement().NewFiles()) { + row_count += file->row_count; + } + ASSERT_EQ(2, row_count); + ASSERT_EQ(0, TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str())); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeOffsetCollision) { + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("_REALTIME_OFFSET", arrow::int64()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(PathUtil::JoinPath(dir->Str(), "foo.db/bar"), "test"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[1, 10]])")), + "PK real-time write schema contains reserved transport field"); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + std::shared_ptr pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + const int64_t allocations_before_write = pool->allocation_count; + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "one"]])"))); + ASSERT_GT(pool->allocation_count, allocations_before_write); + ASSERT_OK(writer->Close()); + writer.reset(); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector retained_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "one", 0, 0}}), retained_rows); + + std::shared_ptr rejecting_pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rejecting_context, + RealtimeContext::Create()); + WriteContextBuilder rejecting_builder(table_path, "rejecting"); + rejecting_builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(rejecting_context) + .WithMemoryPool(rejecting_pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_write_context, + rejecting_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_writer, + FileStoreWrite::Create(std::move(rejecting_write_context))); + ASSERT_OK(rejecting_writer->Write(MakeBatch(schema, "[]"))); + const int64_t rejecting_allocations_before_write = rejecting_pool->allocation_count; + rejecting_pool->reject_allocations = true; + ASSERT_NOK_WITH_MSG(rejecting_writer->Write(MakeBatch(schema, R"([[2, "two"]])")), + "Out of memory"); + ASSERT_GT(rejecting_pool->allocation_count, rejecting_allocations_before_write); + ASSERT_OK_AND_ASSIGN(std::vector rejected_rows, + ReadPreparedRows(rejecting_context)); + ASSERT_TRUE(rejected_rows.empty()); + ASSERT_OK(rejecting_writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { + const int64_t max = std::numeric_limits::max(); + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr initial_context, + RealtimeContext::Create()); + WriteContextBuilder initial_builder(table_path, "initial"); + initial_builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext( + initial_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_write_context, + initial_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_writer, + FileStoreWrite::Create(std::move(initial_write_context))); + ASSERT_OK(initial_writer->Write(MakeBatch(schema, R"([[0, "initial"]])"))); + ASSERT_OK_AND_ASSIGN(std::vector initial_progress, + initial_writer->PrepareCommitWithProgress(0)); + ASSERT_EQ(1, initial_progress.size()); + std::shared_ptr initial_message = + std::dynamic_pointer_cast(initial_progress[0].commit_message); + ASSERT_NE(nullptr, initial_message); + ASSERT_EQ(1, initial_message->GetNewFilesIncrement().NewFiles().size()); + initial_message->GetNewFilesIncrement().NewFiles()[0]->AssignSequenceNumber(max - 2, max - 2); + initial_progress[0].offset_range = OffsetRange(0, max - 1); + + CommitContextBuilder commit_builder(table_path, "initial"); + commit_builder.SetOptions(options); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, commit_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committer, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + committer->CommitWithProgress(initial_progress, 0, std::nullopt)); + ASSERT_OK(initial_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "boundary"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "legal"]])"))); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[2, "overflow"]])")), + "real-time offset range exceeds INT64_MAX"); + ASSERT_OK_AND_ASSIGN(prepared_rows, ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context_impl, + RealtimeContextImpl::Cast(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::vector views, + context_impl->AcquireReadViews()); + ASSERT_EQ(1, views.size()); + ASSERT_EQ(std::optional(OffsetRange(max - 1, max)), + views[0].read_view->GetOffsetRange()); + ASSERT_OK(writer->Close()); + ASSERT_GE(snapshot_id, 1); +} + TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenDefaultCompactRewriterPathEnabled) { ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index b753ea43..c85e75ee 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -78,6 +78,126 @@ struct KeyValue; template class MergeFunctionWrapper; +class MergeFileSplitRead::RealtimeReaderBuilder { + public: + static Result> Create( + MergeFileSplitRead* owner, const std::vector>& disk_splits, + std::vector>&& additional_readers) { + RealtimeReaderBuilder builder(owner); + std::vector> readers; + if (!disk_splits.empty()) { + PAIMON_RETURN_NOT_OK(builder.CollectDiskReaders(disk_splits, &readers)); + } + readers.reserve(readers.size() + additional_readers.size()); + for (std::unique_ptr& additional_reader : additional_readers) { + readers.push_back(std::move(additional_reader)); + } + return builder.CreateMergedReader(std::move(readers)); + } + + private: + explicit RealtimeReaderBuilder(MergeFileSplitRead* owner) : owner_(owner) {} + + Status CollectDiskReaders(const std::vector>& disk_splits, + std::vector>* readers) { + std::shared_ptr first_split = + std::dynamic_pointer_cast(disk_splits.front()); + if (!first_split) { + return Status::Invalid("merge input disk split is not a data split"); + } + const BinaryRow& partition = first_split->Partition(); + const int32_t bucket = first_split->Bucket(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); + + std::vector> data_files; + std::vector> deletion_files; + for (const std::shared_ptr& disk_split : disk_splits) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + if (!data_split || !(data_split->Partition() == partition) || + data_split->Bucket() != bucket) { + return Status::Invalid("merge input disk splits do not share a partition-bucket"); + } + if (!data_split->BeforeFiles().empty() || data_split->IsStreaming() || + data_split->Bucket() == BucketModeDefine::POSTPONE_BUCKET) { + return Status::Invalid("additional merge input requires fixed-bucket batch splits"); + } + const std::vector>& split_files = data_split->DataFiles(); + const std::vector>& split_deletion_files = + data_split->DeletionFiles(); + if (!split_deletion_files.empty() && + split_deletion_files.size() != split_files.size()) { + return Status::Invalid( + "merge input disk split deletion files must be empty or match data files"); + } + data_files.insert(data_files.end(), split_files.begin(), split_files.end()); + if (split_deletion_files.empty()) { + deletion_files.insert(deletion_files.end(), split_files.size(), std::nullopt); + } else { + deletion_files.insert(deletion_files.end(), split_deletion_files.begin(), + split_deletion_files.end()); + } + } + + DeletionVector::Factory dv_factory = DeletionVector::CreateFactory( + owner_->options_.GetFileSystem(), + DeletionVector::CreateDeletionFileMap(data_files, deletion_files), owner_->pool_); + std::vector> disk_sections = + IntervalPartition(data_files, owner_->key_comparator_).Partition(); + for (const std::vector& section : disk_sections) { + for (const SortedRun& run : section) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + owner_->CreateReaderForRun(partition, run, dv_factory, + owner_->predicate_for_keys_, + data_file_path_factory)); + readers->push_back(std::move(disk_reader)); + } + } + return Status::OK(); + } + + Result> CreateMergedReader( + std::vector>&& record_readers) { + if (record_readers.empty()) { + return std::make_unique(std::vector>{}, + owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + owner_->CreateSortMergeReader(std::move(record_readers))); + return CreateProjectedReader(std::move(sort_merge_reader)); + } + + Result> CreateProjectedReader( + std::unique_ptr&& sort_merge_reader) { + if (!owner_->force_keep_delete_) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } + + std::unique_ptr projection_reader; + if (!owner_->context_->EnableMultiThreadRowToBatch()) { + PAIMON_ASSIGN_OR_RAISE( + projection_reader, + KeyValueProjectionReader::Create( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), owner_->pool_)); + } else { + const int32_t thread_number = owner_->context_->GetRowToBatchThreadNumber(); + assert(thread_number > 0); + projection_reader = std::make_unique( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), thread_number, owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(projection_reader, + owner_->ApplyPredicateFilterIfNeeded( + std::move(projection_reader), owner_->context_->GetPredicate())); + return std::make_unique(std::move(projection_reader), + owner_->pool_); + } + + MergeFileSplitRead* owner_; +}; + Result> MergeFileSplitRead::Create( const std::shared_ptr& path_factory, const std::shared_ptr& context, @@ -158,6 +278,12 @@ Result> MergeFileSplitRead::CreateReader( return std::make_unique(std::move(batch_reader), pool_); } +Result> MergeFileSplitRead::CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector>&& additional_readers) { + return RealtimeReaderBuilder::Create(this, disk_splits, std::move(additional_readers)); +} + void MergeFileSplitRead::SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper) { merge_function_wrapper_ = merge_function_wrapper; diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index d4bfa727..3cc63a44 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -117,10 +117,20 @@ class MergeFileSplitRead : public AbstractSplitRead { return value_schema_; } + std::shared_ptr GetKeySchema() const { + return key_schema_; + } + + Result> CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector>&& additional_readers); + void SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper); private: + class RealtimeReaderBuilder; + Result> CreateMergeReader( const std::shared_ptr& data_split, const std::shared_ptr& data_file_path_factory); diff --git a/src/paimon/core/operation/orphan_files_cleaner_impl.cpp b/src/paimon/core/operation/orphan_files_cleaner_impl.cpp index 6b8e5ba0..015d1340 100644 --- a/src/paimon/core/operation/orphan_files_cleaner_impl.cpp +++ b/src/paimon/core/operation/orphan_files_cleaner_impl.cpp @@ -35,6 +35,7 @@ #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/manifest/manifest_list.h" +#include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/operation/metrics/clean_metrics.h" #include "paimon/core/snapshot.h" #include "paimon/core/utils/duration.h" @@ -85,7 +86,7 @@ bool OrphanFilesCleanerImpl::SupportToClean(const std::string& file_name) { return true; } } - return false; + return StringUtils::EndsWith(file_name, ".offsets"); } Result> OrphanFilesCleanerImpl::Clean() { @@ -156,6 +157,10 @@ Result> OrphanFilesCleanerImpl::ListPaimonFileDirs() const std::set paimon_file_dirs; paimon_file_dirs.insert(snapshot_manager_->SnapshotDirectory()); paimon_file_dirs.insert(FileStorePathFactory::ManifestPath(root_path_)); + if (options_.RealtimeEnabled()) { + paimon_file_dirs.insert( + RealtimeCommitProperties::OffsetsDirectory(root_path_, options_.GetBranch())); + } // TODO(jinli.zjw): support clean index, stats, changelog in the future // paimon_file_dirs.insert(FileStorePathFactory::IndexPath(root_path_)); // paimon_file_dirs.insert(FileStorePathFactory::StatisticsPath(root_path_)); @@ -294,6 +299,13 @@ Result> OrphanFilesCleanerImpl::GetUsedFilesBySnapshot( used_files.insert(SnapshotManager::SNAPSHOT_PREFIX + std::to_string(snapshot.Id())); used_files.insert(snapshot.BaseManifestList()); used_files.insert(snapshot.DeltaManifestList()); + if (options_.RealtimeEnabled()) { + std::optional offsets_path = + RealtimeCommitProperties::GetOffsetsPath(snapshot); + if (offsets_path) { + used_files.insert(PathUtil::GetName(offsets_path.value())); + } + } std::vector manifests; PAIMON_RETURN_NOT_OK(manifest_list_->ReadIfFileExist(snapshot.BaseManifestList(), /*filter=*/nullptr, &manifests)); diff --git a/src/paimon/core/operation/raw_file_split_read.h b/src/paimon/core/operation/raw_file_split_read.h index ac211b25..6a97b9b3 100644 --- a/src/paimon/core/operation/raw_file_split_read.h +++ b/src/paimon/core/operation/raw_file_split_read.h @@ -54,8 +54,8 @@ struct DeletionFile; /// splits)->CompleteRowKindBatchReader->(PredicateBatchReader) /// ->ConcatBatchReader across /// files->FieldMappingReader->(ApplyBitmapIndexBatchReader)->(CompleteRowTrackingFieldsBatchReader) -/// ->(ShreddingFileReader)->(MapSharedShreddingFileReader)->(DelegatingPrefetchReader) -/// ->(PrefetchFileBatchReader)->FormatReader +/// ->(ShreddingFileReader)->(MapSharedShreddingFileReader)->(VectorFileBatchReader) +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader class RawFileSplitRead : public AbstractSplitRead { public: diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index eb3d8826..08a854d8 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -42,7 +42,7 @@ ReadContext::ReadContext( const std::shared_ptr& specific_file_system, const std::map& fs_scheme_to_identifier_map, const std::shared_ptr& realtime_context, - const std::map& options, PrefetchCacheMode prefetch_cache_mode, + const std::map& options, bool read_ahead_cache_enabled, const CacheConfig& cache_config, const std::shared_ptr& cache) : path_(path), branch_(branch), @@ -62,7 +62,7 @@ ReadContext::ReadContext( fs_scheme_to_identifier_map_(fs_scheme_to_identifier_map), realtime_context_(realtime_context), options_(options), - prefetch_cache_mode_(prefetch_cache_mode), + read_ahead_cache_enabled_(read_ahead_cache_enabled), cache_config_(cache_config), cache_(cache) {} @@ -97,7 +97,7 @@ class ReadContextBuilder::Impl { predicate_.reset(); enable_predicate_filter_ = false; enable_prefetch_ = false; - prefetch_cache_mode_ = PrefetchCacheMode::ALWAYS; + read_ahead_cache_enabled_ = true; prefetch_batch_count_ = 600; prefetch_max_parallel_num_ = 3; enable_multi_thread_row_to_batch_ = false; @@ -131,7 +131,7 @@ class ReadContextBuilder::Impl { std::shared_ptr executor_; std::shared_ptr specific_file_system_; std::shared_ptr realtime_context_; - PrefetchCacheMode prefetch_cache_mode_ = PrefetchCacheMode::ALWAYS; + bool read_ahead_cache_enabled_ = true; CacheConfig cache_config_; std::shared_ptr cache_; }; @@ -250,8 +250,8 @@ ReadContextBuilder& ReadContextBuilder::WithFileSystem( return *this; } -ReadContextBuilder& ReadContextBuilder::SetPrefetchCacheMode(PrefetchCacheMode mode) { - impl_->prefetch_cache_mode_ = mode; +ReadContextBuilder& ReadContextBuilder::SetReadAheadCacheEnabled(bool enabled) { + impl_->read_ahead_cache_enabled_ = enabled; return *this; } @@ -301,7 +301,7 @@ Result> ReadContextBuilder::Finish() { impl_->enable_multi_thread_row_to_batch_, impl_->row_to_batch_thread_number_, impl_->table_schema_, impl_->memory_pool_, impl_->executor_, impl_->specific_file_system_, impl_->fs_scheme_to_identifier_map_, impl_->realtime_context_, impl_->options_, - impl_->prefetch_cache_mode_, impl_->cache_config_, impl_->cache_); + impl_->read_ahead_cache_enabled_, impl_->cache_config_, impl_->cache_); if (impl_->read_schema_ && impl_->read_schema_->release) { ctx->SetReadSchema(std::move(impl_->read_schema_)); } diff --git a/src/paimon/core/operation/read_context_test.cpp b/src/paimon/core/operation/read_context_test.cpp index 8d5a78e1..c686cccb 100644 --- a/src/paimon/core/operation/read_context_test.cpp +++ b/src/paimon/core/operation/read_context_test.cpp @@ -45,7 +45,7 @@ TEST(ReadContextTest, TestDefaultValue) { ASSERT_FALSE(ctx->GetPredicate()); ASSERT_FALSE(ctx->EnablePredicateFilter()); ASSERT_FALSE(ctx->EnablePrefetch()); - ASSERT_EQ(PrefetchCacheMode::ALWAYS, ctx->GetPrefetchCacheMode()); + ASSERT_TRUE(ctx->ReadAheadCacheEnabled()); ASSERT_EQ(600, ctx->GetPrefetchBatchCount()); ASSERT_EQ(3, ctx->GetPrefetchMaxParallelNum()); ASSERT_FALSE(ctx->EnableMultiThreadRowToBatch()); @@ -59,8 +59,8 @@ TEST(ReadContextTest, TestSetContent) { ReadContextBuilder builder("table_root_path"); std::shared_ptr memory_pool = GetDefaultPool(); std::shared_ptr executor = CreateDefaultExecutor(); - CacheConfig cache_config(/*buffer_size_limit=*/1024, /*range_size_limit=*/512, - /*hole_size_limit=*/128, /*pre_buffer_limit=*/2048); + CacheConfig cache_config(/*range_size_limit=*/512, /*hole_size_limit=*/128, + /*pre_buffer_limit=*/2048); builder.AddOption("key", "value"); builder.SetReadFieldNames({"f1", "f2"}); @@ -70,7 +70,7 @@ TEST(ReadContextTest, TestSetContent) { builder.SetPredicate(predicate); builder.EnablePredicateFilter(true); builder.EnablePrefetch(true); - builder.SetPrefetchCacheMode(PrefetchCacheMode::NEVER); + builder.SetReadAheadCacheEnabled(false); builder.SetPrefetchBatchCount(1200); builder.SetPrefetchMaxParallelNum(6); builder.EnableMultiThreadRowToBatch(true); @@ -95,7 +95,7 @@ TEST(ReadContextTest, TestSetContent) { ASSERT_EQ(*predicate, *(ctx->GetPredicate())); ASSERT_TRUE(ctx->EnablePredicateFilter()); ASSERT_TRUE(ctx->EnablePrefetch()); - ASSERT_EQ(PrefetchCacheMode::NEVER, ctx->GetPrefetchCacheMode()); + ASSERT_FALSE(ctx->ReadAheadCacheEnabled()); ASSERT_EQ(1200, ctx->GetPrefetchBatchCount()); ASSERT_EQ(6, ctx->GetPrefetchMaxParallelNum()); ASSERT_TRUE(ctx->EnableMultiThreadRowToBatch()); @@ -105,7 +105,6 @@ TEST(ReadContextTest, TestSetContent) { ASSERT_TRUE(ctx->GetSpecificTableSchema().has_value()); ASSERT_EQ("table-schema-json", ctx->GetSpecificTableSchema().value()); ASSERT_EQ("rt", ctx->GetBranch()); - ASSERT_EQ(1024U, ctx->GetCacheConfig().GetBufferSizeLimit()); ASSERT_EQ(512U, ctx->GetCacheConfig().GetRangeSizeLimit()); ASSERT_EQ(128U, ctx->GetCacheConfig().GetHoleSizeLimit()); ASSERT_EQ(2048U, ctx->GetCacheConfig().GetPreBufferLimit()); diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index cf1e37aa..18087a29 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -25,12 +25,18 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/compute/api_aggregate.h" +#include "paimon/common/data/columnar/columnar_array.h" +#include "paimon/common/data/columnar/columnar_row.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/predicate/predicate_filter.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/projected_array.h" +#include "paimon/common/utils/projected_row.h" #include "paimon/core/utils/nested_projection_utils.h" #include "paimon/macros.h" @@ -53,6 +59,26 @@ uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { return result; } +bool SupportsMinMax(const std::shared_ptr& type) { + switch (type->id()) { + case arrow::Type::BOOL: + case arrow::Type::INT8: + case arrow::Type::INT16: + case arrow::Type::INT32: + case arrow::Type::INT64: + case arrow::Type::FLOAT: + case arrow::Type::DOUBLE: + case arrow::Type::STRING: + case arrow::Type::BINARY: + case arrow::Type::DATE32: + case arrow::Type::TIMESTAMP: + case arrow::Type::DECIMAL128: + return true; + default: + return false; + } +} + } // namespace class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { @@ -165,11 +191,17 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { public: QueryBatchReader(const ReadView* view, int64_t offset_begin, const std::shared_ptr& read_schema, - const std::shared_ptr& arrow_pool) + const std::shared_ptr& predicate_filter, + std::vector&& statistics_mapping, + const std::shared_ptr& arrow_pool, + const std::shared_ptr& memory_pool) : view_(view), offset_begin_(offset_begin), read_schema_(read_schema), arrow_pool_(arrow_pool), + memory_pool_(memory_pool), + predicate_filter_(predicate_filter), + statistics_mapping_(std::move(statistics_mapping)), metrics_(std::make_shared()) {} Result NextBatch() override { @@ -189,6 +221,10 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { if (stored.offset_range.end <= offset_begin_) { continue; } + PAIMON_ASSIGN_OR_RAISE(bool may_match, MayMatch(stored)); + if (!may_match) { + continue; + } int64_t begin = std::max(0, offset_begin_ - stored.offset_range.begin); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, BuildOutput(stored)); RoaringBitmap32 candidate_rows; @@ -213,6 +249,25 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { } private: + Result MayMatch(const StoredBatch& stored) const { + if (!predicate_filter_ || !stored.statistics) { + return true; + } + const BatchStatistics& statistics = stored.statistics.value(); + std::shared_ptr min_row = std::make_shared( + statistics.min_values, statistics.min_values->fields(), memory_pool_, /*row_id=*/0); + std::shared_ptr max_row = std::make_shared( + statistics.max_values, statistics.max_values->fields(), memory_pool_, /*row_id=*/0); + ProjectedRow projected_min(min_row, statistics_mapping_); + ProjectedRow projected_max(max_row, statistics_mapping_); + std::shared_ptr null_counts = + std::make_shared(statistics.null_counts.get(), memory_pool_, + /*offset=*/0, statistics.null_counts->length()); + ProjectedArray projected_null_counts(null_counts, statistics_mapping_); + return predicate_filter_->Test(read_schema_, stored.data->length(), projected_min, + projected_max, projected_null_counts); + } + Result> BuildOutput(const StoredBatch& stored) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr projected, @@ -231,14 +286,81 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { int64_t offset_begin_; std::shared_ptr read_schema_; std::shared_ptr arrow_pool_; + std::shared_ptr memory_pool_; + std::shared_ptr predicate_filter_; + std::vector statistics_mapping_; std::shared_ptr metrics_; size_t next_batch_ = 0; }; ArrowRealtimeStore::ArrowRealtimeStore(const std::shared_ptr& write_schema, + StatisticsMode statistics_mode, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool) - : write_schema_(write_schema), memory_pool_(memory_pool), arrow_pool_(arrow_pool) {} + : write_schema_(write_schema), + memory_pool_(memory_pool), + arrow_pool_(arrow_pool), + statistics_mode_(statistics_mode) {} + +Result> ArrowRealtimeStore::CollectStatistics( + const std::shared_ptr& data) const { + if (statistics_mode_ == StatisticsMode::NONE) { + return std::optional(); + } + + arrow::ArrayVector min_values; + arrow::ArrayVector max_values; + min_values.reserve(data->num_fields()); + max_values.reserve(data->num_fields()); + arrow::Int64Builder null_count_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(null_count_builder.Reserve(data->num_fields())); + arrow::compute::ScalarAggregateOptions aggregate_options; + aggregate_options.skip_nulls = true; + aggregate_options.min_count = 1; + arrow::compute::ExecContext exec_context(arrow_pool_.get()); + + for (const std::shared_ptr& field : data->fields()) { + null_count_builder.UnsafeAppend(field->null_count()); + if (!SupportsMinMax(field->type())) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr min_value, + arrow::MakeArrayOfNull(field->type(), /*length=*/1, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr max_value, + arrow::MakeArrayOfNull(field->type(), /*length=*/1, arrow_pool_.get())); + min_values.push_back(std::move(min_value)); + max_values.push_back(std::move(max_value)); + continue; + } + + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum min_max, arrow::compute::MinMax(field, aggregate_options, &exec_context)); + std::shared_ptr min_max_scalar = + std::dynamic_pointer_cast(min_max.scalar()); + if (!min_max_scalar || min_max_scalar->value.size() != 2) { + return Status::Invalid("Arrow min_max did not produce min and max scalars"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr min_value, + arrow::MakeArrayFromScalar(*min_max_scalar->value[0], /*length=*/1, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr max_value, + arrow::MakeArrayFromScalar(*min_max_scalar->value[1], /*length=*/1, arrow_pool_.get())); + min_values.push_back(std::move(min_value)); + max_values.push_back(std::move(max_value)); + } + + std::shared_ptr null_counts_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(null_count_builder.Finish(&null_counts_array)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr min_values_struct, + arrow::StructArray::Make(min_values, write_schema_->fields())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr max_values_struct, + arrow::StructArray::Make(max_values, write_schema_->fields())); + return std::optional(BatchStatistics{ + std::move(min_values_struct), std::move(max_values_struct), std::move(null_counts_array)}); +} Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { if (!write_batch.batch) { @@ -264,16 +386,23 @@ Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { } std::shared_ptr struct_array = checked_pointer_cast(data); + PAIMON_ASSIGN_OR_RAISE(std::optional statistics, + CollectStatistics(struct_array)); std::lock_guard lock(mutex_); if (building_range_ && write_batch.offset_range.begin != building_range_->end) { return Status::Invalid("real-time offset ranges must be contiguous"); } uint64_t memory_usage = GetArrayMemoryUsage(struct_array->data()); + if (statistics) { + memory_usage += GetArrayMemoryUsage(statistics->min_values->data()) + + GetArrayMemoryUsage(statistics->max_values->data()) + + GetArrayMemoryUsage(statistics->null_counts->data()); + } building_memory_usage_ += memory_usage; - building_batches_.push_back(StoredBatch{std::move(struct_array), - write_batch.batch->GetRowKind(), - write_batch.offset_range, memory_usage}); + building_batches_.push_back( + StoredBatch{std::move(struct_array), write_batch.batch->GetRowKind(), + write_batch.offset_range, std::move(statistics), memory_usage}); if (!building_range_) { building_range_ = write_batch.offset_range; } else { @@ -329,13 +458,20 @@ Result>> ArrowRealtimeStore::CreateQuer } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, arrow::ImportSchema(context.read_schema)); - // TODO(xinyu.lxy): Support predicate pushdown after adding batch statistics or index metadata. - // The default Arrow store currently ignores context.predicate and - // context.enable_predicate_pushdown, and returns all offset-matching rows as candidates. + std::shared_ptr predicate_filter; + if (context.enable_predicate_pushdown && context.predicate) { + predicate_filter = std::dynamic_pointer_cast(context.predicate); + } + std::vector statistics_mapping; + statistics_mapping.reserve(read_schema->num_fields()); + for (const std::shared_ptr& field : read_schema->fields()) { + statistics_mapping.push_back(write_schema_->GetFieldIndex(field->name())); + } std::vector> readers; if (arrow_view->GetOffsetRange() && arrow_view->GetOffsetRange()->end > offset_begin) { std::unique_ptr reader = std::make_unique( - arrow_view.get(), offset_begin, read_schema, arrow_pool_); + arrow_view.get(), offset_begin, read_schema, predicate_filter, + std::move(statistics_mapping), arrow_pool_, memory_pool_); reader = std::make_unique(std::move(reader), memory_pool_); readers.push_back(std::move(reader)); } diff --git a/src/paimon/core/realtime/arrow_realtime_store.h b/src/paimon/core/realtime/arrow_realtime_store.h index bd9ecfbf..97339852 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.h +++ b/src/paimon/core/realtime/arrow_realtime_store.h @@ -28,6 +28,7 @@ #include "paimon/realtime/realtime_store.h" namespace arrow { +class Array; class MemoryPool; class Schema; class StructArray; @@ -40,6 +41,7 @@ class MemoryPool; class ArrowRealtimeStore : public RealtimeStore { public: ArrowRealtimeStore(const std::shared_ptr& write_schema, + StatisticsMode statistics_mode, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool); @@ -61,10 +63,17 @@ class ArrowRealtimeStore : public RealtimeStore { uint64_t GetMemoryUsage() const override; private: + struct BatchStatistics { + std::shared_ptr min_values; + std::shared_ptr max_values; + std::shared_ptr null_counts; + }; + struct StoredBatch { std::shared_ptr data; std::vector row_kinds; OffsetRange offset_range; + std::optional statistics; uint64_t memory_usage; }; @@ -73,9 +82,13 @@ class ArrowRealtimeStore : public RealtimeStore { class CommitBatchReader; class QueryBatchReader; + Result> CollectStatistics( + const std::shared_ptr& data) const; + std::shared_ptr write_schema_; std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; + StatisticsMode statistics_mode_; mutable std::mutex mutex_; std::vector building_batches_; std::vector> sealed_segments_; diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 0eecb5d0..d0d4ae70 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -25,24 +25,36 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" namespace paimon { Result> ArrowRealtimeStoreFactory::Create( - std::unique_ptr write_schema, const std::map&, - const std::shared_ptr& memory_pool) { - if (!write_schema || !write_schema->release) { + RealtimeStoreCreateRequest&& request) { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("real-time store write schema is null"); } - ScopeGuard schema_guard([schema = write_schema.get()]() { ArrowSchemaRelease(schema); }); - if (!memory_pool) { + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + if (!request.memory_pool) { return Status::Invalid("real-time store memory pool is null"); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, - arrow::ImportSchema(write_schema.get())); - std::shared_ptr arrow_pool = GetArrowPool(memory_pool); - return std::make_shared(imported_schema, memory_pool, arrow_pool); + arrow::ImportSchema(request.write_schema.get())); + switch (request.mode) { + case RealtimeStoreMode::APPEND_ONLY: { + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + return std::make_shared(imported_schema, request.statistics_mode, + request.memory_pool, arrow_pool); + } + case RealtimeStoreMode::PRIMARY_KEY: { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(imported_schema)); + return std::shared_ptr(std::move(store)); + } + } + return Status::Invalid("invalid real-time store mode: ", static_cast(request.mode)); } } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index cc08cc8e..0a3e5235 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -19,6 +19,7 @@ #include "paimon/core/realtime/arrow_realtime_store.h" +#include #include #include #include @@ -28,7 +29,11 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/record_batch.h" #include "paimon/testing/utils/testharness.h" @@ -56,7 +61,11 @@ class ArrowRealtimeStoreTest : public testing::Test { {arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); pool_ = GetDefaultPool(); arrow_pool_ = GetArrowPool(pool_); - store_ = std::make_shared(schema_, pool_, arrow_pool_); + store_ = CreateStore(StatisticsMode::NONE); + } + + std::shared_ptr CreateStore(StatisticsMode statistics_mode) const { + return std::make_shared(schema_, statistics_mode, pool_, arrow_pool_); } std::unique_ptr MakeBatch(const std::string& json) const { @@ -86,6 +95,21 @@ class ArrowRealtimeStoreTest : public testing::Test { return c_schema; } + std::vector ReadIds(const BatchReader::ReadBatchWithBitmap& batch) const { + std::shared_ptr array = + arrow::ImportArray(batch.first.first.get(), batch.first.second.get()).ValueOrDie(); + std::shared_ptr struct_array = + checked_pointer_cast(array); + std::shared_ptr ids = + checked_pointer_cast(struct_array->field(/*pos=*/1)); + std::vector result; + for (RoaringBitmap32::Iterator iter = batch.second.Begin(); iter != batch.second.End(); + ++iter) { + result.push_back(ids->Value(*iter)); + } + return result; + } + protected: std::shared_ptr schema_; std::shared_ptr pool_; @@ -205,6 +229,85 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { << "expected: " << expected_array->ToString() << ", actual: " << actual_array->ToString(); } +TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { + ArrowRealtimeStoreFactory factory; + std::unique_ptr write_schema = MakeReadSchema(schema_); + RealtimeStoreCreateRequest request{std::move(write_schema), + /*options=*/{}, + pool_, + /*partition=*/{}, + /*bucket=*/0, + RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::FULL}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_store, + factory.Create(std::move(request))); + std::shared_ptr store = + std::dynamic_pointer_cast(realtime_store); + ASSERT_NE(nullptr, store); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, "a"], [1, "b"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[10, "c"], [11, "d"]])"), OffsetRange(2, 4)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + std::unique_ptr read_schema = MakeReadSchema(schema_); + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, Literal(int64_t{5})); + RealtimeQueryContext context{read_schema.get(), predicate, /*enable_predicate_pushdown=*/true}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, readers[0]->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(std::vector({10, 11}), ReadIds(batch)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap eof, readers[0]->NextBatchWithBitmap()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + + std::unique_ptr unfiltered_read_schema = MakeReadSchema(schema_); + RealtimeQueryContext unfiltered_context{unfiltered_read_schema.get(), predicate, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> unfiltered_readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, unfiltered_context)); + ASSERT_EQ(1, unfiltered_readers.size()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap unfiltered_batch, + unfiltered_readers[0]->NextBatchWithBitmap()); + ASSERT_EQ(std::vector({0, 1}), ReadIds(unfiltered_batch)); +} + +TEST_F(ArrowRealtimeStoreTest, TestFactoryRejectsInvalidMode) { + ArrowRealtimeStoreFactory factory; + std::unique_ptr write_schema = MakeReadSchema(schema_); + RealtimeStoreCreateRequest request{std::move(write_schema), + /*options=*/{}, + pool_, + /*partition=*/{}, + /*bucket=*/0, + static_cast(-1)}; + ASSERT_NOK_WITH_MSG(factory.Create(std::move(request)), "invalid real-time store mode: -1"); +} + +TEST_F(ArrowRealtimeStoreTest, TestMissingStatisticsRetainsNonMatchingBatch) { + ASSERT_OK( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[0, "a"], [1, "b"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store_->Write( + RealtimeWriteBatch{MakeBatch(R"([[10, "c"], [11, "d"]])"), OffsetRange(2, 4)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + + std::unique_ptr read_schema = MakeReadSchema(schema_); + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, Literal(int64_t{5})); + RealtimeQueryContext context{read_schema.get(), predicate, + /*enable_predicate_pushdown=*/true}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, readers[0]->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(std::vector({0, 1}), ReadIds(batch)); +} + TEST_F(ArrowRealtimeStoreTest, TestRejectsHandlesFromAnotherStoreImplementation) { ASSERT_NOK_WITH_MSG(store_->CreateCommitReaders(std::make_shared()), "segment was not created by the Arrow real-time store"); diff --git a/src/paimon/core/realtime/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp new file mode 100644 index 00000000..a34ccea4 --- /dev/null +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -0,0 +1,698 @@ +/* + * 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 "paimon/core/realtime/prepared_key_value_reader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_nested.h" +#include "arrow/array/array_primitive.h" +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/status.h" + +namespace paimon { + +namespace { + +constexpr int32_t kValueKindIndex = 0; +constexpr int32_t kSequenceNumberIndex = 1; +constexpr int32_t kRealtimeOffsetIndex = 2; +constexpr int32_t kPreparedValueStartIndex = 3; + +template +void CloseReaders(const std::vector>& readers) { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool); + +class RealtimeOffsetCoverage { + public: + static Result> Create(const OffsetRange& sealed_offsets, + size_t reader_count) { + if (sealed_offsets.begin < 0 || sealed_offsets.end < sealed_offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + return std::shared_ptr( + new RealtimeOffsetCoverage(sealed_offsets, reader_count)); + } + + Status Add(const arrow::Int64Array& offsets) { + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (offset < sealed_offsets_.begin || offset >= sealed_offsets_.end) { + return Status::Invalid( + "PK real-time store commit reader offset is outside the sealed range"); + } + min_seen_offset_ = std::min(min_seen_offset_, offset); + max_seen_offset_ = std::max(max_seen_offset_, offset); + ++seen_count_; + } + return Status::OK(); + } + + Status FinishReader() { + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && + (seen_count_ != sealed_offsets_.Count() || + (seen_count_ > 0 && (min_seen_offset_ != sealed_offsets_.begin || + max_seen_offset_ != sealed_offsets_.end - 1)))) { + return Status::Invalid( + "PK real-time store commit readers did not cover the sealed range"); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t reader_count) + : sealed_offsets_(sealed_offsets), reader_count_(reader_count) {} + + OffsetRange sealed_offsets_; + size_t reader_count_; + int64_t min_seen_offset_ = std::numeric_limits::max(); + int64_t max_seen_offset_ = std::numeric_limits::min(); + int64_t seen_count_ = 0; + size_t finished_reader_count_ = 0; +}; + +Status CheckPreparedField(const std::shared_ptr& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid(fmt::format("prepared schema missing transport field {} at index {}", + expected_field.Name(), field_idx)); + } + const std::shared_ptr& field = schema->field(field_idx); + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field->name() != expected_field.Name() || !field->type()->Equals(*expected_field.Type()) || + field->nullable() || field_id != expected_field.Id()) { + return Status::Invalid(fmt::format( + "prepared schema field {} must be non-null {}:{} with field id {}, got {}:{} " + "nullable={} field id {}", + field_idx, expected_field.Name(), expected_field.Type()->ToString(), + expected_field.Id(), field->name(), field->type()->ToString(), field->nullable(), + field_id)); + } + return Status::OK(); +} + +Result FindFieldIndexByPaimonId(const arrow::FieldVector& fields, int32_t field_id) { + std::optional matching_index; + for (int32_t i = 0; i < static_cast(fields.size()); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t candidate_id, + NestedProjectionUtils::GetPaimonFieldId(fields[i])); + if (candidate_id == field_id) { + if (matching_index.has_value()) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + matching_index = i; + } + } + if (matching_index.has_value()) { + return matching_index.value(); + } + return Status::Invalid(fmt::format("cannot find field id {} in prepared schema", field_id)); +} + +Status ValidateProjectionType(const std::shared_ptr& prepared_type, + const std::shared_ptr& query_type) { + if (prepared_type->id() != query_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + switch (query_type->id()) { + case arrow::Type::STRUCT: { + const arrow::FieldVector& prepared_fields = prepared_type->fields(); + for (const std::shared_ptr& query_field : query_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); + } + case arrow::Type::LIST: + return ValidateProjectionType(prepared_type->field(0)->type(), + query_type->field(0)->type()); + case arrow::Type::MAP: { + const std::shared_ptr prepared_map = + checked_pointer_cast(prepared_type); + const std::shared_ptr query_map = + checked_pointer_cast(query_type); + PAIMON_RETURN_NOT_OK( + ValidateProjectionType(prepared_map->key_type(), query_map->key_type())); + return ValidateProjectionType(prepared_map->item_type(), query_map->item_type()); + } + default: + if (!prepared_type->Equals(*query_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + return Status::OK(); + } +} + +Status ValidateProjectionSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema) { + arrow::FieldVector prepared_value_fields( + prepared_schema->fields().begin() + kPreparedValueStartIndex, + prepared_schema->fields().end()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_value_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_value_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); +} + +Status ValidateExactCommitSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& value_schema) { + if (prepared_schema->num_fields() != value_schema->num_fields() + kPreparedValueStartIndex) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + for (int32_t i = 0; i < value_schema->num_fields(); ++i) { + if (!prepared_schema->field(i + kPreparedValueStartIndex) + ->Equals(value_schema->field(i), true)) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + } + return Status::OK(); +} + +Result> AlignStructArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { + const std::shared_ptr data_type = + checked_pointer_cast(array->type()); + std::unordered_map data_field_id_to_idx; + data_field_id_to_idx.reserve(data_type->num_fields()); + for (int32_t i = 0; i < data_type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(data_type->field(i))); + if (!data_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared value struct", field_id)); + } + } + + arrow::ArrayVector aligned_arrays; + aligned_arrays.reserve(read_type->num_fields()); + for (const std::shared_ptr& read_field : read_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t read_field_id, + NestedProjectionUtils::GetPaimonFieldId(read_field)); + auto data_iter = data_field_id_to_idx.find(read_field_id); + if (data_iter == data_field_id_to_idx.end()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr null_child, + arrow::MakeArrayOfNull(read_field->type(), array->offset() + array->length(), + arrow_pool)); + aligned_arrays.push_back(std::move(null_child)); + continue; + } + std::shared_ptr child = + arrow::MakeArray(array->data()->child_data[data_iter->second]); + PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type(), arrow_pool)); + aligned_arrays.push_back(std::move(child)); + } + + std::shared_ptr aligned_data = array->data()->Copy(); + aligned_data->type = read_type; + aligned_data->child_data.clear(); + aligned_data->child_data.reserve(aligned_arrays.size()); + for (const std::shared_ptr& aligned_array : aligned_arrays) { + aligned_data->child_data.push_back(aligned_array->data()); + } + return arrow::MakeArray(std::move(aligned_data)); +} + +Result> AlignListArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { + std::shared_ptr values = array->values(); + PAIMON_ASSIGN_OR_RAISE(values, + AlignArrayByPaimonIds(values, read_type->value_type(), arrow_pool)); + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {values->data()}; + return arrow::MakeArray(new_data); +} + +Result> AlignMapArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { + std::shared_ptr keys = array->keys(); + PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type(), arrow_pool)); + std::shared_ptr items = array->items(); + PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type(), arrow_pool)); + + const std::shared_ptr& entries_data = array->data()->child_data[0]; + std::shared_ptr new_entries = entries_data->Copy(); + new_entries->type = arrow::struct_({read_type->key_field(), read_type->item_field()}); + new_entries->child_data = {keys->data(), items->data()}; + + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {std::move(new_entries)}; + return arrow::MakeArray(new_data); +} + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { + if (array->type()->id() != read_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + switch (read_type->id()) { + case arrow::Type::STRUCT: + return AlignStructArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + case arrow::Type::LIST: + return AlignListArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + case arrow::Type::MAP: + return AlignMapArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + default: + if (!array->type()->Equals(*read_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + return array; + } +} + +Result ProjectFieldsByPaimonIds( + const std::shared_ptr& data_batch, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema, arrow::MemoryPool* arrow_pool) { + std::unordered_map prepared_field_id_to_idx; + prepared_field_id_to_idx.reserve(prepared_schema->num_fields()); + for (int32_t i = kPreparedValueStartIndex; i < prepared_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(prepared_schema->field(i))); + if (!prepared_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + } + + arrow::ArrayVector result; + result.reserve(query_schema->num_fields()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_field_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + auto prepared_iter = prepared_field_id_to_idx.find(query_field_id); + if (prepared_iter == prepared_field_id_to_idx.end()) { + return Status::Invalid( + fmt::format("cannot find field id {} in prepared schema", query_field_id)); + } + std::shared_ptr field_array = data_batch->field(prepared_iter->second); + PAIMON_ASSIGN_OR_RAISE(field_array, + AlignArrayByPaimonIds(field_array, query_field->type(), arrow_pool)); + result.push_back(std::move(field_array)); + } + return result; +} + +class PreparedKeyValueReader final : public KeyValueRecordReader { + public: + PreparedKeyValueReader(std::unique_ptr&& reader, + const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& pool, + const std::shared_ptr& offset_coverage) + : reader_(std::move(reader)), + prepared_schema_(prepared_schema), + visible_offsets_(visible_offsets), + key_schema_(key_schema), + value_schema_(value_schema), + pool_(pool), + arrow_pool_(GetArrowPool(pool)), + offset_coverage_(offset_coverage) {} + + ~PreparedKeyValueReader() override { + Close(); + } + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(PreparedKeyValueReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->RowCount(); + } + + Result Next() override { + if (cursor_ >= reader_->RowCount()) { + return Status::Invalid("No more prepared key values in current iterator"); + } + const int64_t row = reader_->RowAt(cursor_); + std::shared_ptr key = + std::make_shared(reader_->key_ctx_, row); + auto value = std::make_unique(reader_->value_ctx_, row); + PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kind_array_->Value(row))); + int64_t sequence_number = reader_->sequence_number_array_->Value(row); + ++cursor_; + return KeyValue(row_kind, sequence_number, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + PreparedKeyValueReader* reader_; + int64_t cursor_ = 0; + }; + + Result> NextBatch() override { + if (first_error_.has_value()) { + return first_error_.value(); + } + Result> result = NextBatchImpl(); + if (!result.ok()) { + first_error_ = result.status(); + Close(); + } + return result; + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_) { + return; + } + closed_ = true; + ResetBatchState(); + reader_->Close(); + } + + private: + Result> NextBatchImpl() { + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } + return std::unique_ptr(); + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("cannot cast prepared batch to StructArray"); + } + std::shared_ptr data_batch = + checked_pointer_cast(arrow_array); + Status transport_status = + ValidatePreparedTransportSchema(arrow::schema(data_batch->type()->fields())); + if (!transport_status.ok()) { + return Status::Invalid( + "prepared batch field does not match prepared transport " + "schema: ", + transport_status.ToString()); + } + if (visible_offsets_.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema( + arrow::schema(data_batch->type()->fields()), key_schema_)); + PAIMON_ASSIGN_OR_RAISE( + arrow_array, + AlignArrayByPaimonIds(data_batch, arrow::struct_(prepared_schema_->fields()), + arrow_pool_.get())); + data_batch = checked_pointer_cast(arrow_array); + } + PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch)); + + std::shared_ptr> offset_array = + checked_pointer_cast>( + data_batch->field(kRealtimeOffsetIndex)); + if (offset_coverage_) { + PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array)); + } + + row_kind_array_ = checked_pointer_cast>( + data_batch->field(kValueKindIndex)); + sequence_number_array_ = checked_pointer_cast>( + data_batch->field(kSequenceNumberIndex)); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + key_schema_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector value_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + value_schema_, arrow_pool_.get())); + key_ctx_ = std::make_shared(key_fields, pool_); + value_ctx_ = std::make_shared(value_fields, pool_); + if (!SelectVisibleRows(*offset_array)) { + continue; + } + ArrowUtils::TraverseArray(data_batch); + return std::make_unique(this); + } + } + + Status ValidatePreparedBatch(const std::shared_ptr& data_batch) const { + if (data_batch->num_fields() != prepared_schema_->num_fields()) { + return Status::Invalid(fmt::format( + "prepared batch field count {} does not match prepared schema field count {}", + data_batch->num_fields(), prepared_schema_->num_fields())); + } + const arrow::FieldVector& batch_fields = data_batch->type()->fields(); + for (int32_t i = 0; i < data_batch->num_fields(); ++i) { + if (!batch_fields[i]->Equals(prepared_schema_->field(i), true)) { + return Status::Invalid(fmt::format( + "prepared batch field {} does not match declared prepared schema", i)); + } + } + if (!data_batch->field(kValueKindIndex) || + data_batch->field(kValueKindIndex)->type_id() != arrow::Type::INT8) { + return Status::Invalid("cannot cast VALUE_KIND column to int8 arrow array"); + } + if (!data_batch->field(kSequenceNumberIndex) || + data_batch->field(kSequenceNumberIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast SEQUENCE_NUMBER column to int64 arrow array"); + } + if (!data_batch->field(kRealtimeOffsetIndex) || + data_batch->field(kRealtimeOffsetIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast REALTIME_OFFSET column to int64 arrow array"); + } + if (data_batch->field(kValueKindIndex)->null_count() != 0 || + data_batch->field(kSequenceNumberIndex)->null_count() != 0 || + data_batch->field(kRealtimeOffsetIndex)->null_count() != 0) { + return Status::Invalid("prepared transport columns must not contain nulls"); + } + return Status::OK(); + } + + bool SelectVisibleRows(const arrow::Int64Array& offsets) { + if (!visible_offsets_.has_value()) { + return true; + } + visible_rows_.emplace(); + visible_rows_->reserve(offsets.length()); + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (offset >= visible_offsets_->begin && offset < visible_offsets_->end) { + visible_rows_->push_back(row); + } + } + return !visible_rows_->empty(); + } + + int64_t RowCount() const { + return visible_rows_.has_value() ? static_cast(visible_rows_->size()) + : row_kind_array_->length(); + } + + int64_t RowAt(int64_t ordinal) const { + return visible_rows_.has_value() ? (*visible_rows_)[ordinal] : ordinal; + } + + void ResetBatchState() { + key_ctx_.reset(); + value_ctx_.reset(); + row_kind_array_.reset(); + sequence_number_array_.reset(); + visible_rows_.reset(); + } + + private: + bool closed_ = false; + std::optional first_error_; + std::unique_ptr reader_; + std::shared_ptr prepared_schema_; + std::optional visible_offsets_; + std::shared_ptr key_schema_; + std::shared_ptr value_schema_; + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr offset_coverage_; + bool offset_coverage_finished_ = false; + std::shared_ptr key_ctx_; + std::shared_ptr value_ctx_; + std::shared_ptr> row_kind_array_; + std::shared_ptr> sequence_number_array_; + std::optional> visible_rows_; +}; + +} // namespace + +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema) { + if (!prepared_schema || prepared_schema->num_fields() < kPreparedValueStartIndex) { + return Status::Invalid("prepared schema must contain realtime transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kSequenceNumberIndex, SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kRealtimeOffsetIndex, SpecialFields::RealtimeOffset())); + return Status::OK(); +} + +namespace { + +Result> AdaptPreparedBatchReaderImpl( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool, + const std::shared_ptr& offset_coverage) { + std::unique_ptr owned_reader = std::move(reader); + if (!owned_reader) { + return Status::Invalid("prepared batch reader cannot be null"); + } + ScopeGuard close_guard([&owned_reader]() -> void { owned_reader->Close(); }); + if (visible_offsets.has_value() && visible_offsets->begin > visible_offsets->end) { + return Status::Invalid("prepared visible offset range begin exceeds end"); + } + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); + if (!key_schema) { + return Status::Invalid("prepared key schema cannot be null"); + } + if (!value_schema) { + return Status::Invalid("prepared value schema cannot be null"); + } + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, key_schema)); + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, value_schema)); + if (!visible_offsets.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateExactCommitSchema(prepared_schema, value_schema)); + } + std::unique_ptr result( + new PreparedKeyValueReader(std::move(owned_reader), prepared_schema, visible_offsets, + key_schema, value_schema, memory_pool, offset_coverage)); + close_guard.Release(); + return result; +} + +} // namespace + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + return AdaptPreparedBatchReaderImpl(std::move(reader), prepared_schema, visible_offsets, + key_schema, value_schema, memory_pool, + /*offset_coverage=*/nullptr); +} + +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard readers_guard([&readers, &adapted_readers]() { + CloseReaders(readers); + CloseReaders(adapted_readers); + }); + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(sealed_offsets, readers.size())); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr adapted_reader, + AdaptPreparedBatchReaderImpl( + std::move(reader), prepared_schema, std::nullopt, key_schema, + value_schema, memory_pool, offset_coverage)); + adapted_readers.push_back(std::move(adapted_reader)); + } + readers_guard.Release(); + return adapted_readers; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/prepared_key_value_reader.h b/src/paimon/core/realtime/prepared_key_value_reader.h new file mode 100644 index 00000000..4ef4887e --- /dev/null +++ b/src/paimon/core/realtime/prepared_key_value_reader.h @@ -0,0 +1,54 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include "arrow/type_fwd.h" +#include "paimon/core/io/key_value_record_reader.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/result.h" + +namespace paimon { +class BatchReader; +class MemoryPool; + +/// Validates the required leading fields of a prepared real-time transport schema. +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema); + +/// Adapts a plugin query reader and limits its rows to `visible_offsets` when present. +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + +/// Adapts commit readers and validates their offsets against `sealed_offsets`. +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 00000000..7d188609 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,327 @@ +/* + * 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 "paimon/core/realtime/primary_key_realtime_store.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/macros.h" + +namespace paimon { + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime v1 requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime v1 does not support data evolution"); + } + if (!options.GetFieldsSequenceGroups().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence groups"); + } + if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || + options.AggregationRemoveRecordOnDelete() || + !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { + return Status::NotImplemented("PK realtime v1 requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime v1 supports only ascending sequence.field.sort-order"); + } + if (options.NeedLookup() || options.DeletionVectorsEnabled() || + options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector primary_key_fields, + schema.TrimmedPrimaryKeyFields()); + for (const DataField& field : primary_key_fields) { + if (field.Type()->id() == arrow::Type::FLOAT || field.Type()->id() == arrow::Type::DOUBLE) { + return Status::NotImplemented( + "PK realtime v1 does not support FLOAT or DOUBLE primary keys"); + } + } + if (options.GlobalIndexEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(schema)); + if (!definitions.Definitions().empty()) { + return Status::NotImplemented("PK realtime v1 does not support global indexes"); + } + } + return Status::OK(); +} + +namespace { + +uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t total = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + total += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + total += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + total += GetArrayMemoryUsage(data->dictionary); + } + return total; +} + +struct StoredBatch { + std::shared_ptr data; + OffsetRange offset_range; + uint64_t memory_usage; +}; + +class Segment final : public RealtimeSegmentHandle { + public: + Segment(const OffsetRange& range, std::vector&& batches) + : range_(range), batches_(std::move(batches)) {} + + OffsetRange GetOffsetRange() const override { + return range_; + } + const std::vector& Batches() const { + return batches_; + } + + private: + OffsetRange range_; + std::vector batches_; +}; + +class ReadView final : public RealtimeReadView { + public: + explicit ReadView(std::vector>&& segments) + : segments_(std::move(segments)) { + if (!segments_.empty()) { + range_ = OffsetRange(segments_.front()->GetOffsetRange().begin, + segments_.back()->GetOffsetRange().end); + } + } + + std::optional GetOffsetRange() const override { + return range_; + } + const std::vector>& Segments() const { + return segments_; + } + + private: + std::vector> segments_; + std::optional range_; +}; + +class StoredBatchReader final : public BatchReader { + public: + explicit StoredBatchReader(const StoredBatch& batch) + : data_(batch.data), metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!data_) { + return MakeEofBatch(); + } + auto array = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*data_, array.get(), schema.get())); + data_.reset(); + return ReadBatch(std::move(array), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + void Close() override { + data_.reset(); + } + + private: + std::shared_ptr data_; + std::shared_ptr metrics_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + explicit Impl(std::shared_ptr prepared_schema) + : prepared_schema_(std::move(prepared_schema)) {} + + Status Write(RealtimeWriteBatch&& write_batch) { + if (!write_batch.batch || !write_batch.batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t row_count = write_batch.batch->GetData()->length; + if (write_batch.offset_range.begin < 0 || write_batch.offset_range.Count() != row_count || + row_count <= 0) { + return Status::Invalid("PK real-time offset range does not match batch row count"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(write_batch.batch->GetData(), + arrow::struct_(prepared_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time prepared batch is not a StructArray"); + } + std::shared_ptr prepared = + checked_pointer_cast(array); + std::lock_guard lock(mutex_); + building_.push_back( + StoredBatch{prepared, write_batch.offset_range, GetArrayMemoryUsage(prepared->data())}); + building_memory_usage_ += building_.back().memory_usage; + return Status::OK(); + } + + Result>> SealForCommit() { + std::lock_guard lock(mutex_); + if (building_.empty()) { + return std::optional>(); + } + OffsetRange range(building_.front().offset_range.begin, building_.back().offset_range.end); + std::shared_ptr segment = std::make_shared(range, std::move(building_)); + sealed_.push_back(segment); + building_.clear(); + building_memory_usage_ = 0; + return std::optional>(std::move(segment)); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& handle) { + std::shared_ptr segment = std::dynamic_pointer_cast(handle); + if (!segment) { + return Status::Invalid("segment was not created by the PK real-time store"); + } + std::vector> readers; + readers.reserve(segment->Batches().size()); + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back(std::make_unique(batch)); + } + return readers; + } + + Result> AcquireReadView() { + std::lock_guard lock(mutex_); + std::vector> segments = sealed_; + if (!building_.empty()) { + OffsetRange range(building_.front().offset_range.begin, + building_.back().offset_range.end); + segments.push_back( + std::make_shared(range, std::vector(building_))); + } + return std::shared_ptr(new ReadView(std::move(segments))); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t, const RealtimeQueryContext&) { + std::shared_ptr typed = std::dynamic_pointer_cast(view); + if (!typed) { + return Status::Invalid("read view was not created by the PK real-time store"); + } + std::vector> readers; + for (const std::shared_ptr& segment : typed->Segments()) { + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back(std::make_unique(batch)); + } + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_end_offset) { + std::lock_guard lock(mutex_); + while (!sealed_.empty() && sealed_.front()->GetOffsetRange().end <= committed_end_offset) { + sealed_.erase(sealed_.begin()); + } + return Status::OK(); + } + + uint64_t GetMemoryUsage() const { + std::lock_guard lock(mutex_); + uint64_t total = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_) { + for (const StoredBatch& batch : segment->Batches()) { + total += batch.memory_usage; + } + } + return total; + } + + private: + std::shared_ptr prepared_schema_; + mutable std::mutex mutex_; + std::vector building_; + std::vector> sealed_; + uint64_t building_memory_usage_ = 0; +}; + +PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) + : impl_(std::move(impl)) {} +PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; + +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& prepared_schema) { + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); + return std::shared_ptr( + new PrimaryKeyRealtimeStore(std::make_unique(prepared_schema))); +} +Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { + return impl_->Write(std::move(batch)); +} +Result>> +PrimaryKeyRealtimeStore::SealForCommit() { + return impl_->SealForCommit(); +} +Result>> PrimaryKeyRealtimeStore::CreateCommitReaders( + const std::shared_ptr& segment) { + return impl_->CreateCommitReaders(segment); +} +Result> PrimaryKeyRealtimeStore::AcquireReadView() { + return impl_->AcquireReadView(); +} +Result>> PrimaryKeyRealtimeStore::CreateQueryReaders( + const std::shared_ptr& view, int64_t offset, + const RealtimeQueryContext& context) { + return impl_->CreateQueryReaders(view, offset, context); +} +Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_end_offset) { + return impl_->AdvanceCommittedOffset(committed_end_offset); +} +uint64_t PrimaryKeyRealtimeStore::GetMemoryUsage() const { + return impl_->GetMemoryUsage(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h new file mode 100644 index 00000000..35f04485 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,63 @@ +/* + * 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. + */ + +#pragma once + +#include + +#include "paimon/realtime/realtime_store.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class CoreOptions; +class TableSchema; + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema); + +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& prepared_schema); + + ~PrimaryKeyRealtimeStore() override; + + Status Write(RealtimeWriteBatch&& batch) override; + Result>> SealForCommit() override; + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override; + Status AdvanceCommittedOffset(int64_t committed_end_offset) override; + uint64_t GetMemoryUsage() const override; + + private: + class Impl; + explicit PrimaryKeyRealtimeStore(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp new file mode 100644 index 00000000..6ccceba8 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,320 @@ +/* + * 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 "paimon/core/realtime/primary_key_realtime_store.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr PreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}); +} + +std::shared_ptr NestedPreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField(DataField( + 1, + arrow::field("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}))))}); +} + +std::shared_ptr PkSchema( + const std::shared_ptr& key_type = arrow::int64(), + const std::map& options = {}) { + return TableSchema::Create( + /*schema_id=*/0, + arrow::schema({arrow::field("id", key_type), arrow::field("value", arrow::utf8())}), + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options) + .value(); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(PreparedSchema()->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +std::unique_ptr MakeBatch(const std::shared_ptr& schema, + const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +void AssertOffsetsZero(const ArrowArray* array) { + ASSERT_NE(nullptr, array); + ASSERT_EQ(0, array->offset); + for (int64_t child = 0; child < array->n_children; ++child) { + AssertOffsetsZero(array->children[child]); + } + if (array->dictionary) { + AssertOffsetsZero(array->dictionary); + } +} + +Result ReadJson(const std::vector>& readers) { + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + batches.push_back(std::move(array)); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches)); + return result->ToString(); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestSupportedOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); + } +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsFloatingPrimaryKeys) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float32())), + "FLOAT or DOUBLE primary keys"); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float64())), + "FLOAT or DOUBLE primary keys"); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsEnabledGlobalIndex) { + const std::map option_map = {{Options::BUCKET, "1"}, + {Options::PK_BTREE_INDEX_COLUMNS, "id"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK_WITH_MSG( + ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::int64(), option_map)), + "does not support global indexes"); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema())); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_FALSE(segment.has_value()); + ASSERT_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), + "write batch is null"); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 1, "one"]])"), OffsetRange(0, 0)}), + "offset range does not match batch row count"); + + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[0, 1, 0, 1, "one"], [0, 2, 1, 2, "two"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 3, 2, 3, "three"]])"), OffsetRange(2, 3)})); + + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); + ASSERT_GT(store->GetMemoryUsage(), 0); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestBadTransportPrefix) { + const std::shared_ptr valid = PreparedSchema(); + std::vector invalid_fields; + + arrow::FieldVector wrong_type = valid->fields(); + wrong_type[0] = DataField::ConvertDataFieldToArrowField( + DataField(SpecialFields::ValueKind().Id(), + arrow::field("_VALUE_KIND", arrow::int32(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_type)); + + arrow::FieldVector nullable_sequence = valid->fields(); + nullable_sequence[1] = nullable_sequence[1]->WithNullable(true); + invalid_fields.push_back(std::move(nullable_sequence)); + + arrow::FieldVector wrong_offset_id = valid->fields(); + wrong_offset_id[2] = DataField::ConvertDataFieldToArrowField( + DataField(99, arrow::field("_REALTIME_OFFSET", arrow::int64(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_offset_id)); + + for (const arrow::FieldVector& fields : invalid_fields) { + ASSERT_NOK_WITH_MSG(PrimaryKeyRealtimeStore::Create(arrow::schema(fields)), + "prepared schema field"); + } +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[1, 6, 1, 1, "before"], [0, 5, 0, 3, "three"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_EQ( + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 0,\n 2\n ]\n-- " + "child 1 type: int64\n [\n 6,\n 5,\n 7\n ]\n-- child 2 type: int64\n [\n " + "1,\n 0,\n 2\n ]\n-- child 3 type: int64\n [\n 1,\n 3,\n 2\n ]\n-- child " + "4 type: string\n [\n \"before\",\n \"three\",\n \"after\"\n ]", + actual); + for (const std::unique_ptr& reader : readers) { + reader->Close(); + } +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderExportsZeroOffsets) { + std::shared_ptr schema = NestedPreparedSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(schema)); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(schema, R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]]])"), + OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + ASSERT_TRUE(arrow::ImportArray(batch.first.get(), batch.second.get()).ok()); + ASSERT_OK_AND_ASSIGN(batch, readers[0]->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCloseUnreadBatchReaders) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 10, 0, 1, "a"]])"), OffsetRange(0, 1)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 20, 1, 2, "b"]])"), OffsetRange(1, 2)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 30, 2, 3, "c"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(3, readers.size()); + for (const std::unique_ptr& reader : readers) { + reader->Close(); + } +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_EQ(std::optional(OffsetRange(4, 5)), view->GetOffsetRange()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 2, 1, 1, "one"]])"), OffsetRange(1, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + RealtimeQueryContext context{/*read_schema=*/nullptr, /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); +} + +} // namespace +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index af48b289..0fb4e58e 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -47,7 +47,7 @@ Result> RealtimeAppendOnlyWriter::Crea std::unique_ptr<::ArrowSchema> write_schema, const std::shared_ptr& realtime_context, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, + const std::shared_ptr& input_schema, StatisticsMode statistics_mode, const std::map& options, const std::shared_ptr& memory_pool) { if (!realtime_context) { @@ -55,9 +55,11 @@ Result> RealtimeAppendOnlyWriter::Crea } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, RealtimeContextImpl::Cast(realtime_context)); + RealtimeStoreCreateRequest request{ + std::move(write_schema), options, memory_pool, partition, bucket, + RealtimeStoreMode::APPEND_ONLY, statistics_mode}; PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, - realtime_context_impl->GetOrCreateRealtimeStore( - partition, bucket, std::move(write_schema), options, memory_pool)); + realtime_context_impl->GetOrCreateRealtimeStore(std::move(request))); return std::shared_ptr(new RealtimeAppendOnlyWriter( store_state.store, file_writer, input_schema, store_state.initial_offset, memory_pool)); } diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h index 29c198c6..a6190d3b 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.h +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -47,7 +47,7 @@ class RealtimeAppendOnlyWriter : public BatchWriter { std::unique_ptr<::ArrowSchema> write_schema, const std::shared_ptr& realtime_context, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, + const std::shared_ptr& input_schema, StatisticsMode statistics_mode, const std::map& options, const std::shared_ptr& memory_pool); diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 1c7a7cf5..404c0faa 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -34,14 +34,32 @@ #include #include +#include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/arrow/abi.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/uuid.h" #include "paimon/macros.h" #include "paimon/realtime/realtime_store.h" #include "paimon/status.h" namespace paimon { +namespace { + +std::string PartitionToString(const std::map& partition) { + std::string result = "{"; + for (auto iter = partition.begin(); iter != partition.end(); ++iter) { + if (iter != partition.begin()) { + result += ", "; + } + result += iter->first + "=" + iter->second; + } + return result + "}"; +} + +} // namespace RealtimeContextImpl::RealtimeContextImpl(const std::shared_ptr& factory) : factory_(factory) {} @@ -78,30 +96,36 @@ Status RealtimeContextImpl::Start() { } Result RealtimeContextImpl::GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr write_schema, const std::map& options, - const std::shared_ptr& memory_pool) { + RealtimeStoreCreateRequest&& request) { + if (!request.write_schema || !request.write_schema->release) { + return Status::Invalid("real-time store write schema is null"); + } + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr requested_schema, + arrow::ImportSchema(request.write_schema.get())); + schema_guard.Release(); std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); - const RealtimePartitionBucket key(partition, bucket); + const RealtimePartitionBucket key(request.partition, request.bucket); + auto iter = stores_.find(key); int64_t initial_offset = 0; auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); - } return Status::Invalid("real-time offset has reached INT64_MAX"); } initial_offset = offset_iter->second; } - auto iter = stores_.find(key); if (iter != stores_.end()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); + if (iter->second.mode != request.mode || + !iter->second.write_schema->Equals(*requested_schema, /*check_metadata=*/true)) { + return Status::Invalid("real-time store schema or mode mismatch for partition " + + PartitionToString(key.partition) + ", bucket " + + std::to_string(key.bucket) + "; recreate the RealtimeContext"); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - iter->second->AcquireReadView()); + iter->second.store->AcquireReadView()); if (!read_view) { return Status::Invalid("real-time store returned a null read view"); } @@ -116,27 +140,48 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset}; + return RealtimeStoreState{iter->second.store, initial_offset}; } - Result> store_result = - factory_->Create(std::move(write_schema), options, memory_pool); + if (!request.memory_pool) { + return Status::Invalid("real-time store memory pool is null"); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*requested_schema, request.write_schema.get())); + RealtimeStoreMode mode = request.mode; + Result> store_result = factory_->Create(std::move(request)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); - stores_.emplace(key, store); + stores_.emplace(key, StoreEntry{store, requested_schema, mode}); if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); } return RealtimeStoreState{std::move(store), initial_offset}; } +Result RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + std::lock_guard lock(mutex_); + auto iter = stores_.find(partition_bucket); + if (iter == stores_.end()) { + return Status::KeyError("real-time store not found for partition " + + PartitionToString(partition_bucket.partition) + ", bucket " + + std::to_string(partition_bucket.bucket)); + } + StoreEntry& entry = iter->second; + if (max_sequence_number > entry.materialized_max_sequence_number) { + entry.materialized_max_sequence_number = max_sequence_number; + } + return entry.materialized_max_sequence_number; +} + Result> RealtimeContextImpl::AcquireReadViews() { std::lock_guard lock(mutex_); std::vector result; result.reserve(stores_.size()); for (const auto& [partition_bucket, store] : stores_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - store->AcquireReadView()); + store.store->AcquireReadView()); result.push_back( - RealtimePartitionBucketView{partition_bucket, store, std::move(read_view)}); + RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; } @@ -229,12 +274,28 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, if (partition_bucket.bucket < 0 || committed_end_offset < 0) { return Status::Invalid("invalid partition-bucket committed offset"); } + } + // Only stores created by this context can contain state which cannot be restored in + // place. Offsets for other partition-buckets are reference state for lazy store creation + // and may be removed or rolled back without rebuilding the context. + std::lock_guard registry_lock(mutex_); + for (const auto& store_entry : stores_) { + const RealtimePartitionBucket& partition_bucket = store_entry.first; auto previous_iter = committed_offsets_.find(partition_bucket); - if (previous_iter != committed_offsets_.end()) { - if (committed_end_offset < previous_iter->second) { - return Status::Invalid( - "real-time partition-bucket committed offset cannot move backwards"); - } + if (previous_iter == committed_offsets_.end()) { + continue; + } + + auto current_iter = committed_offsets.find(partition_bucket); + if (current_iter == committed_offsets.end()) { + return Status::Invalid( + "real-time committed progress removed an active partition-bucket; recreate " + "RealtimeContext"); + } + if (current_iter->second < previous_iter->second) { + return Status::Invalid( + "real-time committed offset moved backwards for an active partition-bucket; " + "recreate RealtimeContext"); } } committed_offsets_ = committed_offsets; @@ -253,7 +314,7 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, } auto store_iter = stores_.find(partition_bucket); if (store_iter != stores_.end()) { - notifications.emplace_back(partition_bucket, store_iter->second, + notifications.emplace_back(partition_bucket, store_iter->second.store, committed_end_offset); } } diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 13fec77a..29ac7c05 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -32,11 +32,16 @@ #include #include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" #include "paimon/result.h" #include "paimon/visibility.h" struct ArrowSchema; +namespace arrow { +class Schema; +} // namespace arrow + namespace paimon { class RealtimeStore; @@ -64,11 +69,10 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { static Result> Cast( const std::shared_ptr& context); - Result GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, - const std::map& options, - const std::shared_ptr& memory_pool); + Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); + + Result AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number); Result> AcquireReadViews(); @@ -78,6 +82,9 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Status ReleaseReadView(const std::string& opaque_ticket); + // Returns an error requiring a new context if a newer snapshot removes or moves committed + // progress backwards for a store created by this context. Progress for inactive stores is + // only reference state and can be replaced in place. Status AdvanceCommittedProgress(int64_t snapshot_id, const RealtimeOffsetMap& committed_offsets); @@ -89,6 +96,13 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::chrono::steady_clock::time_point expire_at; }; + struct StoreEntry { + std::shared_ptr store; + std::shared_ptr write_schema; + RealtimeStoreMode mode; + int64_t materialized_max_sequence_number = -1; + }; + explicit RealtimeContextImpl(const std::shared_ptr& factory); Status Start(); @@ -98,8 +112,10 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::shared_ptr factory_; std::mutex mutex_; std::mutex progress_mutex_; - std::map> stores_; + std::map stores_; + // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; + // Progress already reflected in stores owned by this context. RealtimeOffsetMap reclaimed_offsets_; std::optional last_refreshed_snapshot_id_; std::mutex read_views_mutex_; diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index ea050094..fbdaa86b 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -9,12 +9,11 @@ * * 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. + * 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 @@ -31,7 +30,6 @@ #include "arrow/c/helpers.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/memory/memory_pool.h" -#include "paimon/realtime/realtime_store.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -49,26 +47,21 @@ class TestingRealtimeStore : public RealtimeStore { Status Write(RealtimeWriteBatch&&) override { return Status::OK(); } - Result>> SealForCommit() override { return std::optional>(); } - Result>> CreateCommitReaders( const std::shared_ptr&) override { return std::vector>(); } - Result> AcquireReadView() override { ++acquire_count; return std::make_shared(); } - Result>> CreateQueryReaders( const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { return std::vector>(); } - Status AdvanceCommittedOffset(int64_t committed_offset) override { ++advance_count; if (fail_next_advance) { @@ -78,7 +71,6 @@ class TestingRealtimeStore : public RealtimeStore { committed_offsets.push_back(committed_offset); return Status::OK(); } - uint64_t GetMemoryUsage() const override { return 0; } @@ -91,13 +83,11 @@ class TestingRealtimeStore : public RealtimeStore { class TestingRealtimeStoreFactory : public RealtimeStoreFactory { public: - Result> Create(std::unique_ptr write_schema, - const std::map&, - const std::shared_ptr&) override { - if (!write_schema || !write_schema->release) { + Result> Create(RealtimeStoreCreateRequest&& request) override { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("testing write schema is null"); } - ArrowSchemaRelease(write_schema.get()); + ArrowSchemaRelease(request.write_schema.get()); auto store = std::make_shared(); stores.push_back(store); return store; @@ -106,12 +96,14 @@ class TestingRealtimeStoreFactory : public RealtimeStoreFactory { std::vector> stores; }; -std::unique_ptr MakeWriteSchema() { - auto c_schema = std::make_unique(); +std::unique_ptr MakeWriteSchema( + const std::shared_ptr& id_type = arrow::int64(), + const std::shared_ptr& metadata = nullptr) { + auto schema = std::make_unique(); EXPECT_TRUE( - arrow::ExportSchema(*arrow::schema({arrow::field("id", arrow::int64())}), c_schema.get()) + arrow::ExportSchema(*arrow::schema({arrow::field("id", id_type)}, metadata), schema.get()) .ok()); - return c_schema; + return schema; } Result> CreateContext( @@ -121,31 +113,41 @@ Result> CreateContext( return RealtimeContextImpl::Cast(context); } -TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { +Result GetOrCreateAppendStore( + const std::shared_ptr& context, + const std::map& partition, int32_t bucket, + std::unique_ptr write_schema, const std::map& options, + const std::shared_ptr& memory_pool, + StatisticsMode statistics_mode = StatisticsMode::NONE) { + return context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(write_schema), options, memory_pool, partition, bucket, + RealtimeStoreMode::APPEND_ONLY, statistics_mode}); +} + +TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - - ASSERT_OK_AND_ASSIGN(RealtimeStoreState first_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, - MakeWriteSchema(), {{"k", "v"}}, pool)); - ASSERT_EQ(0, first_state.initial_offset); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState first, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, + MakeWriteSchema(), {{"k", "v"}}, GetDefaultPool())); + ASSERT_EQ(0, first.initial_offset); ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_again_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, pool)); - ASSERT_EQ(first_state.store, first_again_state.store); - ASSERT_EQ(0, first_again_state.initial_offset); + RealtimeStoreState second, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, + GetDefaultPool(), StatisticsMode::FULL)); + ASSERT_EQ(first.store, second.store); + ASSERT_EQ(0, second.initial_offset); ASSERT_EQ(1, factory->stores.size()); ASSERT_EQ(1, factory->stores[0]->acquire_count); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState second_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 1, MakeWriteSchema(), {}, pool)); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState third_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-03"}}, 0, MakeWriteSchema(), {}, pool)); - ASSERT_NE(first_state.store, second_state.store); - ASSERT_NE(first_state.store, third_state.store); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState third, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 1, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState fourth, + GetOrCreateAppendStore(context, {{"dt", "2026-08-03"}}, 0, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_NE(first.store, third.store); + ASSERT_NE(first.store, fourth.store); ASSERT_EQ(3, factory->stores.size()); ASSERT_OK_AND_ASSIGN(std::vector views, @@ -153,21 +155,94 @@ TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { ASSERT_EQ(3, views.size()); const RealtimePartitionBucket expected_partition_bucket({{"dt", "2026-08-02"}}, 0); ASSERT_EQ(expected_partition_bucket, views[0].partition_bucket); - ASSERT_EQ(first_state.store, views[0].store); + ASSERT_EQ(first.store, views[0].store); ASSERT_TRUE(views[0].read_view); ASSERT_EQ(2, factory->stores[0]->acquire_count); ASSERT_EQ(1, factory->stores[1]->acquire_count); ASSERT_EQ(1, factory->stores[2]->acquire_count); } +TEST(RealtimeContextTest, TestRejectsMismatchedModeOnStoreReuse) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + + ASSERT_NOK_WITH_MSG( + context->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ + MakeWriteSchema(), {}, GetDefaultPool(), partition, 0, RealtimeStoreMode::PRIMARY_KEY}), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_EQ(1, factory->stores.size()); +} + +TEST(RealtimeContextTest, TestRejectsMismatchedSchemaOnStoreReuse) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + std::shared_ptr metadata = + arrow::key_value_metadata({"identity"}, {"v1"}); + ASSERT_OK(GetOrCreateAppendStore( + context, partition, 0, MakeWriteSchema(arrow::int64(), metadata), {}, GetDefaultPool())); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(arrow::int32(), metadata), {}, + GetDefaultPool()), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore( + context, partition, 0, + MakeWriteSchema(arrow::int64(), arrow::key_value_metadata({"identity"}, {"v2"})), {}, + GetDefaultPool()), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_EQ(1, factory->stores.size()); +} + +TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + const RealtimePartitionBucket partition_bucket(partition, /*bucket=*/0); + + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + + ASSERT_OK_AND_ASSIGN(int64_t first, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/4)); + ASSERT_EQ(4, first); + ASSERT_OK_AND_ASSIGN(int64_t second, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/8)); + ASSERT_EQ(8, second); + ASSERT_OK_AND_ASSIGN(int64_t third, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/6)); + ASSERT_EQ(8, third); + ASSERT_OK_AND_ASSIGN(int64_t fourth, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/10)); + ASSERT_EQ(10, fourth); +} + +TEST(RealtimeContextTest, TestMaterializedSequenceRejectsMissingStore) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + + Result result = context->AdvanceMaterializedMaxSequenceNumber( + RealtimePartitionBucket({{"dt", "missing"}}, /*bucket=*/3), + /*max_sequence_number=*/4); + ASSERT_TRUE(result.status().IsKeyError()); + ASSERT_NOK_WITH_MSG(result, "real-time store not found for partition {dt=missing}, bucket 3"); +} + TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(2, factory->stores.size()); ASSERT_NOK_WITH_MSG(context->AdvanceCommittedProgress(-1, {}), @@ -184,9 +259,9 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); ASSERT_TRUE(factory->stores[1]->committed_offsets.empty()); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState restored_state, - context->GetOrCreateRealtimeStore({{"dt", "unknown"}}, 0, MakeWriteSchema(), {}, pool)); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState restored_state, + GetOrCreateAppendStore(context, {{"dt", "unknown"}}, 0, MakeWriteSchema(), + {}, GetDefaultPool())); ASSERT_EQ(9, restored_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress( @@ -203,15 +278,39 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } +TEST(RealtimeContextTest, TestRemovedInactivePartitionDoesNotRequireReopen) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map active_partition = {{"dt", "2026-08-02"}}; + const std::map inactive_partition = {{"dt", "2026-08-03"}}; + const RealtimePartitionBucket active_partition_bucket(active_partition, /*bucket=*/0); + const RealtimePartitionBucket inactive_partition_bucket(inactive_partition, /*bucket=*/0); + + ASSERT_OK(context->AdvanceCommittedProgress( + 5, {{active_partition_bucket, /*offset=*/7}, {inactive_partition_bucket, /*offset=*/9}})); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState active_state, + GetOrCreateAppendStore(context, active_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); + ASSERT_EQ(7, active_state.initial_offset); + + ASSERT_OK(context->AdvanceCommittedProgress(6, {{active_partition_bucket, /*offset=*/7}})); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState inactive_state, + GetOrCreateAppendStore(context, inactive_partition, 0, MakeWriteSchema(), + {}, GetDefaultPool())); + ASSERT_EQ(0, inactive_state.initial_offset); +} + TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 2, MakeWriteSchema(), {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 2, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(3, factory->stores.size()); factory->stores[1]->fail_next_advance = true; @@ -227,7 +326,7 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_OK_AND_ASSIGN( RealtimeStoreState failed_store_state, - context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), {}, pool)); + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(8, failed_store_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress(5, committed_offsets)); @@ -237,11 +336,44 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } +TEST(RealtimeContextTest, TestRequiresReopenWhenCommittedProgressMovesBackwards) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map first_partition = {{"dt", "2026-08-02"}}; + const std::map second_partition = {{"dt", "2026-08-03"}}; + const RealtimePartitionBucket first_partition_bucket(first_partition, /*bucket=*/0); + const RealtimePartitionBucket second_partition_bucket(second_partition, /*bucket=*/0); + + ASSERT_OK(GetOrCreateAppendStore(context, first_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, second_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); + ASSERT_OK(context->AdvanceCommittedProgress( + 5, {{first_partition_bucket, /*offset=*/7}, {second_partition_bucket, /*offset=*/9}})); + ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); + ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); + + ASSERT_NOK_WITH_MSG( + context->AdvanceCommittedProgress( + 6, {{first_partition_bucket, /*offset=*/6}, {second_partition_bucket, /*offset=*/10}}), + "recreate RealtimeContext"); + ASSERT_NOK_WITH_MSG( + context->AdvanceCommittedProgress(6, {{first_partition_bucket, /*offset=*/10}}), + "recreate RealtimeContext"); + ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); + ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); + + ASSERT_OK(context->AdvanceCommittedProgress( + 6, {{first_partition_bucket, /*offset=*/10}, {second_partition_bucket, /*offset=*/11}})); + ASSERT_EQ(std::vector({7, 10}), factory->stores[0]->committed_offsets); + ASSERT_EQ(std::vector({9, 11}), factory->stores[1]->committed_offsets); +} + TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); @@ -264,8 +396,8 @@ TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { TEST(RealtimeContextTest, TestExpiresAbandonedReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp new file mode 100644 index 00000000..ef85e8b6 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,325 @@ +/* + * 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 "paimon/core/realtime/realtime_primary_key_writer.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +struct PreparedArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleasePreparedArray(ArrowArray* array) { + auto* data = static_cast(array->private_data); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); + delete data; +} + +Status RetainPreparedArrayPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release || !arrow_pool) { + return Status::Invalid("cannot retain prepared batch memory pool"); + } + array->private_data = + new PreparedArrayPrivateData{array->release, array->private_data, arrow_pool}; + array->release = ReleasePreparedArray; + return Status::OK(); +} + +Result> PrepareBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, int64_t first_sequence_number, + int64_t first_offset, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr input, + arrow::ImportArray(batch->GetData(), arrow::struct_(write_schema->fields()))); + if (!input || input->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr values = checked_pointer_cast(input); + const int64_t count = values->length(); + arrow::Int8Builder kinds(arrow_pool); + arrow::Int64Builder sequences(arrow_pool); + arrow::Int64Builder offsets(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count)); + const std::vector& row_kinds = batch->GetRowKind(); + for (int64_t row = 0; row < count; ++row) { + const RecordBatch::RowKind kind = + row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row]; + kinds.UnsafeAppend(static_cast(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr kind_array; + std::shared_ptr sequence_array; + std::shared_ptr offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array)); + arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array), + std::move(offset_array)}; + columns.insert(columns.end(), values->fields().begin(), values->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr prepared, + arrow::StructArray::Make(std::move(columns), prepared_schema->fields())); + + std::vector sort_keys; + sort_keys.reserve(trimmed_primary_keys.size() + 1); + for (const std::string& key : trimmed_primary_keys) { + sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + } + sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), + arrow::compute::SortOrder::Ascending); + arrow::compute::ExecContext context(arrow_pool); + arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum indices, + arrow::compute::SortIndices(arrow::Datum(prepared), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(prepared), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + std::shared_ptr sorted_array = sorted.make_array(); + if (!sorted_array || sorted_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time sorted batch is not a StructArray"); + } + return checked_pointer_cast(std::move(sorted_array)); +} + +} // namespace + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool) { + if (!store_state.store || !merge_tree_writer || !write_schema || !key_comparator || + !realtime_context || !memory_pool) { + return Status::Invalid("PK real-time writer received a null dependency"); + } + if (trimmed_primary_keys.empty()) { + return Status::Invalid("PK real-time writer requires at least one primary key"); + } + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr field = write_schema->GetFieldByName(key); + if (!field) { + return Status::Invalid("PK field is missing from write schema: ", key); + } + key_fields.push_back(std::move(field)); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + prepared_fields.insert(prepared_fields.end(), write_schema->fields().begin(), + write_schema->fields().end()); + const RealtimePartitionBucket partition_bucket(partition, bucket); + PAIMON_ASSIGN_OR_RAISE(int64_t initial_max_sequence_number, + realtime_context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, restored_max_sequence_number)); + return std::shared_ptr(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, + arrow::schema(std::move(prepared_fields)), arrow::schema(std::move(key_fields)), + trimmed_primary_keys, key_comparator, store_state.initial_offset, + initial_max_sequence_number, memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, int64_t next_offset, + int64_t last_sequence_number, const std::shared_ptr& memory_pool) + : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + realtime_store_(realtime_store), + merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), + write_schema_(write_schema), + prepared_schema_(prepared_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t count = batch->GetData()->length; + if (count == 0) { + return Status::OK(); + } + const std::vector& row_kinds = batch->GetRowKind(); + if (!row_kinds.empty() && static_cast(row_kinds.size()) != count) { + return Status::Invalid("PK real-time row-kind count does not match batch row count"); + } + for (RecordBatch::RowKind row_kind : row_kinds) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, + RowKind::FromByteValue(static_cast(row_kind))); + static_cast(validated); + } + std::lock_guard lock(realtime_store_mutex_); + if (count > std::numeric_limits::max() - next_offset_) { + return Status::Invalid("real-time offset range exceeds INT64_MAX"); + } + // Reserve INT64_MAX as the exhausted sequence-number sentinel. + if (last_sequence_number_ >= std::numeric_limits::max() - count) { + return Status::Invalid("PK sequence range exceeds INT64_MAX"); + } + const int64_t first_sequence = last_sequence_number_ + 1; + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr prepared, + PrepareBatch(std::move(batch), write_schema_, prepared_schema_, trimmed_primary_keys_, + first_sequence, next_offset_, arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*prepared, output.get())); + PAIMON_RETURN_NOT_OK(RetainPreparedArrayPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ + count)})); + next_offset_ += count; + last_sequence_number_ += count; + PAIMON_RETURN_NOT_OK( + realtime_context_ + ->AdvanceMaterializedMaxSequenceNumber(partition_bucket_, last_sequence_number_) + .status()); + return Status::OK(); +} + +Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard prepare_lock(prepare_mutex_); + std::optional> segment; + { + std::lock_guard store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> sealed, + realtime_store_->SealForCommit()); + segment = std::move(sealed); + } + if (segment && !segment.value()) { + return Status::Invalid("PK real-time store sealed a null segment"); + } + std::optional sealed_range; + if (segment) { + sealed_range = segment.value()->GetOffsetRange(); + if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), sealed_range.value())); + } + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, + merge_tree_writer_->PrepareCommit(wait_compaction)); + if (segment) { + increment.SetRealtimeOffsetRange(sealed_range.value()); + } + return increment; +} + +Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr& segment, + const OffsetRange& sealed_offsets) { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + realtime_store_->CreateCommitReaders(segment)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> prepared_readers, + AdaptPreparedCommitBatchReaders(std::move(readers), prepared_schema_, sealed_offsets, + key_schema_, write_schema_, memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(prepared_readers.size()); + for (std::unique_ptr& prepared_reader : prepared_readers) { + auto merge_function = std::make_unique(/*ignore_delete=*/false); + sorted_readers.push_back(std::make_unique( + std::move(prepared_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); + } + return merge_tree_writer_->WriteSortedReaders(std::move(sorted_readers)); +} + +Status RealtimePrimaryKeyWriter::Compact(bool) { + return Status::Invalid("PK real-time write does not support explicit compaction"); +} +uint64_t RealtimePrimaryKeyWriter::GetMemoryUsage() const { + return realtime_store_->GetMemoryUsage(); +} +Status RealtimePrimaryKeyWriter::FlushMemory() { + return Status::OK(); +} +Result RealtimePrimaryKeyWriter::CompactNotCompleted() { + return merge_tree_writer_->CompactNotCompleted(); +} +Status RealtimePrimaryKeyWriter::Sync() { + return merge_tree_writer_->Sync(); +} +Status RealtimePrimaryKeyWriter::Close() { + return merge_tree_writer_->Close(); +} +std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { + return merge_tree_writer_->GetMetrics(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h new file mode 100644 index 00000000..d65c7e53 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,100 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" + +namespace arrow { +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; +class MergeTreeWriter; +class FieldsComparator; +class RealtimeContextImpl; +struct RealtimeStoreState; + +class RealtimePrimaryKeyWriter final : public BatchWriter { + public: + static Result> Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool); + + Status Write(std::unique_ptr&& batch) override; + Result PrepareCommit(bool wait_compaction) override; + Status Compact(bool full_compaction) override; + uint64_t GetMemoryUsage() const override; + Status FlushMemory() override; + Result CompactNotCompleted() override; + Status Sync() override; + Status Close() override; + std::shared_ptr GetMetrics() const override; + + private: + RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + int64_t next_offset, int64_t last_sequence_number, + const std::shared_ptr& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment, + const OffsetRange& sealed_offsets); + + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr realtime_store_; + std::shared_ptr merge_tree_writer_; + std::shared_ptr realtime_context_; + RealtimePartitionBucket partition_bucket_; + std::shared_ptr write_schema_; + std::shared_ptr prepared_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; + int64_t next_offset_; + int64_t last_sequence_number_; + std::mutex realtime_store_mutex_; + std::mutex prepare_mutex_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_reader_test.cpp b/src/paimon/core/realtime/realtime_reader_test.cpp index ec37cfed..ded06098 100644 --- a/src/paimon/core/realtime/realtime_reader_test.cpp +++ b/src/paimon/core/realtime/realtime_reader_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include "paimon/arrow/abi.h" #include "paimon/testing/utils/testharness.h" @@ -37,6 +38,8 @@ class TestingReadView : public RealtimeReadView { class TestingBatchReader : public BatchReader { public: + explicit TestingBatchReader(int32_t* close_count = nullptr) : close_count_(close_count) {} + Result NextBatch() override { return MakeEofBatch(); } @@ -45,7 +48,14 @@ class TestingBatchReader : public BatchReader { return nullptr; } - void Close() override {} + void Close() override { + if (close_count_) { + ++(*close_count_); + } + } + + private: + int32_t* close_count_; }; TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { @@ -57,5 +67,19 @@ TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { "inner reader is null"); } +TEST(RealtimeReaderTest, TestCloseReleasesResources) { + int32_t close_count = 0; + std::shared_ptr read_view = std::make_shared(); + std::weak_ptr weak_read_view = read_view; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + RealtimeReader::Create(std::move(read_view), + std::make_unique(&close_count))); + ASSERT_FALSE(weak_read_view.expired()); + reader->Close(); + ASSERT_EQ(1, close_count); + ASSERT_TRUE(weak_read_view.expired()); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/schema/arrow_schema_validator.cpp b/src/paimon/core/schema/arrow_schema_validator.cpp index f78e1550..0db1145f 100644 --- a/src/paimon/core/schema/arrow_schema_validator.cpp +++ b/src/paimon/core/schema/arrow_schema_validator.cpp @@ -28,6 +28,7 @@ #include "paimon/common/data/variant/variant_access_utils.h" #include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/types/vector_type.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/decimal_utils.h" #include "paimon/common/utils/string_utils.h" @@ -41,6 +42,7 @@ namespace paimon { bool ArrowSchemaValidator::IsNestedType(const std::shared_ptr& data_type) { return (data_type->id() == arrow::Type::MAP || data_type->id() == arrow::Type::LIST || + data_type->id() == arrow::Type::FIXED_SIZE_LIST || data_type->id() == arrow::Type::STRUCT); } @@ -128,6 +130,14 @@ Status ArrowSchemaValidator::ValidateDataTypeWithFieldId( value_field->type(), value_field->metadata(), /*allow_blob=*/false, field_id_set)); break; } + case arrow::Type::type::FIXED_SIZE_LIST: { + const auto& vector_type = checked_cast(*type); + if (vector_type.list_size() < 1 || + !VectorType::IsValidElementType(vector_type.value_type())) { + return Status::Invalid("Invalid VECTOR type: ", type->ToString()); + } + break; + } case arrow::Type::type::STRUCT: { if (VariantTypeUtils::IsVariantMetadata(key_value_metadata)) { // A variant struct is a leaf type: its value/metadata children carry fixed @@ -203,6 +213,18 @@ Status ArrowSchemaValidator::ValidateField(const std::shared_ptr& PAIMON_RETURN_NOT_OK(ValidateField(value_field, /*allow_blob=*/false)); break; } + case arrow::Type::type::FIXED_SIZE_LIST: { + const auto& vector_type = checked_cast(*field->type()); + if (vector_type.list_size() < 1) { + return Status::Invalid("Vector length must be positive, but was ", + vector_type.list_size()); + } + if (!VectorType::IsValidElementType(vector_type.value_type())) { + return Status::Invalid("Invalid element type for vector: ", + vector_type.value_type()->ToString()); + } + break; + } case arrow::Type::type::STRUCT: { if (VariantTypeUtils::IsVariantField(field)) { if (VariantAccessUtils::IsVariantAccessType(field->type())) { diff --git a/src/paimon/core/schema/arrow_schema_validator_test.cpp b/src/paimon/core/schema/arrow_schema_validator_test.cpp index 0363dff6..ed56b637 100644 --- a/src/paimon/core/schema/arrow_schema_validator_test.cpp +++ b/src/paimon/core/schema/arrow_schema_validator_test.cpp @@ -53,14 +53,29 @@ TEST(ArrowSchemaValidatorTest, TestSimple) { "col16", arrow::struct_({arrow::field("sub1", arrow::int8()), arrow::field("sub2", arrow::int16()), arrow::field("sub3", arrow::int64())})); + auto col17_field = arrow::field("col17", arrow::fixed_size_list(arrow::float32(), 3)); - auto arrow_schema = arrow::schema( - arrow::FieldVector({col1_field, col2_field, col3_field, col4_field, col5_field, col6_field, - col7_field, col8_field, col9_field, col10_field, col11_field, - col12_field, col13_field, col14_field, col15_field, col16_field})); + auto arrow_schema = arrow::schema(arrow::FieldVector( + {col1_field, col2_field, col3_field, col4_field, col5_field, col6_field, col7_field, + col8_field, col9_field, col10_field, col11_field, col12_field, col13_field, col14_field, + col15_field, col16_field, col17_field})); ASSERT_OK(ArrowSchemaValidator::ValidateSchema(*arrow_schema)); } +TEST(ArrowSchemaValidatorTest, TestVectorElementType) { + for (const auto& element_type : + {arrow::boolean(), arrow::int8(), arrow::int16(), arrow::int32(), arrow::int64(), + arrow::float32(), arrow::float64()}) { + auto vector = arrow::field("embedding", arrow::fixed_size_list(element_type, 3)); + ASSERT_OK(ArrowSchemaValidator::ValidateSchema(*arrow::schema({vector}))); + } + for (const auto& element_type : {arrow::utf8()}) { + auto vector = arrow::field("embedding", arrow::fixed_size_list(element_type, 3)); + ASSERT_NOK_WITH_MSG(ArrowSchemaValidator::ValidateSchema(*arrow::schema({vector})), + "Invalid element type for vector"); + } +} + TEST(ArrowSchemaValidatorTest, TestValidateNoRedundantFields) { auto col1_field = arrow::field("col1", arrow::int64()); auto col2_field = arrow::field("col2", arrow::int32()); diff --git a/src/paimon/core/schema/schema_validation.cpp b/src/paimon/core/schema/schema_validation.cpp index 7b8947ea..7342826d 100644 --- a/src/paimon/core/schema/schema_validation.cpp +++ b/src/paimon/core/schema/schema_validation.cpp @@ -38,6 +38,7 @@ #include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/vector_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/object_utils.h" #include "paimon/common/utils/preconditions.h" @@ -98,6 +99,15 @@ Status ValidateSharedShreddingFileFormat(const std::string& option_key, return Status::OK(); } +Status ValidateVectorFileFormat(const std::string& option_key, const std::string& file_format) { + if (StringUtils::ToLowerCase(file_format) != "parquet") { + return Status::Invalid( + fmt::format("VECTOR currently only supports parquet data files, but {} is {}.", + option_key, file_format)); + } + return Status::OK(); +} + Status ValidatePerLevelOption( const std::map& options, const std::string& option_key, const std::function& validator) { @@ -188,6 +198,7 @@ Status SchemaValidation::ValidateTableSchema(const TableSchema& schema) { PAIMON_RETURN_NOT_OK(ValidateRowTracking(schema, options)); PAIMON_RETURN_NOT_OK(ValidateBlobFields(schema, options)); PAIMON_RETURN_NOT_OK(ValidateMapStorageLayout(schema, options)); + PAIMON_RETURN_NOT_OK(ValidateVectorFields(schema, options)); return Status::OK(); } @@ -622,6 +633,9 @@ Status SchemaValidation::ValidateMapStorageLayout(const TableSchema& schema, if (ContainsBlobField(map_type->item_field())) { return Status::Invalid("MAP shared-shredding currently cannot contain BLOB fields."); } + if (VectorUtils::ContainsVectorField(map_type->item_field())) { + return Status::Invalid("MAP shared-shredding currently cannot contain VECTOR fields."); + } // Validate max-columns config PAIMON_RETURN_NOT_OK(options.GetMapSharedShreddingMaxColumns(field_name)); // Validate placement policy config @@ -648,4 +662,30 @@ Status SchemaValidation::ValidateMapStorageLayout(const TableSchema& schema, return Status::OK(); } +Status SchemaValidation::ValidateVectorFields(const TableSchema& schema, + const CoreOptions& options) { + bool has_vector = false; + for (const auto& field : schema.Fields()) { + if (VectorUtils::ContainsVectorField(field.ArrowField())) { + has_vector = true; + break; + } + } + if (!has_vector) { + return Status::OK(); + } + if (!schema.PrimaryKeys().empty()) { + return Status::NotImplemented( + "VECTOR fields in primary-key tables are not implemented yet."); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented( + "VECTOR fields in data-evolution tables are not implemented yet."); + } + PAIMON_RETURN_NOT_OK( + ValidateVectorFileFormat(Options::FILE_FORMAT, options.GetFileFormat()->Identifier())); + return ValidatePerLevelOption(options.ToMap(), Options::FILE_FORMAT_PER_LEVEL, + ValidateVectorFileFormat); +} + } // namespace paimon diff --git a/src/paimon/core/schema/schema_validation.h b/src/paimon/core/schema/schema_validation.h index 613372ff..abf4d5b0 100644 --- a/src/paimon/core/schema/schema_validation.h +++ b/src/paimon/core/schema/schema_validation.h @@ -75,6 +75,8 @@ class SchemaValidation { static Status ValidateMapStorageLayout(const TableSchema& schema, const CoreOptions& options); + static Status ValidateVectorFields(const TableSchema& schema, const CoreOptions& options); + static bool IsComplexType(const std::shared_ptr& field); }; diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 4c878113..050f0970 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -46,6 +46,76 @@ TEST(SchemaValidationTest, TestSimple) { ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestRealtimeOffsetIsNotGloballyReserved) { + auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(0, schema, {}, {}, {})); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); +} + +TEST(SchemaValidationTest, TestVectorType) { + auto vector_field = arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)); + auto schema = arrow::schema({arrow::field("id", arrow::int64()), vector_field}); + std::map parquet_options = {{Options::BUCKET, "-1"}, + {Options::FILE_FORMAT, "parquet"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, parquet_options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); + + std::map orc_options = {{Options::BUCKET, "-1"}, + {Options::FILE_FORMAT, "orc"}}; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, orc_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR currently only supports parquet data files"); + + std::map primary_key_options = {{Options::BUCKET, "1"}}; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"embedding"}, primary_key_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "in primary key field embedding is unsupported"); + + primary_key_options[Options::FILE_FORMAT] = "parquet"; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"id"}, primary_key_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR fields in primary-key tables are not implemented yet."); + + auto nested_schema = arrow::schema({ + arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_field->type())})), + }); + ASSERT_OK_AND_ASSIGN( + table_schema, + TableSchema::Create(/*schema_id=*/0, nested_schema, + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, primary_key_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR fields in primary-key tables are not implemented yet."); + + std::map data_evolution_options = { + {Options::BUCKET, "-1"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}, + }; + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, schema, + /*partition_keys=*/{}, + /*primary_keys=*/{}, data_evolution_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR fields in data-evolution tables are not implemented yet."); + ASSERT_OK_AND_ASSIGN(table_schema, + TableSchema::Create(/*schema_id=*/0, nested_schema, + /*partition_keys=*/{}, + /*primary_keys=*/{}, data_evolution_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "VECTOR fields in data-evolution tables are not implemented yet."); +} + TEST(SchemaValidationTest, TestRowTracking) { auto f0 = arrow::field("f0", arrow::utf8()); auto f1 = arrow::field("f1", arrow::int32()); diff --git a/src/paimon/core/schema/table_schema.cpp b/src/paimon/core/schema/table_schema.cpp index d7be7f11..6e2d8747 100644 --- a/src/paimon/core/schema/table_schema.cpp +++ b/src/paimon/core/schema/table_schema.cpp @@ -118,6 +118,14 @@ Result> TableSchema::AssignFieldIdsRecursively( /*set_field_id=*/false, field_id)); return arrow::field(field->name(), arrow::list(new_value_field), field->nullable(), metadata); + } else if (type->id() == arrow::Type::FIXED_SIZE_LIST) { + auto vector_type = checked_pointer_cast(field->type()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr new_value_field, + AssignFieldIdsRecursively(vector_type->value_field(), + /*set_field_id=*/false, field_id)); + return arrow::field(field->name(), + arrow::fixed_size_list(new_value_field, vector_type->list_size()), + field->nullable(), metadata); } else if (field->type()->id() == arrow::Type::MAP) { auto map_type = checked_pointer_cast(field->type()); std::shared_ptr key_field = map_type->key_field(); diff --git a/src/paimon/core/table/bucket_mode.cpp b/src/paimon/core/table/bucket_mode.cpp index d46739f8..429a787d 100644 --- a/src/paimon/core/table/bucket_mode.cpp +++ b/src/paimon/core/table/bucket_mode.cpp @@ -24,15 +24,13 @@ namespace paimon { BucketMode ResolveBucketMode(int32_t bucket, const std::shared_ptr& table_schema) { - if (bucket == BucketModeDefine::POSTPONE_BUCKET) { + bool has_primary_keys = !table_schema->PrimaryKeys().empty(); + // Postpone bucket is only valid for primary key tables. + if (has_primary_keys && bucket == BucketModeDefine::POSTPONE_BUCKET) { return BucketMode::POSTPONE_MODE; } if (bucket == -1) { - return table_schema->PrimaryKeys().empty() ? BucketMode::BUCKET_UNAWARE - : BucketMode::HASH_DYNAMIC; - } - if (bucket == BucketModeDefine::UNAWARE_BUCKET) { - return BucketMode::BUCKET_UNAWARE; + return has_primary_keys ? BucketMode::HASH_DYNAMIC : BucketMode::BUCKET_UNAWARE; } return BucketMode::HASH_FIXED; } diff --git a/src/paimon/core/table/bucket_mode.h b/src/paimon/core/table/bucket_mode.h index f87ab346..e6e6491c 100644 --- a/src/paimon/core/table/bucket_mode.h +++ b/src/paimon/core/table/bucket_mode.h @@ -63,10 +63,16 @@ enum class BucketMode { class BucketModeDefine { public: + /// The bucket id that all data of a `BucketMode::BUCKET_UNAWARE` table is written to. Note that + /// this is a bucket id, not a valid value of the 'bucket' option. static constexpr int32_t UNAWARE_BUCKET = 0; + /// The value of the 'bucket' option which enables `BucketMode::POSTPONE_MODE`, it is also used + /// as the bucket id of the data waiting to be assigned to a real bucket. static constexpr int32_t POSTPONE_BUCKET = -2; }; +/// Resolves the bucket mode from the 'bucket' option and the table schema. Note that an invalid +/// 'bucket' value is rejected by `SchemaValidation::ValidateBucket` instead of here. BucketMode ResolveBucketMode(int32_t bucket, const std::shared_ptr& table_schema); } // namespace paimon diff --git a/src/paimon/core/table/bucket_mode_test.cpp b/src/paimon/core/table/bucket_mode_test.cpp index 2a77e5ea..8feeb6e8 100644 --- a/src/paimon/core/table/bucket_mode_test.cpp +++ b/src/paimon/core/table/bucket_mode_test.cpp @@ -50,13 +50,17 @@ TEST(BucketModeTest, TestResolveBucketMode) { std::shared_ptr append_schema = CreateTableSchema(/*primary_keys=*/{}); std::shared_ptr pk_schema = CreateTableSchema(/*primary_keys=*/{"f0"}); + // Postpone bucket only applies to primary key tables. EXPECT_EQ(BucketMode::POSTPONE_MODE, + ResolveBucketMode(BucketModeDefine::POSTPONE_BUCKET, pk_schema)); + EXPECT_EQ(BucketMode::HASH_FIXED, ResolveBucketMode(BucketModeDefine::POSTPONE_BUCKET, append_schema)); + EXPECT_EQ(BucketMode::BUCKET_UNAWARE, ResolveBucketMode(-1, append_schema)); EXPECT_EQ(BucketMode::HASH_DYNAMIC, ResolveBucketMode(-1, pk_schema)); - EXPECT_EQ(BucketMode::BUCKET_UNAWARE, - ResolveBucketMode(BucketModeDefine::UNAWARE_BUCKET, pk_schema)); + EXPECT_EQ(BucketMode::HASH_FIXED, ResolveBucketMode(4, append_schema)); + EXPECT_EQ(BucketMode::HASH_FIXED, ResolveBucketMode(4, pk_schema)); } } // namespace paimon::test diff --git a/src/paimon/core/table/source/append_count_reader.cpp b/src/paimon/core/table/source/append_count_reader.cpp index 8f9af90a..5d684af6 100644 --- a/src/paimon/core/table/source/append_count_reader.cpp +++ b/src/paimon/core/table/source/append_count_reader.cpp @@ -21,6 +21,7 @@ #include "paimon/core/deletionvectors/deletion_vector.h" #include "paimon/core/table/source/data_split_impl.h" +#include "paimon/core/table/source/realtime_split.h" #include "paimon/status.h" namespace paimon { @@ -35,6 +36,16 @@ Result AppendCountReader::CountRows() { } Result AppendCountReader::CountSingleSplit(const std::shared_ptr& split) const { + std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); + if (realtime_split) { + int64_t total = realtime_split->MemoryEndOffset() - realtime_split->CommittedEndOffset(); + for (const std::shared_ptr& disk_split : realtime_split->DiskSplits()) { + PAIMON_ASSIGN_OR_RAISE(int64_t disk_count, CountSingleSplit(disk_split)); + total += disk_count; + } + return total; + } + auto data_split = std::dynamic_pointer_cast(split); if (!data_split) { return Status::Invalid("split cannot be cast to DataSplitImpl"); diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index 20786c9f..12b3823a 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -77,6 +77,13 @@ Result> AppendOnlyTableRead::CreateReader( std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::shared_ptr& split : splits) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); @@ -92,8 +99,6 @@ Result> AppendOnlyTableRead::CreateReader( } } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (!realtime_splits.empty()) { const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { @@ -106,7 +111,7 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return result; + return std::make_unique(std::move(readers), GetMemoryPool()); } Result> AppendOnlyTableRead::CreateRealtimeReader( @@ -124,6 +129,13 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); std::vector> readers; readers.reserve(realtime_split->DiskSplits().size() + 1); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), realtime_split->Bucket()); if (memory.partition_bucket != expected_partition_bucket) { @@ -150,8 +162,17 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( std::vector> memory_readers, memory.store->CreateQueryReaders(memory.read_view, realtime_split->CommittedEndOffset(), query_context)); - + const size_t first_memory_reader = readers.size(); + readers.reserve(readers.size() + memory_readers.size()); for (std::unique_ptr& memory_reader : memory_readers) { + readers.push_back(std::move(memory_reader)); + } + + for (size_t i = first_memory_reader; i < readers.size(); ++i) { + std::unique_ptr& memory_reader = readers[i]; + if (!memory_reader) { + return Status::Invalid("append-only real-time store returned a null query reader"); + } if (context_->EnablePredicateFilter() && context_->GetPredicate()) { PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create( std::move(memory_reader), @@ -159,14 +180,15 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, RealtimeReader::Create(memory.read_view, std::move(memory_reader))); - readers.push_back(std::move(realtime_reader)); + memory_reader = std::move(realtime_reader); } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (release_ticket) { PAIMON_RETURN_NOT_OK( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + readers_guard.Release(); return result; } @@ -183,17 +205,54 @@ Result> AppendOnlyTableRead::CreateDiskReader( Result> AppendOnlyTableRead::CreateCountReader( const std::vector>& splits) { - for (const std::shared_ptr& split : splits) { - if (std::dynamic_pointer_cast(split)) { - return Status::NotImplemented( - "CreateCountReader does not support process-local real-time splits"); - } - } if (context_->GetPredicate() != nullptr) { return Status::NotImplemented( "CreateCountReader with predicate pushdown is not supported yet"); } + std::vector> realtime_splits; + for (const std::shared_ptr& split : splits) { + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(split); + if (realtime_split) { + realtime_splits.push_back(std::move(realtime_split)); + } + } + if (!realtime_splits.empty()) { + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + for (const std::shared_ptr& realtime_split : realtime_splits) { + if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { + return Status::Invalid("unsupported real-time split version"); + } + if (realtime_split->MemoryEndOffset() < realtime_split->CommittedEndOffset()) { + return Status::Invalid("real-time memory upper offset is behind committed offset"); + } + PAIMON_ASSIGN_OR_RAISE( + RealtimePartitionBucketView memory, + realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); + const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), + realtime_split->Bucket()); + if (memory.partition_bucket != expected_partition_bucket) { + return Status::Invalid( + "real-time read-view ticket belongs to another partition-bucket"); + } + const std::optional memory_range = memory.read_view->GetOffsetRange(); + if (!memory_range || memory_range->end != realtime_split->MemoryEndOffset()) { + return Status::Invalid( + "real-time read-view ticket does not match the split offset range"); + } + } + for (const std::shared_ptr& realtime_split : realtime_splits) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + } + return std::make_unique(splits, context_->GetCoreOptions().GetFileSystem(), GetMemoryPool()); } diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 20880749..0f6b2189 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -20,12 +20,28 @@ #include "paimon/core/table/source/key_value_table_read.h" #include +#include +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" +#include "paimon/core/table/source/realtime_split.h" #include "paimon/status.h" namespace paimon { @@ -34,6 +50,61 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +struct ColumnarBatchContext; + +namespace { + +Result>> CreateMemoryReaders( + const std::shared_ptr& split, const RealtimePartitionBucketView& memory, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& context, + const std::shared_ptr& memory_pool) { + std::shared_ptr full_value_schema = + DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields()); + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + prepared_fields.insert(prepared_fields.end(), full_value_schema->fields().begin(), + full_value_schema->fields().end()); + std::shared_ptr prepared_schema = arrow::schema(std::move(prepared_fields)); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema, c_schema.get())); + ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, + memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); + ScopeGuard batch_readers_guard([&batch_readers]() { + for (const std::unique_ptr& reader : batch_readers) { + if (reader) { + reader->Close(); + } + } + }); + std::vector> result; + result.reserve(batch_readers.size()); + for (std::unique_ptr& reader : batch_readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null query reader"); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_reader, + AdaptPreparedBatchReader(std::move(reader), prepared_schema, + OffsetRange(split->CommittedEndOffset(), + split->MemoryEndOffset()), + key_schema, value_schema, memory_pool)); + auto merge = std::make_unique(false); + result.push_back(std::make_unique( + std::move(prepared_reader), key_comparator, + std::make_shared(std::move(merge)))); + } + batch_readers_guard.Release(); + return result; +} + +} // namespace KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, @@ -75,6 +146,11 @@ void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { + std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); + if (realtime_split) { + return CreateRealtimeReader(realtime_split, /*release_ticket=*/true); + } + std::shared_ptr dispatch_split = split; if (auto indexed_split = std::dynamic_pointer_cast(split)) { PAIMON_RETURN_NOT_OK(indexed_split->Validate()); @@ -126,8 +202,104 @@ Result> KeyValueTableRead::CreateReader( return Status::Invalid("create reader failed, not read match with data split."); } +Result> KeyValueTableRead::CreateReader( + const std::vector>& splits) { + std::vector> readers; + readers.reserve(splits.size()); + std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); + for (const std::shared_ptr& split : splits) { + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(split); + if (realtime_split) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + CreateRealtimeReader(realtime_split, + /*release_ticket=*/false)); + readers.push_back(std::move(reader)); + realtime_splits.push_back(std::move(realtime_split)); + } else { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); + readers.push_back(std::move(reader)); + } + } + + if (!realtime_splits.empty()) { + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + for (const std::shared_ptr& realtime_split : realtime_splits) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + } + return std::make_unique(std::move(readers), GetMemoryPool()); +} + +Result> KeyValueTableRead::CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket) { + if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { + return Status::Invalid("unsupported real-time split version"); + } + if (realtime_split->MemoryEndOffset() < realtime_split->CommittedEndOffset()) { + return Status::Invalid("real-time split memory end offset precedes committed end offset"); + } + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(RealtimePartitionBucketView memory, + realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); + const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), + realtime_split->Bucket()); + if (memory.partition_bucket != expected_partition_bucket) { + return Status::Invalid("real-time read-view ticket belongs to another partition-bucket"); + } + const std::optional memory_range = memory.read_view->GetOffsetRange(); + if (!memory_range || memory_range->end != realtime_split->MemoryEndOffset()) { + return Status::Invalid("real-time read-view ticket does not match the split offset range"); + } + for (const std::unique_ptr& read : split_reads_) { + auto* merge_read = dynamic_cast(read.get()); + if (merge_read) { + PAIMON_ASSIGN_OR_RAISE( + std::vector> memory_readers, + CreateMemoryReaders(realtime_split, memory, merge_read->GetKeySchema(), + merge_read->GetValueSchema(), merge_read->GetKeyComparator(), + context_, GetMemoryPool())); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + merge_read->CreateRealtimeReader(realtime_split->DiskSplits(), + std::move(memory_readers))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, + RealtimeReader::Create(memory.read_view, std::move(reader))); + if (release_ticket) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + return std::unique_ptr(std::move(realtime_reader)); + } + } + return Status::Invalid("create reader failed, merge file split read not found"); +} + Result> KeyValueTableRead::CreateCountReader( const std::vector>& splits) { + for (const std::shared_ptr& split : splits) { + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented( + "CreateCountReader does not support process-local real-time splits"); + } + } if (context_->GetPredicate() != nullptr) { return Status::NotImplemented( "CreateCountReader with predicate pushdown is not supported yet"); diff --git a/src/paimon/core/table/source/key_value_table_read.h b/src/paimon/core/table/source/key_value_table_read.h index d6a1c83d..6824ae59 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -35,6 +35,7 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +class RealtimeSplit; class KeyValueTableRead : public TableRead { public: @@ -45,6 +46,9 @@ class KeyValueTableRead : public TableRead { Result> CreateReader(const std::shared_ptr& split) override; + Result> CreateReader( + const std::vector>& splits) override; + Result> CreateCountReader( const std::vector>& splits) override; @@ -57,6 +61,9 @@ class KeyValueTableRead : public TableRead { const std::shared_ptr& memory_pool, const std::shared_ptr& executor); + Result> CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket); + std::vector> split_reads_; std::shared_ptr path_factory_; std::shared_ptr context_; diff --git a/src/paimon/core/table/source/realtime_split.h b/src/paimon/core/table/source/realtime_split.h index 2e264e82..7f571453 100644 --- a/src/paimon/core/table/source/realtime_split.h +++ b/src/paimon/core/table/source/realtime_split.h @@ -31,7 +31,10 @@ namespace paimon { -/// Split combining committed disk splits and a ticket for one immutable memory view. +/// Split combining disk splits and one immutable memory view. +/// +/// Append scans keep earlier disk splits independently schedulable and place only the tail disk +/// split in this wrapper. Other table semantics may choose a different disk grouping policy. /// /// `committed_end_offset` and `memory_end_offset` are exclusive bounds. Disk covers the committed /// prefix and memory readers return the remaining `[committed_end_offset, memory_end_offset)` diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index 5a76e81f..4c3968dc 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -19,6 +19,7 @@ #include "paimon/core/table/source/realtime_table_scan.h" +#include #include #include #include @@ -107,7 +108,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl for (const std::shared_ptr& split : disk_splits) { std::shared_ptr data_split = std::dynamic_pointer_cast(split); if (!data_split) { - return Status::Invalid("real-time append scan requires process-local data splits"); + return Status::Invalid("real-time scan requires process-local data splits"); } std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, @@ -119,7 +120,6 @@ Result>> RealtimeTableScan::CreateRealtimeSpl .push_back(split); } - // TODO(xinyu.lxy): Support splitting one partition-bucket into multiple real-time splits. std::vector> result; std::vector pinned_tickets; ScopeGuard ticket_guard([this, &pinned_tickets]() { @@ -151,7 +151,16 @@ Result>> RealtimeTableScan::CreateRealtimeSpl result.insert(result.end(), grouped_disk_splits.begin(), grouped_disk_splits.end()); continue; } + RealtimePartitionBucketView& memory = memory_iter->second; + if (!pk_table_) { + // Append tables can schedule all but the tail disk split independently. The tail split + // carries the immutable memory view so disk and memory are still concatenated by one + // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit. + auto tail_disk_split = std::prev(grouped_disk_splits.end()); + result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split); + grouped_disk_splits.erase(grouped_disk_splits.begin(), tail_disk_split); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_split, create_realtime_split(key, std::move(grouped_disk_splits), memory)); result.push_back(std::move(realtime_split)); @@ -168,7 +177,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl return result; } -RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, +RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -176,6 +185,7 @@ RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, const std::shared_ptr& scan_filter, int64_t read_view_ttl_millis) : disk_scan_(std::move(disk_scan)), + pk_table_(pk_table), realtime_context_(realtime_context), path_factory_(path_factory), snapshot_manager_(snapshot_manager), diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h index 7d036d42..692b749e 100644 --- a/src/paimon/core/table/source/realtime_table_scan.h +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -35,10 +35,10 @@ class FileSystem; class ScanFilter; class SnapshotManager; -/// Adds process-local memory splits to a normal append-table batch scan. +/// Adds process-local memory splits to a normal data-table batch scan. class RealtimeTableScan : public TableScan { public: - RealtimeTableScan(std::unique_ptr&& disk_scan, + RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -67,6 +67,7 @@ class RealtimeTableScan : public TableScan { const std::optional& snapshot_id) const; std::unique_ptr disk_scan_; + bool pk_table_; std::shared_ptr realtime_context_; std::shared_ptr path_factory_; std::shared_ptr snapshot_manager_; diff --git a/src/paimon/core/table/source/table_read.cpp b/src/paimon/core/table/source/table_read.cpp index e7e3f02f..3d6a3676 100644 --- a/src/paimon/core/table/source/table_read.cpp +++ b/src/paimon/core/table/source/table_read.cpp @@ -90,6 +90,9 @@ Result> CreateTableRead( const std::shared_ptr& memory_pool, const std::shared_ptr& executor) { const auto& core_options = internal_context->GetCoreOptions(); const auto& table_schema = internal_context->GetTableSchema(); + if (internal_context->GetRealtimeContext() && !core_options.RealtimeEnabled()) { + return Status::Invalid("real-time read requires realtime.enabled=true"); + } auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); PAIMON_ASSIGN_OR_RAISE(std::vector external_paths, core_options.CreateExternalPaths()); diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index a543a051..7f9fe568 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -41,6 +41,7 @@ #include "paimon/core/operation/data_evolution_file_store_scan.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/schema_validation.h" @@ -218,12 +219,12 @@ Result> TableScan::Create(std::unique_ptr> NewDataTableScan(const std::shared_ptrGetSpecificFileSystem(), {})); core_options.WithCache(context->GetCache()); - PAIMON_RETURN_NOT_OK(ValidateRealtimeScan(*table_schema, core_options, *context)); + PAIMON_RETURN_NOT_OK( + ValidateRealtimeScan(*table_schema, core_options, *context, read_optimized)); // validate options if (core_options.GetBucket() == -1) { if (!table_schema->PrimaryKeys().empty()) { @@ -340,7 +349,7 @@ Result> NewDataTableScan(const std::shared_ptr realtime_context, RealtimeContextImpl::Cast(context->GetRealtimeContext())); return std::make_unique( - std::move(batch_scan), realtime_context, path_factory, + std::move(batch_scan), pk_table, realtime_context, path_factory, snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(), context->GetScanFilters(), core_options.GetRealtimeReadViewTtlMillis()); } diff --git a/src/paimon/core/table/system/audit_log_system_table.cpp b/src/paimon/core/table/system/audit_log_system_table.cpp index e384146e..7572f906 100644 --- a/src/paimon/core/table/system/audit_log_system_table.cpp +++ b/src/paimon/core/table/system/audit_log_system_table.cpp @@ -440,7 +440,7 @@ Result> AuditLogSystemTable::NewChangelogRead( .SetPrefetchMaxParallelNum(context->GetPrefetchMaxParallelNum()) .EnableMultiThreadRowToBatch(context->EnableMultiThreadRowToBatch()) .SetRowToBatchThreadNumber(context->GetRowToBatchThreadNumber()) - .SetPrefetchCacheMode(context->GetPrefetchCacheMode()) + .SetReadAheadCacheEnabled(context->ReadAheadCacheEnabled()) .WithCacheConfig(context->GetCacheConfig()) .WithCache(context->GetCache()); diff --git a/src/paimon/core/table/system/read_optimized_system_table.cpp b/src/paimon/core/table/system/read_optimized_system_table.cpp index 516861cb..d7bfa091 100644 --- a/src/paimon/core/table/system/read_optimized_system_table.cpp +++ b/src/paimon/core/table/system/read_optimized_system_table.cpp @@ -58,6 +58,10 @@ std::map ReadOptimizedSystemTable::ReadOptimizedOption Result> ReadOptimizedSystemTable::NewScan( const std::shared_ptr& context) const { + if (context->GetRealtimeContext() && !table_schema_->PrimaryKeys().empty()) { + return Status::NotImplemented( + "PK real-time union read does not support read-optimized scans"); + } auto options = ReadOptimizedOptions(); ScanContextBuilder builder(table_path_); builder.SetOptions(options) @@ -109,7 +113,7 @@ Result> ReadOptimizedSystemTable::NewRead( .WithExecutor(context->GetExecutor()) .WithFileSystem(context->GetSpecificFileSystem()) .WithFileSystemSchemeToIdentifierMap(context->GetFileSystemSchemeToIdentifierMap()) - .SetPrefetchCacheMode(context->GetPrefetchCacheMode()) + .SetReadAheadCacheEnabled(context->ReadAheadCacheEnabled()) .WithCacheConfig(context->GetCacheConfig()) .WithCache(context->GetCache()) .SetReadFieldNames(context->GetReadFieldNames()) diff --git a/src/paimon/core/utils/field_mapping.cpp b/src/paimon/core/utils/field_mapping.cpp index 447d8d58..be7287dd 100644 --- a/src/paimon/core/utils/field_mapping.cpp +++ b/src/paimon/core/utils/field_mapping.cpp @@ -185,7 +185,7 @@ Result>> FieldMappingBuilder::CreateDa if (!read_fields[i].Type()->Equals(data_fields[i].Type())) { auto read_type_id = read_fields[i].Type()->id(); if (read_type_id == arrow::Type::STRUCT || read_type_id == arrow::Type::LIST || - read_type_id == arrow::Type::MAP) { + read_type_id == arrow::Type::MAP || read_type_id == arrow::Type::FIXED_SIZE_LIST) { // Nested type differs by pruning/evolution; the reader's reshape // handles it, no scalar cast. cast_executors.push_back(nullptr); diff --git a/src/paimon/core/utils/objects_file.h b/src/paimon/core/utils/objects_file.h index f8509fe2..b3135b31 100644 --- a/src/paimon/core/utils/objects_file.h +++ b/src/paimon/core/utils/objects_file.h @@ -19,7 +19,6 @@ #pragma once #include -#include #include #include #include @@ -78,6 +77,10 @@ class ObjectsFile { Result> WriteWithoutRolling(const std::vector& records); protected: + Status ReadArrowBatches( + const std::string& file_name, + const std::function&)>& consumer) const; + std::shared_ptr path_factory_; std::shared_ptr pool_; std::unique_ptr> serializer_; @@ -127,6 +130,31 @@ template Status ObjectsFile::Read(const std::string& file_name, const std::function(const T&)>& filter, std::vector* result) const { + return ReadArrowBatches( + file_name, + [this, &filter, result](const std::shared_ptr& struct_array) -> Status { + result->reserve(result->size() + struct_array->length()); + const arrow::ArrayVector& fields = struct_array->fields(); + for (int64_t i = 0; i < struct_array->length(); i++) { + ColumnarRow row(fields, pool_, i); + PAIMON_ASSIGN_OR_RAISE(T obj, serializer_->FromRow(row)); + if (filter) { + PAIMON_ASSIGN_OR_RAISE(bool filter_res, filter(obj)); + if (filter_res) { + result->push_back(std::move(obj)); + } + } else { + result->push_back(std::move(obj)); + } + } + return Status::OK(); + }); +} + +template +Status ObjectsFile::ReadArrowBatches( + const std::string& file_name, + const std::function&)>& consumer) const { std::string file_path = path_factory_->ToPath(file_name); std::shared_ptr file_input_stream; std::shared_ptr cached_bytes; @@ -171,22 +199,10 @@ Status ObjectsFile::Read(const std::string& file_name, if (!typed_array || typed_array->type_id() != arrow::Type::STRUCT) { return Status::Invalid(fmt::format("file {}, cannot cast to struct array", file_name)); } - auto* struct_array = checked_cast(typed_array.get()); - result->reserve(struct_array->length()); - for (int64_t i = 0; i < struct_array->length(); i++) { - ColumnarRow row(struct_array->fields(), pool_, i); - PAIMON_ASSIGN_OR_RAISE(T obj, serializer_->FromRow(row)); - if (filter) { - PAIMON_ASSIGN_OR_RAISE(bool filter_res, filter(obj)); - if (filter_res) { - result->push_back(std::move(obj)); - } - } else { - result->push_back(std::move(obj)); - } - } + std::shared_ptr struct_array = + checked_pointer_cast(typed_array); + PAIMON_RETURN_NOT_OK(consumer(struct_array)); } - reader->Close(); return Status::OK(); } diff --git a/src/paimon/core/utils/partition_utils.h b/src/paimon/core/utils/partition_utils.h new file mode 100644 index 00000000..b576d674 --- /dev/null +++ b/src/paimon/core/utils/partition_utils.h @@ -0,0 +1,67 @@ +/* + * 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. + */ + +#pragma once + +#include +#include + +#include "paimon/common/utils/binary_row_partition_computer.h" +#include "paimon/result.h" + +namespace paimon { + +class PartitionUtils { + public: + PartitionUtils() = delete; + ~PartitionUtils() = delete; + + static Result MatchPartitionSpec(const std::map& partition, + const std::map& partition_spec, + const BinaryRowPartitionComputer& partition_computer) { + for (const auto& entry : partition_spec) { + if (partition.find(entry.first) == partition.end()) { + return false; + } + } + // Dynamic overwrite already supplies canonical partition values. Avoid trying to parse + // legacy DATE names such as "19723" as user-facing DATE literals again. + if (MatchNormalizedPartitionSpec(partition, partition_spec)) { + return true; + } + std::map normalized_partition_spec; + PAIMON_ASSIGN_OR_RAISE(normalized_partition_spec, + partition_computer.NormalizePartitionSpec(partition_spec)); + return MatchNormalizedPartitionSpec(partition, normalized_partition_spec); + } + + static bool MatchNormalizedPartitionSpec( + const std::map& partition, + const std::map& normalized_partition_spec) { + for (const auto& [key, value] : normalized_partition_spec) { + auto iter = partition.find(key); + if (iter == partition.end() || iter->second != value) { + return false; + } + } + return true; + } +}; + +} // namespace paimon diff --git a/src/paimon/core/utils/partition_utils_test.cpp b/src/paimon/core/utils/partition_utils_test.cpp new file mode 100644 index 00000000..2ac62b92 --- /dev/null +++ b/src/paimon/core/utils/partition_utils_test.cpp @@ -0,0 +1,71 @@ +/* + * 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 "paimon/core/utils/partition_utils.h" + +#include +#include +#include + +#include "arrow/type.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/binary_row_partition_computer.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(PartitionUtilsTest, MatchNormalizedPartitionSpec) { + const std::map partition = {{"dt", "2026-08-21"}, {"region", "cn"}}; + + ASSERT_TRUE( + PartitionUtils::MatchNormalizedPartitionSpec(partition, /*normalized_partition_spec=*/{})); + ASSERT_TRUE(PartitionUtils::MatchNormalizedPartitionSpec(partition, {{"dt", "2026-08-21"}})); + ASSERT_TRUE(PartitionUtils::MatchNormalizedPartitionSpec( + partition, {{"dt", "2026-08-21"}, {"region", "cn"}})); + ASSERT_FALSE(PartitionUtils::MatchNormalizedPartitionSpec(partition, {{"dt", "2026-08-22"}})); + ASSERT_FALSE(PartitionUtils::MatchNormalizedPartitionSpec(partition, {{"hour", "12"}})); +} + +TEST(PartitionUtilsTest, MatchPartitionSpecNormalizesPartialSpec) { + std::shared_ptr schema = + arrow::schema({arrow::field("dt", arrow::date32()), arrow::field("region", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr partition_computer, + BinaryRowPartitionComputer::Create( + /*partition_keys=*/{"dt", "region"}, schema, + /*default_part_value=*/"__DEFAULT_PARTITION__", + /*legacy_partition_name_enabled=*/true, GetDefaultPool())); + const std::map partition = {{"dt", "19723"}, {"region", "cn"}}; + ASSERT_OK_AND_ASSIGN( + bool raw_spec_matches, + PartitionUtils::MatchPartitionSpec(partition, {{"dt", "2024-01-01"}}, *partition_computer)); + ASSERT_TRUE(raw_spec_matches); + + ASSERT_OK_AND_ASSIGN( + bool normalized_spec_matches, + PartitionUtils::MatchPartitionSpec(partition, {{"dt", "19723"}}, *partition_computer)); + ASSERT_TRUE(normalized_spec_matches); + + ASSERT_OK_AND_ASSIGN(bool unknown_key_matches, + PartitionUtils::MatchPartitionSpec( + partition, {{"unknown_partition_key", "value"}}, *partition_computer)); + ASSERT_FALSE(unknown_key_matches); +} + +} // namespace paimon::test diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 12713ca5..1a7345fd 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -19,10 +19,11 @@ #include "paimon/core/utils/primary_key_table_utils.h" #include -#include #include #include +#include #include +#include #include "arrow/type.h" #include "gtest/gtest.h" diff --git a/src/paimon/format/avro/avro_direct_decoder.cpp b/src/paimon/format/avro/avro_direct_decoder.cpp index f837eed0..f9c8a9a4 100644 --- a/src/paimon/format/avro/avro_direct_decoder.cpp +++ b/src/paimon/format/avro/avro_direct_decoder.cpp @@ -33,6 +33,22 @@ namespace paimon::avro { +const AvroDirectDecoder::DecodeContext::BuilderMetadata& +AvroDirectDecoder::DecodeContext::GetBuilderMetadata(const arrow::ArrayBuilder* builder) { + auto iter = builder_metadata_.find(builder); + if (iter != builder_metadata_.end()) { + return iter->second; + } + + std::shared_ptr data_type = builder->type(); + BuilderMetadata metadata{data_type->id(), std::nullopt}; + if (data_type->id() == arrow::Type::TIMESTAMP) { + metadata.timestamp_unit = + checked_cast(data_type.get())->unit(); + } + return builder_metadata_.emplace(builder, metadata).first->second; +} + namespace { /// Forward declaration for mutual recursion. @@ -266,8 +282,8 @@ Status DecodeAvroValueToBuilder(const ::avro::NodePtr& avro_node, case ::avro::AVRO_INT: { int32_t value = decoder->decodeInt(); - auto arrow_type = array_builder->type(); - switch (arrow_type->id()) { + const auto& builder_metadata = ctx->GetBuilderMetadata(array_builder); + switch (builder_metadata.type) { case arrow::Type::INT8: { auto* builder = checked_cast(array_builder); PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Append(value)); @@ -287,7 +303,7 @@ Status DecodeAvroValueToBuilder(const ::avro::NodePtr& avro_node, if (logical_type.type() != ::avro::LogicalType::Type::DATE) { return Status::TypeError( fmt::format("Unexpected avro type [{}] with arrow type [{}].", - ::avro::toString(type), arrow_type->ToString())); + ::avro::toString(type), array_builder->type()->ToString())); } auto* builder = checked_cast(array_builder); PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Append(value)); @@ -296,7 +312,7 @@ Status DecodeAvroValueToBuilder(const ::avro::NodePtr& avro_node, default: return Status::TypeError( fmt::format("Unexpected avro type [{}] with arrow type [{}].", - ::avro::toString(type), arrow_type->ToString())); + ::avro::toString(type), array_builder->type()->ToString())); } } @@ -315,9 +331,9 @@ Status DecodeAvroValueToBuilder(const ::avro::NodePtr& avro_node, case ::avro::LogicalType::Type::LOCAL_TIMESTAMP_MICROS: case ::avro::LogicalType::Type::LOCAL_TIMESTAMP_NANOS: { auto* builder = checked_cast(array_builder); - auto ts_type = checked_cast(builder->type().get()); // for arrow second, we need to convert it from avro millisecond - if (ts_type->unit() == arrow::TimeUnit::type::SECOND) { + const auto& builder_metadata = ctx->GetBuilderMetadata(builder); + if (builder_metadata.timestamp_unit == arrow::TimeUnit::type::SECOND) { value /= DateTimeUtils::CONVERSION_FACTORS[DateTimeUtils::MILLISECOND]; } PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Append(value)); diff --git a/src/paimon/format/avro/avro_direct_decoder.h b/src/paimon/format/avro/avro_direct_decoder.h index c507091a..6422f915 100644 --- a/src/paimon/format/avro/avro_direct_decoder.h +++ b/src/paimon/format/avro/avro_direct_decoder.h @@ -22,7 +22,11 @@ #pragma once +#include #include +#include +#include +#include #include "arrow/array/builder_base.h" #include "avro/Decoder.hh" @@ -41,10 +45,26 @@ class AvroDirectDecoder { /// Avoids frequent small allocations by reusing temporary buffers across multiple decode /// operations. This is particularly important for string, binary, and decimal data types. struct DecodeContext { + struct BuilderMetadata { + arrow::Type::type type; + std::optional timestamp_unit; + }; + + /// Returns immutable type metadata without repeatedly copying the builder's DataType. + const BuilderMetadata& GetBuilderMetadata(const arrow::ArrayBuilder* builder); + + /// Clears metadata before the builder tree is replaced or destroyed. + void ClearBuilderMetadata() { + builder_metadata_.clear(); + } + // Scratch buffer for string decoding (reused across rows) std::string string_scratch; // Scratch buffer for binary/decimal data (reused across rows) std::vector bytes_scratch; + + private: + std::unordered_map builder_metadata_; }; /// Directly decode Avro data to Arrow array builders without GenericDatum diff --git a/src/paimon/format/avro/avro_direct_encoder_decoder_test.cpp b/src/paimon/format/avro/avro_direct_encoder_decoder_test.cpp index f276d946..78f4ca48 100644 --- a/src/paimon/format/avro/avro_direct_encoder_decoder_test.cpp +++ b/src/paimon/format/avro/avro_direct_encoder_decoder_test.cpp @@ -62,6 +62,7 @@ class AvroDirectEncoderDecoderTest : public ::testing::Test { auto decoder = ::avro::binaryDecoder(); decoder->init(*input_stream); + decode_ctx_.ClearBuilderMetadata(); for (int32_t i = 0; i < expected_count; ++i) { PAIMON_RETURN_NOT_OK(AvroDirectDecoder::DecodeAvroToBuilder( avro_node, projection, decoder.get(), builder, &decode_ctx_)); @@ -157,6 +158,18 @@ TEST_F(AvroDirectEncoderDecoderTest, TestIntegerTypes) { CheckResult(schema_json, input_array, &builder); } + // Test INT16 + { + std::string schema_json = R"({"type": "int"})"; + arrow::Int16Builder builder; + ASSERT_TRUE(builder.Append(1).ok()); + ASSERT_TRUE(builder.Append(-32768).ok()); + ASSERT_TRUE(builder.Append(32767).ok()); + std::shared_ptr input_array; + ASSERT_TRUE(builder.Finish(&input_array).ok()); + CheckResult(schema_json, input_array, &builder); + } + // Test INT32 { std::string schema_json = R"({"type": "int"})"; @@ -182,6 +195,16 @@ TEST_F(AvroDirectEncoderDecoderTest, TestIntegerTypes) { } } +TEST_F(AvroDirectEncoderDecoderTest, TestDecodeContextBuilderMetadataLifecycle) { + arrow::Int8Builder int8_builder; + ASSERT_EQ(decode_ctx_.GetBuilderMetadata(&int8_builder).type, arrow::Type::INT8); + + decode_ctx_.ClearBuilderMetadata(); + + arrow::Int16Builder int16_builder; + ASSERT_EQ(decode_ctx_.GetBuilderMetadata(&int16_builder).type, arrow::Type::INT16); +} + TEST_F(AvroDirectEncoderDecoderTest, TestFloatingPointTypes) { // Test FLOAT { diff --git a/src/paimon/format/avro/avro_file_batch_reader.cpp b/src/paimon/format/avro/avro_file_batch_reader.cpp index f48ec4cc..1e217f5c 100644 --- a/src/paimon/format/avro/avro_file_batch_reader.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader.cpp @@ -172,6 +172,7 @@ Status AvroFileBatchReader::SetReadSchema(::ArrowSchema* read_schema, } reader_ = std::move(reader); array_builder_ = std::move(array_builder); + decode_context_.ClearBuilderMetadata(); previous_first_row_ = std::numeric_limits::max(); previous_batch_row_count_ = 0; next_row_to_read_ = std::numeric_limits::max(); diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index cd627eb5..a201839b 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader.cpp @@ -46,6 +46,16 @@ #include "paimon/format/orc/predicate_converter.h" namespace paimon::orc { +namespace { + +void CollectAllColumnIds(const ::orc::Type* type, std::vector* column_ids) { + column_ids->push_back(type->getColumnId()); + for (uint64_t i = 0; i < type->getSubtypeCount(); ++i) { + CollectAllColumnIds(type->getSubtype(i), column_ids); + } +} + +} // namespace OrcFileBatchReader::OrcFileBatchReader(std::unique_ptr<::orc::ReaderMetrics>&& reader_metrics, std::unique_ptr&& reader, @@ -228,13 +238,15 @@ Status OrcFileBatchReader::CollectTargetColumnIds(const ::orc::Type* src_type, } break; } - // Do not support partial field recall inside list/map types. default: { + // Partial field recall inside list/map types is unsupported, so the target must match + // the complete source subtree. Include the container and every descendant because all + // of their streams are recalled by the ORC reader. if (src_type->toString() != target_type->toString()) { return Status::Invalid(fmt::format("type mismatch: src {} vs target {}", src_type->toString(), target_type->toString())); } - target_column_ids->push_back(src_type->getColumnId()); + CollectAllColumnIds(src_type, target_column_ids); break; } } diff --git a/src/paimon/format/orc/orc_file_batch_reader_test.cpp b/src/paimon/format/orc/orc_file_batch_reader_test.cpp index 038c93e2..c39ba516 100644 --- a/src/paimon/format/orc/orc_file_batch_reader_test.cpp +++ b/src/paimon/format/orc/orc_file_batch_reader_test.cpp @@ -368,8 +368,9 @@ TEST_F(OrcFileBatchReaderTest, TestCreateRowReaderOptions) { OrcFileBatchReader::CreateRowReaderOptions( src_type.get(), target_type.get(), /*search_arg=*/nullptr, options, &target_column_ids)); - // Struct IDs (0, 1) not included. Selected: sub1(2), sub2-list(3), sub3(6), col3-map(8). - ASSERT_EQ(target_column_ids, (std::vector{2, 3, 6, 8})); + // Struct IDs (0, 1) are not included. LIST/MAP containers include their complete + // subtrees: sub1(2), sub2(3, 4, 5), sub3(6), col3(8, 9, 10). + ASSERT_EQ(target_column_ids, (std::vector{2, 3, 4, 5, 6, 8, 9, 10})); } { // read with type mismatch in nested field @@ -485,6 +486,98 @@ TEST_F(OrcFileBatchReaderTest, TestCreateRowReaderOptions) { } } +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsPrimitiveList) { + std::unique_ptr<::orc::Type> src_type = + ::orc::Type::buildTypeFromString("struct,ignored:string>"); + std::unique_ptr<::orc::Type> target_type = + ::orc::Type::buildTypeFromString("struct>"); + std::vector target_column_ids; + + ASSERT_OK(OrcFileBatchReader::CollectTargetColumnIds(src_type.get(), target_type.get(), + &target_column_ids)); + // root struct(0), items-list(1), element(2), ignored(3) + ASSERT_EQ(target_column_ids, (std::vector{1, 2})); +} + +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsDeeplyNestedList) { + std::string schema = "struct>>>"; + std::unique_ptr<::orc::Type> src_type = ::orc::Type::buildTypeFromString(schema); + std::unique_ptr<::orc::Type> target_type = ::orc::Type::buildTypeFromString(schema); + std::vector target_column_ids; + + ASSERT_OK(OrcFileBatchReader::CollectTargetColumnIds(src_type.get(), target_type.get(), + &target_column_ids)); + // root struct(0), outer list(1), inner list(2), element struct(3), value(4), label(5) + ASSERT_EQ(target_column_ids, (std::vector{1, 2, 3, 4, 5})); +} + +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsPrimitiveMap) { + std::unique_ptr<::orc::Type> src_type = + ::orc::Type::buildTypeFromString("struct,ignored:double>"); + std::unique_ptr<::orc::Type> target_type = + ::orc::Type::buildTypeFromString("struct>"); + std::vector target_column_ids; + + ASSERT_OK(OrcFileBatchReader::CollectTargetColumnIds(src_type.get(), target_type.get(), + &target_column_ids)); + // root struct(0), attributes-map(1), key(2), value(3), ignored(4) + ASSERT_EQ(target_column_ids, (std::vector{1, 2, 3})); +} + +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsDeeplyNestedMap) { + std::string schema = + "struct>>>>"; + std::unique_ptr<::orc::Type> src_type = ::orc::Type::buildTypeFromString(schema); + std::unique_ptr<::orc::Type> target_type = ::orc::Type::buildTypeFromString(schema); + std::vector target_column_ids; + + ASSERT_OK(OrcFileBatchReader::CollectTargetColumnIds(src_type.get(), target_type.get(), + &target_column_ids)); + // root struct(0), map(1), key(2), value-list(3), element struct(4), score(5), + // tags-list(6), tag element(7) + ASSERT_EQ(target_column_ids, (std::vector{1, 2, 3, 4, 5, 6, 7})); +} + +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsStructProjectionWithListAndMap) { + std::unique_ptr<::orc::Type> src_type = ::orc::Type::buildTypeFromString( + "struct,plain:double,attributes:map>," + "ignored:string>"); + std::unique_ptr<::orc::Type> target_type = ::orc::Type::buildTypeFromString( + "struct,attributes:map>>"); + std::vector target_column_ids; + + ASSERT_OK(OrcFileBatchReader::CollectTargetColumnIds(src_type.get(), target_type.get(), + &target_column_ids)); + // root struct(0) and outer struct(1) are not included. Selected: items(2, 3) and + // attributes(5, 6, 7). plain(4) and ignored(8) are skipped. + ASSERT_EQ(target_column_ids, (std::vector{2, 3, 5, 6, 7})); +} + +TEST_F(OrcFileBatchReaderTest, TestCollectTargetColumnIdsRejectsPartialListAndMapProjection) { + { + std::unique_ptr<::orc::Type> src_type = + ::orc::Type::buildTypeFromString("struct>>"); + std::unique_ptr<::orc::Type> target_type = + ::orc::Type::buildTypeFromString("struct>>"); + std::vector target_column_ids; + ASSERT_NOK_WITH_MSG(OrcFileBatchReader::CollectTargetColumnIds( + src_type.get(), target_type.get(), &target_column_ids), + "type mismatch"); + ASSERT_TRUE(target_column_ids.empty()); + } + { + std::unique_ptr<::orc::Type> src_type = ::orc::Type::buildTypeFromString( + "struct>>"); + std::unique_ptr<::orc::Type> target_type = + ::orc::Type::buildTypeFromString("struct>>"); + std::vector target_column_ids; + ASSERT_NOK_WITH_MSG(OrcFileBatchReader::CollectTargetColumnIds( + src_type.get(), target_type.get(), &target_column_ids), + "type mismatch"); + ASSERT_TRUE(target_column_ids.empty()); + } +} + TEST_P(OrcFileBatchReaderTest, TestNextBatchSimple) { std::string file_name = paimon::test::GetDataDir() + "/orc/append_09.db/append_09/f1=10/bucket-1/" diff --git a/src/paimon/format/parquet/CMakeLists.txt b/src/paimon/format/parquet/CMakeLists.txt index a1a566c0..96854658 100644 --- a/src/paimon/format/parquet/CMakeLists.txt +++ b/src/paimon/format/parquet/CMakeLists.txt @@ -55,6 +55,7 @@ if(PAIMON_BUILD_TESTS) file_reader_wrapper_test.cpp page_filtered_row_group_reader_test.cpp parquet_timestamp_converter_test.cpp + parquet_vector_io_test.cpp parquet_field_id_converter_test.cpp parquet_file_batch_reader_test.cpp parquet_format_writer_test.cpp diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp index 48a4430a..fb5b6574 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper.cpp @@ -21,12 +21,14 @@ #include #include #include +#include #include "arrow/io/interfaces.h" #include "arrow/record_batch.h" #include "arrow/util/range.h" #include "fmt/format.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/format/parquet/column_index_filter.h" #include "paimon/format/parquet/page_filtered_row_group_reader.h" #include "paimon/format/parquet/parquet_format_defs.h" @@ -202,6 +204,17 @@ Status FileReaderWrapper::SeekToRow(uint64_t row_number) { current_row_group_idx_ = i; next_row_to_read_ = rg_start; + if (!reader_initialized_) { + // PrepareForReading (first Next()) will build batch_reader_, so just + // record the seeked start for it. Building batch_reader_ here would be + // discarded by PrepareForReading, and the arrow GetRecordBatchReader + // eagerly reads every column chunk, so building twice doubles the + // requested bytes. + pending_start_idx_ = i; + batch_reader_.reset(); + return Status::OK(); + } + // Rebuild batch_reader_ for non-page-filtered RGs at/after seek position. std::vector fully_matched_indices; for (uint64_t j = i; j < target_row_groups_.size(); j++) { @@ -221,6 +234,12 @@ Status FileReaderWrapper::SeekToRow(uint64_t row_number) { } next_row_to_read_ = num_rows_; current_row_group_idx_ = target_row_groups_.size(); + if (!reader_initialized_) { + // Seek past the last row group before initialization: the deferred + // PrepareForReading must start at EOF as well. + pending_start_idx_ = target_row_groups_.size(); + batch_reader_.reset(); + } return Status::OK(); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::SeekToRow") @@ -376,39 +395,77 @@ Status FileReaderWrapper::PrepareForReadingLazy( target_row_groups_ = target_row_groups; target_column_indices_ = column_indices; reader_initialized_ = false; + pending_start_idx_.reset(); return Status::OK(); } -std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( - const std::vector& column_indices) { - std::vector<::arrow::io::ReadRange> ranges; - auto file_metadata = file_reader_->parquet_reader()->metadata(); - - for (const auto& trg : target_row_groups_) { - if (trg.IsExcludedByReadRange()) continue; - - if (trg.IsPartiallyMatched()) { - // Page-filtered RGs: only matching page byte ranges. - auto row_group_page_index_reader = GetRowGroupPageIndexReader(trg.GetRowGroupIndex()); - auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - trg, column_indices, row_group_page_index_reader, file_reader_->parquet_reader()); - ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()), - std::make_move_iterator(page_ranges.end())); - } else { - // Fully-matched RGs: entire column chunk ranges. - auto rg_metadata = file_metadata->RowGroup(trg.GetRowGroupIndex()); - for (int32_t col_idx : column_indices) { - auto col_chunk = rg_metadata->ColumnChunk(col_idx); - int64_t offset = col_chunk->data_page_offset(); - if (col_chunk->has_dictionary_page() && col_chunk->dictionary_page_offset() > 0 && - offset > col_chunk->dictionary_page_offset()) { - offset = col_chunk->dictionary_page_offset(); +Result> FileReaderWrapper::CollectPreBufferRanges( + const std::vector& column_indices, uint64_t start_idx) { + return DoCollectPreBufferRanges(column_indices, /*skip_read_range_excluded=*/true, start_idx); +} + +Result> FileReaderWrapper::DoCollectPreBufferRanges( + const std::vector& column_indices, bool skip_read_range_excluded, uint64_t start_idx) { + try { + std::vector<::arrow::io::ReadRange> ranges; + auto file_metadata = file_reader_->parquet_reader()->metadata(); + + for (uint64_t idx = start_idx; idx < target_row_groups_.size(); idx++) { + const auto& trg = target_row_groups_[idx]; + if (skip_read_range_excluded && trg.IsExcludedByReadRange()) { + continue; + } + + if (trg.IsPartiallyMatched()) { + // Page-filtered RGs: only matching page byte ranges. + auto row_group_page_index_reader = + GetRowGroupPageIndexReader(trg.GetRowGroupIndex()); + auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( + trg, column_indices, row_group_page_index_reader, + file_reader_->parquet_reader()); + ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()), + std::make_move_iterator(page_ranges.end())); + } else { + // Fully-matched RGs: entire column chunk ranges. + auto rg_metadata = file_metadata->RowGroup(trg.GetRowGroupIndex()); + for (int32_t col_idx : column_indices) { + auto col_chunk = rg_metadata->ColumnChunk(col_idx); + int64_t offset = col_chunk->data_page_offset(); + if (col_chunk->has_dictionary_page() && + col_chunk->dictionary_page_offset() > 0 && + offset > col_chunk->dictionary_page_offset()) { + offset = col_chunk->dictionary_page_offset(); + } + ranges.push_back({offset, col_chunk->total_compressed_size()}); } - ranges.push_back({offset, col_chunk->total_compressed_size()}); } } + return ranges; + } + PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::DoCollectPreBufferRanges") +} + +Result>> FileReaderWrapper::GetPreBufferRanges() { + PAIMON_ASSIGN_OR_RAISE(std::vector<::arrow::io::ReadRange> ranges, + DoCollectPreBufferRanges(target_column_indices_, + /*skip_read_range_excluded=*/false, + /*start_idx=*/0)); + std::vector> pre_buffer_ranges; + pre_buffer_ranges.reserve(ranges.size()); + for (const auto& range : ranges) { + // Ranges come from signed parquet metadata; a corrupt footer may hold negative or + // overflowing values. Validate before converting to uint64_t, since downstream + // range coalescing does unchecked offset + length arithmetic on them. + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(range.offset, "pre-buffer range offset")); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(range.length, "pre-buffer range length")); + if (range.offset > std::numeric_limits::max() - range.length) { + return Status::Invalid(fmt::format("pre-buffer range overflows: offset={}, length={}", + range.offset, range.length)); + } + pre_buffer_ranges.emplace_back(static_cast(range.offset), + static_cast(range.length)); } - return ranges; + return pre_buffer_ranges; } void FileReaderWrapper::DispatchPreBuffer(std::vector<::arrow::io::ReadRange> ranges) { @@ -429,10 +486,24 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t target_row_groups_ = target_row_groups; target_column_indices_ = column_indices; + // Find the first row group to read: skip read-range-excluded ones, and honor a + // seek issued while the reader was still uninitialized (SeekToRow defers reader + // construction to here). + uint64_t first_active_idx = 0; + while (first_active_idx < target_row_groups_.size() && + target_row_groups_[first_active_idx].IsExcludedByReadRange()) { + first_active_idx++; + } + if (pending_start_idx_.has_value()) { + first_active_idx = std::max(first_active_idx, pending_start_idx_.value()); + pending_start_idx_.reset(); + } + // Partition into fully-matched and page-filtered row groups, skipping excluded ones. std::vector fully_matched_row_groups; uint64_t active_count = 0; - for (const auto& trg : target_row_groups_) { + for (uint64_t i = first_active_idx; i < target_row_groups_.size(); i++) { + const auto& trg = target_row_groups_[i]; if (trg.IsExcludedByReadRange()) { continue; } @@ -462,16 +533,12 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t // When page-filtered RGs exist, issue a single PreBuffer covering both kinds. // Otherwise GetRecordBatchReader already issued PreBuffer internally. if (has_partially_matched) { - auto all_ranges = CollectPreBufferRanges(column_indices); + PAIMON_ASSIGN_OR_RAISE(std::vector<::arrow::io::ReadRange> all_ranges, + CollectPreBufferRanges(column_indices, first_active_idx)); DispatchPreBuffer(std::move(all_ranges)); } - // Reset read state. Find the first non-excluded row group. - uint64_t first_active_idx = 0; - while (first_active_idx < target_row_groups_.size() && - target_row_groups_[first_active_idx].IsExcludedByReadRange()) { - first_active_idx++; - } + // Reset read state to the first row group that will be read. if (first_active_idx >= target_row_groups_.size()) { next_row_to_read_ = num_rows_; } else { @@ -489,6 +556,8 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t Status FileReaderWrapper::ApplyReadRanges( const std::vector>& read_ranges) { + // A read-range change invalidates any seek recorded before initialization. + pending_start_idx_.reset(); if (read_ranges.empty()) { for (auto& trg : target_row_groups_) { trg.SetExcludedByReadRange(true); diff --git a/src/paimon/format/parquet/file_reader_wrapper.h b/src/paimon/format/parquet/file_reader_wrapper.h index 78a07c79..02e6ae5b 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.h +++ b/src/paimon/format/parquet/file_reader_wrapper.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -63,6 +64,9 @@ class FileReaderWrapper { /// Seek to the specified row number. /// @param row_number The row to seek to (must be at a row group boundary). + /// When the reader is not yet initialized (before the first Next()), the reader + /// construction is deferred to PrepareForReading to avoid building a batch reader + /// that would be immediately discarded. Status SeekToRow(uint64_t row_number); /// Read the next batch of rows. @@ -147,6 +151,13 @@ class FileReaderWrapper { std::shared_ptr<::parquet::RowGroupPageIndexReader> GetRowGroupPageIndexReader( int32_t row_group_index); + /// Compute the (offset, length) byte ranges required by the current target row groups + /// and columns. Unlike the arrow-internal PreBuffer path, this covers row groups that + /// are excluded by read-range dispatch as well, because the shared prefetch cache must + /// serve data consumed by all sub-readers. Relies only on file metadata, so it is safe + /// to call before the lazy reader initialization. + Result>> GetPreBufferRanges(); + private: FileReaderWrapper(std::unique_ptr<::parquet::arrow::FileReader>&& file_reader, const std::vector>& all_row_group_ranges, @@ -165,9 +176,21 @@ class FileReaderWrapper { /// Read next batch from the fully-matched batch_reader_. Returns nullptr when exhausted. Result> NextFullyMatched(); - /// Collect all byte ranges that need pre-buffering (page-filtered + fully-matched). - std::vector<::arrow::io::ReadRange> CollectPreBufferRanges( - const std::vector& column_indices); + /// Collect all byte ranges that need pre-buffering (page-filtered + fully-matched), + /// skipping row groups excluded by ApplyReadRanges and row groups before start_idx + /// (already skipped by a deferred seek). + Result> CollectPreBufferRanges( + const std::vector& column_indices, uint64_t start_idx); + + /// Core byte-range collection shared by CollectPreBufferRanges and GetPreBufferRanges. + /// When skip_read_range_excluded is true, row groups excluded by ApplyReadRanges are + /// skipped (arrow-internal PreBuffer for this reader); when false, they are included + /// (shared prefetch cache covering all sub-readers). Ranges before start_idx are + /// never collected. Metadata and page index lookups throw on malformed files or IO + /// failures, so the exceptions are converted to a Status here. + Result> DoCollectPreBufferRanges( + const std::vector& column_indices, bool skip_read_range_excluded, + uint64_t start_idx); /// Dispatch a single PreBufferRanges call with merged ranges. void DispatchPreBuffer(std::vector<::arrow::io::ReadRange> ranges); @@ -186,6 +209,10 @@ class FileReaderWrapper { uint64_t previous_first_row_ = std::numeric_limits::max(); uint64_t current_row_group_idx_ = 0; bool reader_initialized_ = false; + // Target index recorded by SeekToRow when the reader was still uninitialized; + // consumed by PrepareForReading so the deferred initialization starts at the seeked + // position instead of rebuilding readers twice. + std::optional pending_start_idx_; // Streaming reader for the currently-active page-filtered row group. Created lazily // on the first Next() call into a page-filtered RG, drained batch-by-batch, then reset diff --git a/src/paimon/format/parquet/file_reader_wrapper_test.cpp b/src/paimon/format/parquet/file_reader_wrapper_test.cpp index aaef711e..41dcde19 100644 --- a/src/paimon/format/parquet/file_reader_wrapper_test.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper_test.cpp @@ -18,8 +18,13 @@ #include "paimon/format/parquet/file_reader_wrapper.h" +#include +#include +#include #include +#include #include +#include #include "arrow/api.h" #include "arrow/array/builder_binary.h" @@ -52,6 +57,64 @@ class Array; namespace paimon::parquet::test { +// Tracks positional reads (Read at offset / ReadAsync) issued through the stream. +class ReadTrackingInputStream : public InputStream { + public: + explicit ReadTrackingInputStream(std::shared_ptr input) + : input_(std::move(input)) {} + + Status Seek(int64_t offset, SeekOrigin origin) override { + return input_->Seek(offset, origin); + } + + Result GetPos() const override { + return input_->GetPos(); + } + + Result Read(char* buffer, int64_t size) override { + return input_->Read(buffer, size); + } + + Result Read(char* buffer, int64_t size, int64_t offset) override { + RecordPositionalRead(offset, size); + return input_->Read(buffer, size, offset); + } + + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + RecordPositionalRead(offset, size); + input_->ReadAsync(buffer, size, offset, std::move(callback)); + } + + Result GetUri() const override { + return input_->GetUri(); + } + + Result Length() const override { + return input_->Length(); + } + + Status Close() override { + return input_->Close(); + } + + int64_t GetPositionalReadBytes() const { + std::lock_guard lock(mutex_); + return positional_read_bytes_; + } + + private: + void RecordPositionalRead(int64_t offset, int64_t size) { + (void)offset; + std::lock_guard lock(mutex_); + positional_read_bytes_ += size; + } + + std::shared_ptr input_; + mutable std::mutex mutex_; + int64_t positional_read_bytes_ = 0; +}; + class FileReaderWrapperTest : public ::testing::Test { public: void SetUp() override { @@ -121,8 +184,14 @@ class FileReaderWrapperTest : public ::testing::Test { Result> PrepareReaderWrapper( const std::string& file_path, int64_t wrapper_batch_size = 0) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, fs_->Open(file_path)); + return PrepareReaderWrapperOnStream(std::move(in), wrapper_batch_size); + } + + Result> PrepareReaderWrapperOnStream( + std::shared_ptr in, int64_t wrapper_batch_size = 0) { PAIMON_ASSIGN_OR_RAISE(int64_t file_length, in->Length()); - auto input_stream = std::make_unique(in, file_length, arrow_pool_); + auto input_stream = + std::make_unique(std::move(in), file_length, arrow_pool_); ::parquet::arrow::FileReaderBuilder file_reader_builder; ::parquet::ReaderProperties reader_properties; reader_properties.enable_buffered_stream(); @@ -250,6 +319,52 @@ TEST_F(FileReaderWrapperTest, Simple) { ASSERT_EQ(5500, reader_wrapper->GetPreviousBatchFirstRowNumber().value()); } +/// The prefetch framework always issues SeekToRow right after SetReadRanges, while the +/// wrapper is still uninitialized (before the first Next()). That seek must not build +/// the arrow batch reader eagerly: building it once in SeekToRow and again in +/// PrepareForReading makes the arrow reader request every column chunk twice (2x read +/// amplification). The deferred construction must also honor the seeked start position. +TEST_F(FileReaderWrapperTest, SeekBeforeInitIssuesNoReadsAndStartsAtSeekPosition) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "seek_before_init.parquet"); + PrepareParquetFile(file_path, /*row_count=*/5500); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); + auto tracking_stream = std::make_shared(std::move(in)); + auto* tracking = tracking_stream.get(); + ASSERT_OK_AND_ASSIGN(auto reader_wrapper, + PrepareReaderWrapperOnStream(std::move(tracking_stream))); + ASSERT_EQ(6, reader_wrapper->GetNumberOfRowGroups()); + + // Baseline: only metadata reads happened so far (footer etc. during Open/Build). + int64_t baseline_read_bytes = tracking->GetPositionalReadBytes(); + + // Seek to the start of RG2 while still uninitialized. This must only record the + // position, not build a batch reader that would eagerly read column chunks. + ASSERT_OK(reader_wrapper->SeekToRow(2000)); + ASSERT_EQ(2000, reader_wrapper->GetNextRowToRead()); + ASSERT_EQ(baseline_read_bytes, tracking->GetPositionalReadBytes()) + << "SeekToRow before initialization issued eager column chunk reads; the deferred " + "PrepareForReader would build a second reader and read everything twice"; + + // The first Next() performs the single deferred initialization at the seeked position. + int64_t total_rows = 0; + bool checked_first_batch = false; + while (true) { + ASSERT_OK_AND_ASSIGN(auto batch, reader_wrapper->Next()); + if (!batch) { + break; + } + if (!checked_first_batch) { + ASSERT_EQ(2000, reader_wrapper->GetPreviousBatchFirstRowNumber().value()); + checked_first_batch = true; + } + total_rows += batch->num_rows(); + } + // RG2..RG5 cover rows [2000, 5500). + ASSERT_EQ(3500, total_rows); + ASSERT_EQ(5500, reader_wrapper->GetNextRowToRead()); +} + /// Regression: when batch_size_ is 0 (the default) and a row group is consumed via /// the page-filtered streaming path, we must not pass 0 to TableBatchReader::set_chunksize /// — that would make ReadNext spin forever on zero-row batches. The wrapper now @@ -569,4 +684,218 @@ TEST_F(FileReaderWrapperTest, PrepareForReading) { reader_wrapper->GetPreviousBatchFirstRowNumber().value()); } +namespace { + +// A minimal Thrift compact-protocol walker, just enough to locate and corrupt one +// i64 field inside a Parquet footer. Only the field types that may appear in an +// unencrypted footer are supported; anything else makes the walk fail. +class CompactThriftFooter { + public: + CompactThriftFooter(std::string* data, size_t pos) : data_(data), pos_(pos) {} + + size_t pos() const { + return pos_; + } + + // Advance to the struct field with the given id, checking it has the expected + // type, and leave the cursor at the start of its value. + bool SeekField(int32_t target_id, uint8_t target_type) { + int32_t field_id = 0; + while (true) { + uint8_t type = 0; + if (!NextField(&field_id, &type)) return false; + if (field_id == 0) return false; // STOP without finding the field + if (field_id == target_id) return type == target_type; + if (!SkipValue(type)) return false; + } + } + + // Enter the body of the first element of the list at the cursor; the element + // must be a struct. + bool EnterFirstListElement() { + uint64_t header = 0; + if (!ReadByte(&header)) return false; + uint64_t size = (header >> 4) & 0x0F; + if ((header & 0x0F) != kStruct) return false; + if (size == 15 && !ReadVarint(&size)) return false; + return size > 0; // Cursor is now at the first element's struct body. + } + + static constexpr uint8_t kI64 = 6; + static constexpr uint8_t kList = 9; + static constexpr uint8_t kStruct = 12; + + private: + static constexpr uint8_t kBoolTrue = 1; + static constexpr uint8_t kBoolFalse = 2; + static constexpr uint8_t kByte = 3; + static constexpr uint8_t kI16 = 4; + static constexpr uint8_t kI32 = 5; + static constexpr uint8_t kDouble = 7; + static constexpr uint8_t kBinary = 8; + static constexpr uint8_t kSet = 10; + static constexpr uint8_t kMap = 11; + + bool ReadByte(uint64_t* out) { + if (pos_ >= data_->size()) return false; + *out = static_cast((*data_)[pos_++]); + return true; + } + + bool ReadVarint(uint64_t* out) { + *out = 0; + for (int shift = 0; shift < 64; shift += 7) { + uint64_t byte = 0; + if (!ReadByte(&byte)) return false; + *out |= (byte & 0x7F) << shift; + if ((byte & 0x80) == 0) return true; + } + return false; // Varints longer than 10 bytes are malformed. + } + + bool NextField(int32_t* field_id, uint8_t* type) { + uint64_t header = 0; + if (!ReadByte(&header)) return false; + if (header == 0) { + *field_id = 0; // STOP + return true; + } + *type = header & 0x0F; + uint64_t delta = (header >> 4) & 0x0F; + if (delta != 0) { + *field_id += static_cast(delta); + return true; + } + uint64_t zigzag = 0; + if (!ReadVarint(&zigzag)) return false; + *field_id = static_cast((zigzag >> 1) ^ -(zigzag & 1)); + return true; + } + + bool SkipValue(uint8_t type) { + switch (type) { + case kBoolTrue: + case kBoolFalse: + return true; // The value is encoded in the field header itself. + case kByte: { + uint64_t unused = 0; + return ReadByte(&unused); + } + case kI16: + case kI32: + case kI64: { + uint64_t unused = 0; + return ReadVarint(&unused); + } + case kDouble: + if (pos_ + 8 > data_->size()) return false; + pos_ += 8; + return true; + case kBinary: { + uint64_t length = 0; + if (!ReadVarint(&length)) return false; + if (pos_ + length > data_->size()) return false; + pos_ += length; + return true; + } + case kList: + case kSet: + return SkipCollection(); + case kMap: { + uint64_t size = 0; + if (!ReadVarint(&size)) return false; + if (size == 0) return true; + uint64_t kv_types = 0; + if (!ReadByte(&kv_types)) return false; + for (uint64_t i = 0; i < size; ++i) { + if (!SkipValue((kv_types >> 4) & 0x0F)) return false; + if (!SkipValue(kv_types & 0x0F)) return false; + } + return true; + } + case kStruct: { + int32_t field_id = 0; + while (true) { + uint8_t field_type = 0; + if (!NextField(&field_id, &field_type)) return false; + if (field_id == 0) return true; // STOP + if (!SkipValue(field_type)) return false; + } + } + default: + return false; + } + } + + bool SkipCollection() { + uint64_t header = 0; + if (!ReadByte(&header)) return false; + uint64_t size = (header >> 4) & 0x0F; + uint8_t elem_type = header & 0x0F; + if (size == 15 && !ReadVarint(&size)) return false; + for (uint64_t i = 0; i < size; ++i) { + if (elem_type == kBoolTrue || elem_type == kBoolFalse) { + uint64_t unused = 0; + if (!ReadByte(&unused)) return false; + continue; + } + if (!SkipValue(elem_type)) return false; + } + return true; + } + + std::string* data_; + size_t pos_; +}; + +} // namespace + +// A corrupt footer may carry negative column chunk offsets. GetPreBufferRanges must +// reject them instead of casting them into huge uint64_t ranges that would blow up +// the downstream range coalescing. +TEST_F(FileReaderWrapperTest, GetPreBufferRangesRejectsNegativeMetadataOffset) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "test.parquet"); + PrepareParquetFile(file_path, /*row_count=*/100); + + std::ifstream file(file_path, std::ios::binary); + std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + file.close(); + ASSERT_GT(content.size(), size_t{12}); + + // Footer layout: [thrift FileMetaData][footer length (4-byte LE)]["PAR1"]. + size_t tail = content.size(); + ASSERT_EQ("PAR1", content.substr(tail - 4)); + uint32_t footer_length = static_cast(content[tail - 8]) | + (static_cast(content[tail - 7]) << 8) | + (static_cast(content[tail - 6]) << 16) | + (static_cast(content[tail - 5]) << 24); + ASSERT_LT(static_cast(footer_length) + 8, content.size()); + size_t footer_start = tail - 8 - footer_length; + + // Walk to FileMetaData.row_groups[0].columns[0].meta_data.data_page_offset and + // flip the sign of its zigzag varint (positive -> negative, byte length kept). + CompactThriftFooter footer(&content, footer_start); + ASSERT_TRUE(footer.SeekField(/*FileMetaData.row_groups=*/4, CompactThriftFooter::kList)); + ASSERT_TRUE(footer.EnterFirstListElement()); + ASSERT_TRUE(footer.SeekField(/*RowGroup.columns=*/1, CompactThriftFooter::kList)); + ASSERT_TRUE(footer.EnterFirstListElement()); + ASSERT_TRUE(footer.SeekField(/*ColumnChunk.meta_data=*/3, CompactThriftFooter::kStruct)); + ASSERT_TRUE(footer.SeekField(/*ColumnMetaData.data_page_offset=*/9, CompactThriftFooter::kI64)); + size_t offset_pos = footer.pos(); + ASSERT_EQ(0, content[offset_pos] & 0x01); // Positive value: even zigzag encoding. + content[offset_pos] |= 0x01; // Now decodes to a negative offset. + + std::string corrupt_path = PathUtil::JoinPath(dir_->Str(), "corrupt.parquet"); + std::ofstream corrupt_file(corrupt_path, std::ios::binary); + corrupt_file.write(content.data(), static_cast(content.size())); + corrupt_file.close(); + + ASSERT_OK_AND_ASSIGN(auto reader_wrapper, PrepareReaderWrapper(corrupt_path)); + ASSERT_OK(reader_wrapper->PrepareForReadingLazy( + {TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false, + /*ranges=*/RowRanges())}, + /*column_indices=*/{0, 1, 2})); + ASSERT_NOK_WITH_MSG(reader_wrapper->GetPreBufferRanges(), "pre-buffer range offset"); +} + } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp index 1b4bdd30..f20f224f 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -284,9 +284,9 @@ Status PageFilteredRowGroupReader::WaitForPreBuffer( Result> PageFilteredRowGroupReader::ReadFilteredField( const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, - int32_t row_group_index, int32_t field_index, const std::vector& column_indices, - const RowRanges& row_ranges, int64_t row_group_row_count, - ::parquet::arrow::FileReader* arrow_file_reader) { + int32_t row_group_index, int32_t field_index, + std::shared_ptr> column_indices, const RowRanges& row_ranges, + int64_t row_group_row_count, ::parquet::arrow::FileReader* arrow_file_reader) { // Factory: set a direct data page read plan on every leaf (per-leaf OffsetIndex). // The plan lets Arrow jump over unselected page headers as well as page bodies. auto factory = @@ -397,12 +397,14 @@ Result> PageFilteredRowGroupReader::Re std::vector> result_arrays; result_arrays.reserve(field_indices.size()); + std::shared_ptr> col_indices_set = + std::make_shared>(column_indices.begin(), column_indices.end()); // TODO(zhouhongfeng.zhf): This loop could be parallelized. for (int field_idx : field_indices) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr chunked_array, ReadFilteredField(row_group_page_index_reader, row_group_index, field_idx, - column_indices, row_ranges, row_group_row_count, arrow_file_reader)); + col_indices_set, row_ranges, row_group_row_count, arrow_file_reader)); if (chunked_array->length() != expected_rows) { return Status::Invalid( diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.h b/src/paimon/format/parquet/page_filtered_row_group_reader.h index 683bde71..a143ae5a 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.h +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -117,11 +118,13 @@ class PageFilteredRowGroupReader { /// Sets a direct page read plan on all leaves via factory, then drives each leaf /// independently via ResetLeaf/SkipRecords/ReadRecords using its own /// compressed_ranges. + /// `column_indices` holds `int` rather than `int32_t` because the set is + /// handed straight to Arrow's `FileReader::GetColumn` (to avoid reconstruction and deep copy) static Result> ReadFilteredField( const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, - int32_t row_group_index, int32_t field_index, const std::vector& column_indices, - const RowRanges& row_ranges, int64_t row_group_row_count, - ::parquet::arrow::FileReader* arrow_file_reader); + int32_t row_group_index, int32_t field_index, + std::shared_ptr> column_indices, const RowRanges& row_ranges, + int64_t row_group_row_count, ::parquet::arrow::FileReader* arrow_file_reader); }; } // namespace paimon::parquet diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp index 67b24ac5..daacacee 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp @@ -207,7 +207,8 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto batch_reader, ParquetFileBatchReader::Create( std::move(in_stream), options, batch_size, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), predicate, @@ -235,7 +236,8 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN( auto batch_reader, ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size, nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), predicate, bitmap)); @@ -2120,7 +2122,8 @@ TEST_F(PageFilteredRowGroupReaderTest, BitmapInvalidStrategyTest) { ASSERT_OK_AND_ASSIGN(auto batch_reader, ParquetFileBatchReader::Create( std::move(in_stream), options, 1024, nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); auto read_schema = arrow::schema({arrow::field("val", arrow::int32())}); auto c_schema = std::make_unique(); diff --git a/src/paimon/format/parquet/parquet_field_id_converter.cpp b/src/paimon/format/parquet/parquet_field_id_converter.cpp index 56d36d36..adb27173 100644 --- a/src/paimon/format/parquet/parquet_field_id_converter.cpp +++ b/src/paimon/format/parquet/parquet_field_id_converter.cpp @@ -105,6 +105,12 @@ arrow::Result> ParquetFieldIdConverter::ProcessField( ProcessField(list_type->value_field(), convert_type)); auto new_type = arrow::list(new_value_field); return field->WithType(new_type)->WithMergedMetadata(updated_metadata); + } else if (type->id() == arrow::Type::FIXED_SIZE_LIST) { + auto vector_type = checked_pointer_cast(type); + ARROW_ASSIGN_OR_RAISE(auto new_value_field, + ProcessField(vector_type->value_field(), convert_type)); + auto new_type = arrow::fixed_size_list(new_value_field, vector_type->list_size()); + return field->WithType(new_type)->WithMergedMetadata(updated_metadata); } else if (type->id() == arrow::Type::MAP) { auto map_type = checked_pointer_cast(type); ARROW_ASSIGN_OR_RAISE(auto new_key_field, diff --git a/src/paimon/format/parquet/parquet_field_id_converter_test.cpp b/src/paimon/format/parquet/parquet_field_id_converter_test.cpp index 7f7c114e..d2053f12 100644 --- a/src/paimon/format/parquet/parquet_field_id_converter_test.cpp +++ b/src/paimon/format/parquet/parquet_field_id_converter_test.cpp @@ -186,7 +186,8 @@ TEST_F(ParquetFieldIdConverterTest, TestNestedType) { arrow::field("sub2", arrow::timestamp(arrow::TimeUnit::NANO)), arrow::field("sub3", arrow::decimal128(23, 5)), arrow::field("sub4", arrow::binary()), - arrow::field("sub5", arrow::binary())})))}; + arrow::field("sub5", arrow::binary())}))), + arrow::field("f3", arrow::fixed_size_list(arrow::float32(), 7))}; auto schema = arrow::schema(fields); ASSERT_OK_AND_ASSIGN( auto table_schema, @@ -211,8 +212,11 @@ TEST_F(ParquetFieldIdConverterTest, TestNestedType) { {"sub4", arrow::Type::BINARY, "16"}, {"sub5", arrow::Type::BINARY, "17"}, {"sub1", arrow::Type::DATE32, "18"}, {"sub2", arrow::Type::TIMESTAMP, "19"}, {"sub3", arrow::Type::DECIMAL128, "20"}, {"sub4", arrow::Type::BINARY, "21"}, - {"sub5", arrow::Type::BINARY, "22"}}; + {"sub5", arrow::Type::BINARY, "22"}, {"f3", arrow::Type::FIXED_SIZE_LIST, "23"}}; ASSERT_EQ(expected_field_infos, field_infos); + auto new_vector = + checked_pointer_cast(new_schema->GetFieldByName("f3")->type()); + ASSERT_EQ(new_vector->list_size(), 7); // convert to paimon.id ASSERT_OK_AND_ASSIGN(auto old_schema, ParquetFieldIdConverter::GetPaimonIdsFromParquetIds(new_schema)); @@ -220,6 +224,9 @@ TEST_F(ParquetFieldIdConverterTest, TestNestedType) { PrintFieldMetadata(old_schema, ParquetFieldIdConverter::IdConvertType::PARQUET_TO_PAIMON_ID, &old_field_infos); ASSERT_EQ(expected_field_infos, old_field_infos); + auto old_vector = + checked_pointer_cast(old_schema->GetFieldByName("f3")->type()); + ASSERT_EQ(old_vector->list_size(), 7); } } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index c0cd40e1..0c9d065e 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include "arrow/acero/options.h" @@ -40,6 +41,7 @@ #include "fmt/format.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/string_utils.h" @@ -113,6 +115,12 @@ bool HasSameNestedProjectionShape(const std::shared_ptr& read_t const auto& file_list = static_cast(*file_type); return HasSameNestedProjectionShape(read_list.value_type(), file_list.value_type()); } + case arrow::Type::FIXED_SIZE_LIST: { + const auto& read_vector = checked_cast(*read_type); + const auto& file_vector = checked_cast(*file_type); + return read_vector.list_size() == file_vector.list_size() && + HasSameNestedProjectionShape(read_vector.value_type(), file_vector.value_type()); + } case arrow::Type::MAP: { const auto& read_map = static_cast(*read_type); const auto& file_map = static_cast(*file_type); @@ -123,6 +131,18 @@ bool HasSameNestedProjectionShape(const std::shared_ptr& read_t return false; } } + +// Resolve whether parquet-level pre-buffering should be enabled. When the framework +// provides runtime hints, they describe the authoritative state of this read: once the +// shared read-ahead cache takes over prefetching, disable parquet's own pre-buffering so +// the same byte ranges are not buffered twice. Without hints, fall back to the option. +Result ResolvePreBufferEnabled(const std::map& options, + const std::optional& hints) { + if (hints.has_value() && hints->prefetch_enabled && hints->read_ahead_cache_enabled) { + return false; + } + return OptionsUtils::GetValueFromMap(options, PARQUET_READ_ENABLE_PRE_BUFFER, true); +} } // namespace ParquetFileBatchReader::ParquetFileBatchReader( @@ -143,14 +163,14 @@ Result> ParquetFileBatchReader::Create( const std::map& options, int32_t batch_size, std::shared_ptr<::parquet::FileMetaData> file_metadata, std::shared_ptr> storage_read_bytes, - const std::shared_ptr& pool) { + const std::shared_ptr& pool, const std::optional& hints) { try { assert(input_stream); PAIMON_ASSIGN_OR_RAISE(::parquet::ReaderProperties reader_properties, - CreateReaderProperties(pool, options)); + CreateReaderProperties(pool, options, hints)); PAIMON_ASSIGN_OR_RAISE(::parquet::ArrowReaderProperties arrow_reader_properties, - CreateArrowReaderProperties(pool, options, batch_size)); + CreateArrowReaderProperties(pool, options, batch_size, hints)); ::parquet::arrow::FileReaderBuilder file_reader_builder; PAIMON_RETURN_NOT_OK_FROM_ARROW( @@ -635,14 +655,16 @@ Result>> ParquetFileBatchReader::GenRe PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetFileBatchReader::GenReadRanges") } +Result>> ParquetFileBatchReader::PreBufferRange() { + return reader_->GetPreBufferRanges(); +} + Result<::parquet::ReaderProperties> ParquetFileBatchReader::CreateReaderProperties( const std::shared_ptr& pool, - const std::map& options) { + const std::map& options, const std::optional& hints) { ::parquet::ReaderProperties reader_properties; // TODO(jinli.zjw): set more ReaderProperties (compare with java) - PAIMON_ASSIGN_OR_RAISE( - bool enable_pre_buffer, - OptionsUtils::GetValueFromMap(options, PARQUET_READ_ENABLE_PRE_BUFFER, true)); + PAIMON_ASSIGN_OR_RAISE(bool enable_pre_buffer, ResolvePreBufferEnabled(options, hints)); if (enable_pre_buffer) { reader_properties.enable_buffered_stream(); } else { @@ -653,7 +675,8 @@ Result<::parquet::ReaderProperties> ParquetFileBatchReader::CreateReaderProperti Result<::parquet::ArrowReaderProperties> ParquetFileBatchReader::CreateArrowReaderProperties( const std::shared_ptr& pool, - const std::map& options, int32_t batch_size) { + const std::map& options, int32_t batch_size, + const std::optional& hints) { PAIMON_ASSIGN_OR_RAISE( uint32_t executor_thread_count, OptionsUtils::GetValueFromMap(options, PARQUET_READ_EXECUTOR_THREAD_COUNT, @@ -661,9 +684,7 @@ Result<::parquet::ArrowReaderProperties> ParquetFileBatchReader::CreateArrowRead ::parquet::ArrowReaderProperties arrow_reader_props; // TODO(jinli.zjw): set more ArrowReaderProperties (compare with java) - PAIMON_ASSIGN_OR_RAISE( - bool enable_pre_buffer, - OptionsUtils::GetValueFromMap(options, PARQUET_READ_ENABLE_PRE_BUFFER, true)); + PAIMON_ASSIGN_OR_RAISE(bool enable_pre_buffer, ResolvePreBufferEnabled(options, hints)); arrow_reader_props.set_pre_buffer(enable_pre_buffer); arrow_reader_props.set_batch_size(static_cast(batch_size)); if (executor_thread_count != 0) { @@ -740,6 +761,16 @@ Status ParquetFileBatchReader::CollectLeafIndices(const std::shared_ptr(*file_type); PAIMON_RETURN_NOT_OK(CollectLeafIndices(read_list.value_type(), file_list.value_type(), leaf_index, indices)); + } else if (file_type->id() == arrow::Type::FIXED_SIZE_LIST) { + if (!HasSameNestedProjectionShape(read_type, file_type)) { + return Status::Invalid(fmt::format( + "Parquet does not support partial projection inside list/map: src {} vs target {}", + file_type->ToString(), read_type->ToString())); + } + const auto& read_vector = checked_cast(*read_type); + const auto& file_vector = checked_cast(*file_type); + PAIMON_RETURN_NOT_OK(CollectLeafIndices(read_vector.value_type(), file_vector.value_type(), + leaf_index, indices)); } else if (file_type->id() == arrow::Type::MAP) { if (!HasSameNestedProjectionShape(read_type, file_type)) { return Status::Invalid(fmt::format( @@ -761,8 +792,7 @@ Status ParquetFileBatchReader::CollectLeafIndices(const std::shared_ptr& file_type, int32_t* leaf_index) { - if (file_type->id() == arrow::Type::STRUCT || file_type->id() == arrow::Type::LIST || - file_type->id() == arrow::Type::MAP) { + if (ArrowSchemaValidator::IsNestedType(file_type)) { for (int32_t i = 0; i < file_type->num_fields(); i++) { SkipLeafIndices(file_type->field(i)->type(), leaf_index); } diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 7e5e9afa..daa18f04 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -44,6 +44,7 @@ #include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/format/parquet/row_ranges.h" #include "paimon/format/parquet/target_row_group.h" +#include "paimon/format/read_hints.h" #include "paimon/logging.h" #include "paimon/reader/prefetch_file_batch_reader.h" #include "paimon/result.h" @@ -76,11 +77,11 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { const std::map& options, int32_t batch_size, std::shared_ptr<::parquet::FileMetaData> file_metadata, std::shared_ptr> storage_read_bytes, - const std::shared_ptr& pool); + const std::shared_ptr& pool, const std::optional& hints); static Result<::parquet::ReaderProperties> CreateReaderProperties( const std::shared_ptr& pool, - const std::map& options); + const std::map& options, const std::optional& hints); // For timestamp type, we return the schema stored in file, e.g., second in parquet file will // store as milli. @@ -102,6 +103,8 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { Result>> GenReadRanges( bool* need_prefetch) const override; + Result>> PreBufferRange() override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override { if (row_mapping_.empty()) { PAIMON_ASSIGN_OR_RAISE(uint64_t previous_first_row, @@ -162,7 +165,8 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { static Result<::parquet::ArrowReaderProperties> CreateArrowReaderProperties( const std::shared_ptr& pool, - const std::map& options, int32_t batch_size); + const std::map& options, int32_t batch_size, + const std::optional& hints); static void FlattenSchema(const std::shared_ptr& type, int32_t* index, std::vector* index_vector) { diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index bcd4af21..1409f24b 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -18,12 +18,15 @@ #include "paimon/format/parquet/parquet_file_batch_reader.h" +#include #include #include #include #include #include #include +#include +#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -37,6 +40,8 @@ #include "arrow/ipc/api.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/io/cache_input_stream.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" #include "paimon/common/utils/arrow/arrow_utils.h" @@ -44,11 +49,13 @@ #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/read_ahead_cache.h" #include "paimon/defs.h" #include "paimon/format/parquet/parquet_field_id_converter.h" #include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/format/parquet/parquet_format_writer.h" #include "paimon/format/parquet/parquet_reader_builder.h" +#include "paimon/format/read_hints.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" @@ -60,6 +67,7 @@ #include "paimon/testing/utils/testharness.h" #include "paimon/testing/utils/timezone_guard.h" #include "paimon/utils/roaring_bitmap32.h" +#include "parquet/file_reader.h" #include "parquet/properties.h" namespace paimon { @@ -225,7 +233,8 @@ class ParquetFileBatchReaderTest : public ::testing::Test, EXPECT_OK_AND_ASSIGN(auto parquet_batch_reader, ParquetFileBatchReader::Create( std::move(in_stream), options, batch_size, - /*file_metadata=*/nullptr, std::move(storage_read_bytes), pool_)); + /*file_metadata=*/nullptr, std::move(storage_read_bytes), pool_, + /*hints=*/std::nullopt)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); EXPECT_TRUE(arrow_status.ok()); @@ -392,7 +401,8 @@ TEST_F(ParquetFileBatchReaderTest, TestSetReadSchema) { ASSERT_OK_AND_ASSIGN(auto parquet_batch_reader, ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size_, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, pool_)); + /*storage_read_bytes=*/nullptr, pool_, + /*hints=*/std::nullopt)); // test GetFileSchema() ASSERT_OK_AND_ASSIGN(auto c_file_schema, parquet_batch_reader->GetFileSchema()); auto arrow_file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); @@ -838,8 +848,9 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateReaderProperties) { { // test default options std::map options; - ASSERT_OK_AND_ASSIGN(auto reader_properties, - ParquetFileBatchReader::CreateReaderProperties(pool_, options)); + ASSERT_OK_AND_ASSIGN(auto reader_properties, ParquetFileBatchReader::CreateReaderProperties( + pool_, options, + /*hints=*/std::nullopt)); ASSERT_EQ(reader_properties.is_buffered_stream_enabled(), true); } } @@ -851,7 +862,8 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { int32_t batch_size = 1024; ASSERT_OK_AND_ASSIGN( auto arrow_reader_properties, - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size, + /*hints=*/std::nullopt)); ASSERT_EQ(arrow_reader_properties.pre_buffer(), true); ASSERT_EQ(arrow_reader_properties.batch_size(), 1024); ASSERT_EQ(arrow_reader_properties.use_threads(), true); @@ -865,7 +877,8 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { int32_t batch_size = 1024; ASSERT_OK_AND_ASSIGN( auto arrow_reader_properties, - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size, + /*hints=*/std::nullopt)); ASSERT_EQ(arrow_reader_properties.use_threads(), false); } { @@ -873,7 +886,8 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { int32_t batch_size = 1024; ASSERT_OK_AND_ASSIGN( auto arrow_reader_properties, - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size)); + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, batch_size, + /*hints=*/std::nullopt)); ASSERT_EQ(arrow_reader_properties.use_threads(), true); ASSERT_EQ(arrow::GetCpuThreadPoolCapacity(), 6); } @@ -886,7 +900,8 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { }; ASSERT_OK_AND_ASSIGN( auto arrow_reader_properties, - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024)); + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024, + /*hints=*/std::nullopt)); const auto& cache_options = arrow_reader_properties.cache_options(); ASSERT_TRUE(cache_options.lazy); ASSERT_EQ(cache_options.prefetch_limit, 2); @@ -898,7 +913,8 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { {PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT, "-1"}, }; ASSERT_NOK_WITH_MSG( - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024), + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024, + /*hints=*/std::nullopt), "parquet.read.cache-option.hole-size-limit must be non-negative"); } { @@ -907,12 +923,79 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) { {PARQUET_READ_CACHE_OPTION_RANGE_SIZE_LIMIT, "1048576"}, }; ASSERT_NOK_WITH_MSG( - ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024), + ParquetFileBatchReader::CreateArrowReaderProperties(pool_, options, 1024, + /*hints=*/std::nullopt), "parquet.read.cache-option.range-size-limit must be greater than " "parquet.read.cache-option.hole-size-limit"); } } +TEST_F(ParquetFileBatchReaderTest, TestPreBufferReadHints) { + const int32_t batch_size = 1024; + // When both framework prefetch and the shared read-ahead cache are active, parquet's own + // pre-buffering must be disabled even if the option explicitly enables it (runtime state + // takes precedence over the option). + { + std::map options = {{PARQUET_READ_ENABLE_PRE_BUFFER, "true"}}; + ReadHints hints; + hints.prefetch_enabled = true; + hints.read_ahead_cache_enabled = true; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, hints)); + ASSERT_EQ(arrow_props.pre_buffer(), false); + ASSERT_OK_AND_ASSIGN(auto reader_props, + ParquetFileBatchReader::CreateReaderProperties(pool_, options, hints)); + ASSERT_EQ(reader_props.is_buffered_stream_enabled(), false); + } + // Only prefetch enabled (cache disabled): fall back to the option, which defaults to true. + { + std::map options; + ReadHints hints; + hints.prefetch_enabled = true; + hints.read_ahead_cache_enabled = false; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, hints)); + ASSERT_EQ(arrow_props.pre_buffer(), true); + } + // Only cache enabled (prefetch disabled): fall back to the option, which defaults to true. + { + std::map options; + ReadHints hints; + hints.prefetch_enabled = false; + hints.read_ahead_cache_enabled = true; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, hints)); + ASSERT_EQ(arrow_props.pre_buffer(), true); + } + // Neither active and the option explicitly disabled: honor the option. + { + std::map options = {{PARQUET_READ_ENABLE_PRE_BUFFER, "false"}}; + ReadHints hints; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, hints)); + ASSERT_EQ(arrow_props.pre_buffer(), false); + ASSERT_OK_AND_ASSIGN(auto reader_props, + ParquetFileBatchReader::CreateReaderProperties(pool_, options, hints)); + ASSERT_EQ(reader_props.is_buffered_stream_enabled(), false); + } + // Neither active and no option: default to enabled. + { + std::map options; + ReadHints hints; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, hints)); + ASSERT_EQ(arrow_props.pre_buffer(), true); + } + // No hints provided at all (builder used without WithReadHints): fall back to the option. + { + std::map options = {{PARQUET_READ_ENABLE_PRE_BUFFER, "false"}}; + ASSERT_OK_AND_ASSIGN(auto arrow_props, ParquetFileBatchReader::CreateArrowReaderProperties( + pool_, options, batch_size, + /*hints=*/std::nullopt)); + ASSERT_EQ(arrow_props.pre_buffer(), false); + } +} + TEST_F(ParquetFileBatchReaderTest, TestBitmapRowGroupPushDownWithMultiRowGroups) { arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; auto arrow_type = arrow::struct_(fields); @@ -1472,4 +1555,179 @@ TEST_F(ParquetFileBatchReaderTest, TestRowMappingSetReadSchemaTwice) { ASSERT_EQ(parquet_batch_reader->GetPreviousBatchFileRowId(2).value(), 5); } +// The shared prefetch cache is initialized from a single sub-reader's PreBufferRange(), +// so the returned byte ranges must cover the column chunks of all target row groups, +// including those excluded by read-range dispatch. +TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeCoversDispatchExcludedRowGroups) { + arrow::FieldVector fields = {arrow::field("c0", arrow::int32()), + arrow::field("c1", arrow::int32())}; + arrow::Int32Builder c0_builder; + arrow::Int32Builder c1_builder; + ASSERT_TRUE(c0_builder.Reserve(20).ok()); + ASSERT_TRUE(c1_builder.Reserve(20).ok()); + for (int32_t i = 0; i < 20; ++i) { + c0_builder.UnsafeAppend(i); + c1_builder.UnsafeAppend(i % 4); + } + auto c0_array = c0_builder.Finish().ValueOrDie(); + auto c1_array = c1_builder.Finish().ValueOrDie(); + auto src_array = arrow::StructArray::Make({c0_array, c1_array}, fields).ValueOrDie(); + auto arrow_schema = arrow::schema(fields); + WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/10, + /*enable_dictionary=*/true, /*max_row_group_length=*/10); + + auto parquet_batch_reader = PrepareParquetFileBatchReader( + file_path_, arrow_schema, /*predicate=*/nullptr, std::nullopt, batch_size_); + + bool need_prefetch = false; + ASSERT_OK_AND_ASSIGN(auto row_group_ranges, + parquet_batch_reader->GenReadRanges(&need_prefetch)); + ASSERT_EQ(2u, row_group_ranges.size()); + + // Simulate prefetch dispatch: this sub-reader owns only the first row group. + ASSERT_OK(parquet_batch_reader->SetReadRanges({row_group_ranges[0]})); + + ASSERT_OK_AND_ASSIGN(auto pre_buffer_ranges, parquet_batch_reader->PreBufferRange()); + + // Expected: column chunk ranges of every row group, including the dispatch-excluded one. + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path_)); + ASSERT_OK_AND_ASSIGN(int64_t file_length, in->Length()); + auto adapter = std::make_shared(std::move(in), file_length, pool_); + auto parquet_reader = ::parquet::ParquetFileReader::Open(adapter); + ASSERT_TRUE(parquet_reader); + auto file_metadata = parquet_reader->metadata(); + ASSERT_EQ(2, file_metadata->num_row_groups()); + + std::vector> expected_ranges; + for (int32_t rg = 0; rg < file_metadata->num_row_groups(); ++rg) { + auto rg_metadata = file_metadata->RowGroup(rg); + for (int32_t col = 0; col < rg_metadata->num_columns(); ++col) { + auto col_chunk = rg_metadata->ColumnChunk(col); + int64_t offset = col_chunk->data_page_offset(); + if (col_chunk->has_dictionary_page() && col_chunk->dictionary_page_offset() > 0 && + offset > col_chunk->dictionary_page_offset()) { + offset = col_chunk->dictionary_page_offset(); + } + expected_ranges.emplace_back(static_cast(offset), + static_cast(col_chunk->total_compressed_size())); + } + } + ASSERT_EQ(4u, expected_ranges.size()); + + std::vector> actual_ranges = pre_buffer_ranges; + std::sort(expected_ranges.begin(), expected_ranges.end()); + std::sort(actual_ranges.begin(), actual_ranges.end()); + ASSERT_EQ(expected_ranges, actual_ranges); +} + +// A page-index partially-matched row group should contribute page-level byte ranges +// that are strictly smaller than its full column chunk. +TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeWithPageFilteredRowGroup) { + arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; + auto src_array = MakeSequentialIntData(12); + auto arrow_schema = arrow::schema(fields); + // One row per page, three row groups of four rows each. + WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/1, + /*enable_dictionary=*/false, /*max_row_group_length=*/4, /*max_page_size=*/1); + + // Only rows 10 and 11 of RowGroup 2 match, making it partially matched; RowGroups 0 + // and 1 are excluded by the predicate entirely. + std::shared_ptr predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"f0", FieldType::INT, Literal(10)); + auto parquet_batch_reader = PrepareParquetFileBatchReader(file_path_, arrow_schema, predicate, + std::nullopt, batch_size_, + /*enable_page_level_filter=*/true); + + ASSERT_OK_AND_ASSIGN(auto pre_buffer_ranges, parquet_batch_reader->PreBufferRange()); + ASSERT_FALSE(pre_buffer_ranges.empty()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path_)); + ASSERT_OK_AND_ASSIGN(int64_t file_length, in->Length()); + auto adapter = std::make_shared(std::move(in), file_length, pool_); + auto parquet_reader = ::parquet::ParquetFileReader::Open(adapter); + ASSERT_TRUE(parquet_reader); + auto col_chunk = parquet_reader->metadata()->RowGroup(2)->ColumnChunk(0); + auto chunk_offset = static_cast(col_chunk->data_page_offset()); + uint64_t chunk_end = chunk_offset + static_cast(col_chunk->total_compressed_size()); + + uint64_t filtered_total = 0; + for (const auto& range : pre_buffer_ranges) { + ASSERT_GE(range.first, chunk_offset); + ASSERT_LE(range.first + range.second, chunk_end); + filtered_total += range.second; + } + ASSERT_LT(filtered_total, chunk_end - chunk_offset); +} + +// End-to-end: PreBufferRange() feeds the shared ReadAheadCache through CacheInputStream, +// and data reads are served from the cache. +TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeFeedsReadAheadCache) { + arrow::FieldVector fields = {arrow::field("f0", arrow::int32())}; + auto src_array = MakeSequentialIntData(20); + auto arrow_schema = arrow::schema(fields); + WriteArray(file_path_, src_array, arrow_schema, /*write_batch_size=*/10, + /*enable_dictionary=*/true, /*max_row_group_length=*/10); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr cache_stream, fs_->Open(file_path_)); + auto cache = std::make_shared(cache_stream, CacheConfig(), GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_stream, fs_->Open(file_path_)); + auto cache_input_stream = std::make_shared(std::move(reader_stream), cache); + + std::map options; + ParquetReaderBuilder builder(options, batch_size_); + builder.WithMemoryPool(GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_reader, + builder.Build(cache_input_stream)); + auto parquet_batch_reader = dynamic_cast(base_reader.get()); + ASSERT_TRUE(parquet_batch_reader); + std::unique_ptr c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*arrow_schema, c_schema.get()).ok()); + ASSERT_OK( + parquet_batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(auto pre_buffer_ranges, parquet_batch_reader->PreBufferRange()); + ASSERT_FALSE(pre_buffer_ranges.empty()); + std::vector byte_ranges; + byte_ranges.reserve(pre_buffer_ranges.size()); + for (const auto& range : pre_buffer_ranges) { + byte_ranges.emplace_back(range.first, range.second); + } + ASSERT_OK(cache->Init(std::move(byte_ranges))); + // Dispatch the prefetch immediately so every pre-buffered range is covered + // before the reads below consume them. + cache->Warmup(); + + // Baseline before draining: the Build() phase may have issued reads that no + // prefetch range can cover (e.g. footer parsing when no metadata cache is + // configured). Only the data-consumption reads below must be miss-free. + std::shared_ptr baseline_metrics = std::make_shared(); + cache->CollectMetrics(&baseline_metrics); + ASSERT_OK_AND_ASSIGN(uint64_t baseline_misses, + baseline_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_OK_AND_ASSIGN(uint64_t baseline_miss_bytes, + baseline_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); + + // Drain the file through the cache-backed stream. + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + paimon::test::ReadResultCollector::CollectResult(parquet_batch_reader)); + ASSERT_EQ(20, result->length()); + + // The cache metrics must show that the data reads were served by the cache: + // at least one hit, and no additional miss falling back to the wrapped stream. + std::shared_ptr cache_metrics = std::make_shared(); + cache->CollectMetrics(&cache_metrics); + ASSERT_OK_AND_ASSIGN(uint64_t hits, + cache_metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_GT(hits, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t hit_bytes, + cache_metrics->GetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES)); + ASSERT_GT(hit_bytes, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t misses, + cache_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_EQ(misses, baseline_misses); + ASSERT_OK_AND_ASSIGN(uint64_t miss_bytes, + cache_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); + ASSERT_EQ(miss_bytes, baseline_miss_bytes); +} + } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index 51976372..112167bf 100644 --- a/src/paimon/format/parquet/parquet_reader_builder.h +++ b/src/paimon/format/parquet/parquet_reader_builder.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -34,6 +35,7 @@ #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/parquet/parquet_file_batch_reader.h" #include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/read_hints.h" #include "paimon/format/reader_builder.h" #include "paimon/memory/memory_pool.h" #include "paimon/memory/memory_segment.h" @@ -59,6 +61,11 @@ class ParquetReaderBuilder : public ReaderBuilder { return this; } + ReaderBuilder* WithReadHints(const std::optional& hints) override { + hints_ = hints; + return this; + } + Result> Build( const std::shared_ptr& path) const override { try { @@ -78,9 +85,9 @@ class ParquetReaderBuilder : public ReaderBuilder { std::move(unique_input_stream)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<::parquet::FileMetaData> file_metadata, GetCachedParquetMetadata(input_stream, file_uri, arrow_pool)); - return ParquetFileBatchReader::Create(std::move(input_stream), options_, batch_size_, - std::move(file_metadata), - std::move(storage_read_bytes), arrow_pool); + return ParquetFileBatchReader::Create( + std::move(input_stream), options_, batch_size_, std::move(file_metadata), + std::move(storage_read_bytes), arrow_pool, hints_); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetReaderBuilder::Build") } @@ -137,7 +144,7 @@ class ParquetReaderBuilder : public ReaderBuilder { } PAIMON_ASSIGN_OR_RAISE( ::parquet::ReaderProperties reader_properties, - ParquetFileBatchReader::CreateReaderProperties(arrow_pool, options_)); + ParquetFileBatchReader::CreateReaderProperties(arrow_pool, options_, hints_)); auto cache_key = CacheKey::ForKind(file_uri, /*position=*/-1, /*length=*/-1, CacheKind::DATA_FILE_FOOTER); @@ -162,6 +169,7 @@ class ParquetReaderBuilder : public ReaderBuilder { std::shared_ptr pool_; std::map options_; std::shared_ptr cache_; + std::optional hints_; }; } // namespace paimon::parquet diff --git a/src/paimon/format/parquet/parquet_stats_extractor.cpp b/src/paimon/format/parquet/parquet_stats_extractor.cpp index 4b8f97f0..8dfe7faa 100644 --- a/src/paimon/format/parquet/parquet_stats_extractor.cpp +++ b/src/paimon/format/parquet/parquet_stats_extractor.cpp @@ -296,7 +296,9 @@ ParquetStatsExtractor::ExtractWithFileInfo(const std::shared_ptr& fi // nested type do not have parquet stats const auto& logical_type = node->logical_type(); FieldType nested_type = FieldType::UNKNOWN; - if (logical_type->is_list()) { + if (write_schema_->field(field_idx)->type()->id() == arrow::Type::FIXED_SIZE_LIST) { + nested_type = FieldType::VECTOR; + } else if (logical_type->is_list()) { nested_type = FieldType::ARRAY; } else if (logical_type->is_map()) { nested_type = FieldType::MAP; diff --git a/src/paimon/format/parquet/parquet_stats_extractor_test.cpp b/src/paimon/format/parquet/parquet_stats_extractor_test.cpp index 4dc7fbe5..65998426 100644 --- a/src/paimon/format/parquet/parquet_stats_extractor_test.cpp +++ b/src/paimon/format/parquet/parquet_stats_extractor_test.cpp @@ -62,7 +62,8 @@ class ParquetStatsExtractorTest : public ::testing::Test { void TearDown() override {} void CheckStats(const arrow::FieldVector& fields, const std::string& input, - const std::vector& expected_stats, int64_t expect_row_count) { + const std::vector& expected_stats, int64_t expect_row_count, + const std::vector& expected_types = {}) { auto arrow_schema = arrow::schema(fields); auto struct_type = arrow::struct_(fields); std::map options; @@ -95,6 +96,12 @@ class ParquetStatsExtractorTest : public ::testing::Test { for (size_t i = 0; i < expected_stats.size(); i++) { ASSERT_EQ(expected_stats[i], col_stats_vec[i]->ToString()); } + if (!expected_types.empty()) { + ASSERT_EQ(col_stats_vec.size(), expected_types.size()); + for (size_t i = 0; i < expected_types.size(); ++i) { + ASSERT_EQ(col_stats_vec[i]->GetFieldType(), expected_types[i]); + } + } auto row_count = result.second.GetRowCount(); ASSERT_EQ(row_count, expect_row_count); } @@ -237,6 +244,13 @@ TEST_F(ParquetStatsExtractorTest, TestExtractStatsComplexType) { CheckStats(fields, data_str, expected_stats_str, /*expect_row_count=*/6); } +TEST_F(ParquetStatsExtractorTest, TestExtractVectorStats) { + arrow::FieldVector fields = { + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))}; + CheckStats(fields, R"([[[1.0, 2.0, 3.0]], [null]])", {"min null, max null, null count null"}, + /*expect_row_count=*/2, {FieldType::VECTOR}); +} + TEST_F(ParquetStatsExtractorTest, TestNullForAllType) { auto timezone = DateTimeUtils::GetLocalTimezoneName(); arrow::FieldVector fields = { diff --git a/src/paimon/format/parquet/parquet_vector_io_test.cpp b/src/paimon/format/parquet/parquet_vector_io_test.cpp new file mode 100644 index 00000000..ad60caef --- /dev/null +++ b/src/paimon/format/parquet/parquet_vector_io_test.cpp @@ -0,0 +1,463 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" +#include "paimon/common/utils/arrow/arrow_output_stream_adapter.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/io/vector_file_batch_reader.h" +#include "paimon/defs.h" +#include "paimon/format/parquet/parquet_file_batch_reader.h" +#include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/parquet/parquet_format_writer.h" +#include "paimon/fs/file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "parquet/arrow/writer.h" +#include "parquet/properties.h" + +namespace paimon { +class Predicate; +} // namespace paimon + +namespace paimon::parquet::test { + +class ParquetVectorIoTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + arrow_pool_ = GetArrowPool(pool_); + dir_ = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); + fs_ = dir_->GetFileSystem(); + } + + void WriteAndCheck(const std::string& file_name, + const std::shared_ptr& write_type, + const std::shared_ptr& read_type, + const std::string& json) { + std::string file_path = dir_->Str() + "/" + file_name; + WriteWithFormatWriter(file_path, write_type, json, /*max_row_group_length=*/1024); + + std::shared_ptr file_type; + ReadFileType(file_path, &file_type); + std::shared_ptr physical_value_type = file_type->field(1)->type(); + if (physical_value_type->id() == arrow::Type::STRUCT) { + physical_value_type = physical_value_type->field(0)->type(); + } + ASSERT_EQ(physical_value_type->id(), arrow::Type::LIST); + + std::unique_ptr vector_reader; + CreateVectorReader(file_path, arrow::schema(read_type->fields()), /*predicate=*/nullptr, + /*options=*/{}, /*batch_size=*/10, &vector_reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(vector_reader.get())); + + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON(read_type, json); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + std::shared_ptr expected = std::move(expected_result).ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)) + << actual->ToString(); + } + + /// Writes the JSON rows through the Paimon Parquet writer, which stores VECTOR values as + /// Parquet LIST. + void WriteWithFormatWriter(const std::string& file_path, + const std::shared_ptr& write_type, + const std::string& json, int64_t max_row_group_length) { + arrow::Result> write_array_result = + arrow::ipc::internal::json::ArrayFromJSON(write_type, json); + ASSERT_TRUE(write_array_result.ok()) << write_array_result.status().ToString(); + std::shared_ptr write_array = std::move(write_array_result).ValueOrDie(); + auto c_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*write_array, c_array.get()).ok()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path, /*overwrite=*/false)); + ::parquet::WriterProperties::Builder properties_builder; + properties_builder.max_row_group_length(max_row_group_length); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr writer, + ParquetFormatWriter::Create(out, arrow::schema(write_type->fields()), + properties_builder.build(), + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); + ASSERT_OK(writer->AddBatch(c_array.get())); + ASSERT_OK(writer->Finish()); + ASSERT_OK(out->Close()); + } + + /// Writes `array` with the plain Arrow Parquet writer, storing the Arrow schema so that + /// FixedSizeList columns are read back as FixedSizeList, the way Paimon Rust and Python + /// writers store them. + void WriteWithArrowWriter(const std::string& file_path, + const std::shared_ptr& type, + const std::string& json) { + arrow::Result> array_result = + arrow::ipc::internal::json::ArrayFromJSON(type, json); + ASSERT_TRUE(array_result.ok()) << array_result.status().ToString(); + arrow::Result> batch_result = + arrow::RecordBatch::FromStructArray(std::move(array_result).ValueOrDie()); + ASSERT_TRUE(batch_result.ok()) << batch_result.status().ToString(); + arrow::Result> table_result = + arrow::Table::FromRecordBatches({std::move(batch_result).ValueOrDie()}); + ASSERT_TRUE(table_result.ok()) << table_result.status().ToString(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path, /*overwrite=*/false)); + auto arrow_out = std::make_shared(out); + ::parquet::WriterProperties::Builder properties_builder; + std::shared_ptr<::parquet::ArrowWriterProperties> arrow_properties = + ::parquet::ArrowWriterProperties::Builder().store_schema()->build(); + arrow::Status status = ::parquet::arrow::WriteTable( + *std::move(table_result).ValueOrDie(), arrow_pool_.get(), arrow_out, + /*chunk_size=*/1024, properties_builder.build(), arrow_properties); + ASSERT_TRUE(status.ok()) << status.ToString(); + ASSERT_OK(out->Close()); + } + + void ReadFileType(const std::string& file_path, + std::shared_ptr* file_type_out) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); + ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); + auto in_stream = std::make_shared(in, length, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + ParquetFileBatchReader::Create(std::move(in_stream), /*options=*/{}, + /*batch_size=*/10, /*file_metadata=*/nullptr, + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr c_file_schema, reader->GetFileSchema()); + arrow::Result> file_type_result = + arrow::ImportType(c_file_schema.get()); + ASSERT_TRUE(file_type_result.ok()) << file_type_result.status().ToString(); + *file_type_out = + checked_pointer_cast(std::move(file_type_result).ValueOrDie()); + } + + void CreateVectorReader(const std::string& file_path, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::map& options, int32_t batch_size, + std::unique_ptr* vector_reader_out) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file_path)); + ASSERT_OK_AND_ASSIGN(int64_t length, in->Length()); + auto in_stream = std::make_shared(in, length, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size, + /*file_metadata=*/nullptr, + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); + std::unique_ptr vector_reader = + std::make_unique(std::move(reader), pool_); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_schema.get()).ok()); + ASSERT_OK(vector_reader->SetReadSchema(c_schema.get(), predicate, + /*selection_bitmap=*/std::nullopt)); + *vector_reader_out = std::move(vector_reader); + } + + void ReadFixtureAndCheck( + const std::string& file_name, arrow::Type::type expected_file_vector_type, + int32_t vector_length, const std::vector& expected_ids, + const std::vector>>& expected_vectors) { + std::string file_path = + paimon::test::GetDataDir() + "/parquet/vector_compatibility/" + file_name; + std::shared_ptr file_type; + ReadFileType(file_path, &file_type); + std::shared_ptr file_vector_field = file_type->GetFieldByName("embedding"); + ASSERT_TRUE(file_vector_field); + ASSERT_EQ(file_vector_field->type()->id(), expected_file_vector_type); + std::shared_ptr file_id_field = file_type->GetFieldByName("id"); + ASSERT_TRUE(file_id_field); + + auto vector_type = arrow::fixed_size_list( + arrow::field("element", arrow::float32(), /*nullable=*/false), vector_length); + auto logical_schema = + arrow::schema({file_id_field, file_vector_field->WithType(vector_type)}); + std::unique_ptr vector_reader; + CreateVectorReader(file_path, logical_schema, /*predicate=*/nullptr, /*options=*/{}, + /*batch_size=*/10, &vector_reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(vector_reader.get())); + ASSERT_EQ(actual->num_chunks(), 1); + ASSERT_EQ(actual->type()->id(), arrow::Type::STRUCT); + auto struct_array = checked_pointer_cast(actual->chunk(0)); + std::shared_ptr id_field = struct_array->GetFieldByName("id"); + std::shared_ptr vector_field = struct_array->GetFieldByName("embedding"); + ASSERT_TRUE(id_field); + ASSERT_TRUE(vector_field); + ASSERT_EQ(id_field->type_id(), arrow::Type::INT32); + ASSERT_EQ(vector_field->type_id(), arrow::Type::FIXED_SIZE_LIST); + auto ids = checked_pointer_cast(id_field); + auto vector_array = checked_pointer_cast(vector_field); + ASSERT_EQ(ids->length(), static_cast(expected_ids.size())); + ASSERT_EQ(vector_array->length(), static_cast(expected_vectors.size())); + for (int64_t i = 0; i < ids->length(); ++i) { + ASSERT_FALSE(ids->IsNull(i)); + ASSERT_EQ(ids->Value(i), expected_ids[i]); + if (!expected_vectors[i]) { + ASSERT_TRUE(vector_array->IsNull(i)); + continue; + } + ASSERT_FALSE(vector_array->IsNull(i)); + ASSERT_EQ(vector_array->value_length(i), + static_cast(expected_vectors[i]->size())); + auto values = checked_pointer_cast(vector_array->value_slice(i)); + for (int64_t j = 0; j < values->length(); ++j) { + ASSERT_FALSE(values->IsNull(j)); + ASSERT_FLOAT_EQ(values->Value(j), expected_vectors[i].value()[j]); + } + } + } + + private: + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr fs_; + std::unique_ptr dir_; +}; + +TEST_F(ParquetVectorIoTest, WriteAndReadVector) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto struct_type = checked_pointer_cast(arrow::struct_( + {arrow::field("id", arrow::int32()), arrow::field("embedding", vector_type)})); + WriteAndCheck("vector.parquet", struct_type, struct_type, + R"([[1, [1.0, 2.0, 3.0]], [2, null], [3, [4.0, 5.0, 6.0]]])"); +} + +TEST_F(ParquetVectorIoTest, WriteAndReadAllNullVector) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto struct_type = checked_pointer_cast(arrow::struct_( + {arrow::field("id", arrow::int32()), arrow::field("embedding", vector_type)})); + WriteAndCheck("all-null-vector-list.parquet", struct_type, struct_type, + R"([[1, null], [2, null], [3, null]])"); +} + +TEST_F(ParquetVectorIoTest, WriteAndReadAllNullFixedSizeListWithArrowSchema) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto logical_type = checked_pointer_cast(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("embedding", vector_type), + })); + const std::string json = R"([[1, null], [2, null], [3, null]])"; + std::string file_path = dir_->Str() + "/all-null-vector.parquet"; + WriteWithArrowWriter(file_path, logical_type, json); + + std::shared_ptr file_type; + ReadFileType(file_path, &file_type); + std::shared_ptr file_vector_field = file_type->GetFieldByName("embedding"); + ASSERT_TRUE(file_vector_field); + ASSERT_EQ(file_vector_field->type()->id(), arrow::Type::FIXED_SIZE_LIST); + + std::unique_ptr reader; + CreateVectorReader(file_path, arrow::schema(logical_type->fields()), /*predicate=*/nullptr, + /*options=*/{}, /*batch_size=*/10, &reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON(logical_type, json); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + ASSERT_TRUE(std::make_shared(std::move(expected_result).ValueOrDie()) + ->Equals(actual)) + << actual->ToString(); +} + +TEST_F(ParquetVectorIoTest, ReadOrdinaryParquetListAsVector) { + auto physical_type = checked_pointer_cast( + arrow::struct_({arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::list(arrow::float32()))})); + auto logical_type = checked_pointer_cast(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)), + })); + WriteAndCheck("list.parquet", physical_type, logical_type, + R"([[1, [1.0, 2.0, 3.0]], [2, null], [3, [4.0, 5.0, 6.0]]])"); +} + +TEST_F(ParquetVectorIoTest, WriteAndReadNestedDoubleVector) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float64(), /*nullable=*/false), 2); + auto struct_type = checked_pointer_cast(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_type), + arrow::field("history", arrow::list(vector_type)), + arrow::field("by_name", arrow::map(arrow::utf8(), + vector_type))})), + })); + WriteAndCheck("nested-vector.parquet", struct_type, struct_type, + R"([[1, [[1.0, 2.0], [[3.0, 4.0], null], [["a", [5.0, 6.0]]]]], + [2, [null, null, [["b", null]]]]])"); +} + +// Vectors nested in a LIST keep their Arrow type when a third-party writer stores them as +// FixedSizeList, so the Parquet reader must accept a FixedSizeList read type as well. +TEST_F(ParquetVectorIoTest, ReadNestedFixedSizeListFile) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto logical_type = checked_pointer_cast(arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("history", arrow::list(vector_type)), + })); + const std::string json = R"([[1, [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], [2, []]])"; + std::string file_path = dir_->Str() + "/nested-fixed-size-list.parquet"; + WriteWithArrowWriter(file_path, logical_type, json); + + // Without this the file would expose the column as list> and the read would take + // the LIST to VECTOR conversion instead of the nested FixedSizeList path under test. + std::shared_ptr file_type; + ReadFileType(file_path, &file_type); + std::shared_ptr file_history_field = file_type->GetFieldByName("history"); + ASSERT_TRUE(file_history_field); + ASSERT_EQ(file_history_field->type()->id(), arrow::Type::LIST); + ASSERT_EQ(file_history_field->type()->field(0)->type()->id(), arrow::Type::FIXED_SIZE_LIST); + + std::unique_ptr reader; + CreateVectorReader(file_path, arrow::schema(logical_type->fields()), /*predicate=*/nullptr, + /*options=*/{}, /*batch_size=*/10, &reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON(logical_type, json); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + ASSERT_TRUE(std::make_shared(std::move(expected_result).ValueOrDie()) + ->Equals(actual)) + << actual->ToString(); +} + +TEST_F(ParquetVectorIoTest, ReadVectorWithPredicatePushdown) { + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + auto logical_type = checked_pointer_cast(arrow::struct_( + {arrow::field("id", arrow::int32()), arrow::field("embedding", vector_type)})); + // One row per row group, so the predicate on `id` prunes row groups while reading. + std::string file_path = dir_->Str() + "/vector-predicate.parquet"; + WriteWithFormatWriter(file_path, logical_type, + R"([[1, [1.0, 2.0, 3.0]], [2, null], [3, [4.0, 5.0, 6.0]], + [4, [7.0, 8.0, 9.0]]])", + /*max_row_group_length=*/1); + + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(2)); + std::unique_ptr reader; + CreateVectorReader(file_path, arrow::schema(logical_type->fields()), predicate, + /*options=*/{}, /*batch_size=*/10, &reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON( + logical_type, R"([[3, [4.0, 5.0, 6.0]], [4, [7.0, 8.0, 9.0]]])"); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + ASSERT_TRUE(std::make_shared(std::move(expected_result).ValueOrDie()) + ->Equals(actual)) + << actual->ToString(); +} + +TEST_F(ParquetVectorIoTest, ReadJavaFixture) { + ReadFixtureAndCheck( + "java_vector.parquet", arrow::Type::LIST, /*vector_length=*/2, + /*expected_ids=*/{0, 1, 2, 3, 4}, + /*expected_vectors=*/ + {{{0.0f, 0.0f}}, {{1.0f, 0.0f}}, {{2.0f, 0.0f}}, {{3.0f, 0.0f}}, {{4.0f, 0.0f}}}); +} + +TEST_F(ParquetVectorIoTest, ReadRustFixture) { + ReadFixtureAndCheck("rust_vector.parquet", arrow::Type::FIXED_SIZE_LIST, + /*vector_length=*/3, /*expected_ids=*/{1, 2, 3}, + /*expected_vectors=*/ + {{{1.0f, 2.0f, 3.0f}}, {{7.0f, 8.0f, 9.0f}}, {{4.0f, 5.0f, 6.0f}}}); +} + +TEST_F(ParquetVectorIoTest, ReadNullableJavaFixture) { + ReadFixtureAndCheck("java_vector_nullable.parquet", arrow::Type::LIST, /*vector_length=*/3, + /*expected_ids=*/{1, 2, 3}, + /*expected_vectors=*/ + {{{1.0f, 2.0f, 3.0f}}, std::nullopt, {{4.0f, 5.0f, 6.0f}}}); +} + +TEST_F(ParquetVectorIoTest, ReadNullableRustFixture) { + ReadFixtureAndCheck("rust_vector_nullable.parquet", arrow::Type::FIXED_SIZE_LIST, + /*vector_length=*/3, /*expected_ids=*/{1, 2, 3}, + /*expected_vectors=*/ + {{{1.0f, 2.0f, 3.0f}}, std::nullopt, {{4.0f, 5.0f, 6.0f}}}); +} + +// A table can hold files from several writers, and Paimon Java stores VECTOR as Parquet LIST +// while Paimon Rust stores it as FixedSizeList. Reading both with the table schema must produce +// batches of one Arrow type, otherwise they cannot be combined into a single result. +TEST_F(ParquetVectorIoTest, ReadMixedListAndFixedSizeListFixtures) { + // The Arrow type a Paimon schema builds for `id INT, embedding VECTOR`. The Rust + // fixture instead names the element field `element` and marks it non-nullable. + auto logical_schema = + arrow::schema({arrow::field("id", arrow::int32()), + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))}); + std::shared_ptr logical_type = arrow::struct_(logical_schema->fields()); + + // A reader owns the memory pool that its batches are allocated from, so it has to outlive + // the chunks collected from it. This mirrors a scan, which holds every split reader until + // the whole result has been consumed. + std::vector> readers; + arrow::ArrayVector chunks; + for (const char* file_name : {"java_vector_nullable.parquet", "rust_vector_nullable.parquet"}) { + std::string file_path = + paimon::test::GetDataDir() + "/parquet/vector_compatibility/" + file_name; + std::unique_ptr reader; + CreateVectorReader(file_path, logical_schema, /*predicate=*/nullptr, /*options=*/{}, + /*batch_size=*/10, &reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + readers.push_back(std::move(reader)); + ASSERT_TRUE(actual->type()->Equals(logical_type)) + << file_name << ": " << actual->type()->ToString(); + chunks.insert(chunks.end(), actual->chunks().begin(), actual->chunks().end()); + } + + arrow::Result> merged_result = + arrow::ChunkedArray::Make(chunks); + ASSERT_TRUE(merged_result.ok()) << merged_result.status().ToString(); + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON( + logical_type, R"([[1, [1.0, 2.0, 3.0]], [2, null], [3, [4.0, 5.0, 6.0]], + [1, [1.0, 2.0, 3.0]], [2, null], [3, [4.0, 5.0, 6.0]]])"); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + std::shared_ptr merged = std::move(merged_result).ValueOrDie(); + ASSERT_TRUE(std::make_shared(std::move(expected_result).ValueOrDie()) + ->Equals(merged)) + << merged->ToString(); +} + +} // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/predicate_pushdown_test.cpp b/src/paimon/format/parquet/predicate_pushdown_test.cpp index 9723794e..26175e1d 100644 --- a/src/paimon/format/parquet/predicate_pushdown_test.cpp +++ b/src/paimon/format/parquet/predicate_pushdown_test.cpp @@ -130,7 +130,8 @@ class PredicatePushdownTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto batch_reader, ParquetFileBatchReader::Create( std::move(in_stream), options, batch_size_, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); std::unique_ptr c_schema = std::make_unique(); auto arrow_status = arrow::ExportSchema(*read_schema, c_schema.get()); ASSERT_TRUE(arrow_status.ok()); diff --git a/src/paimon/format/parquet/variant_parquet_test.cpp b/src/paimon/format/parquet/variant_parquet_test.cpp index fd60e7b2..404b0521 100644 --- a/src/paimon/format/parquet/variant_parquet_test.cpp +++ b/src/paimon/format/parquet/variant_parquet_test.cpp @@ -386,7 +386,8 @@ class VariantParquetTest : public ::testing::Test { std::move(in_stream), options, /*batch_size=*/1024, /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); *file_reader = std::move(parquet_reader); ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> c_file_schema, (*file_reader)->GetFileSchema()); @@ -569,11 +570,12 @@ TEST_F(VariantParquetTest, WriteAndReadRoundTrip) { auto in_stream = std::make_unique(std::move(input_stream), length, arrow_pool_); std::map options = {}; - ASSERT_OK_AND_ASSIGN(auto batch_reader, ParquetFileBatchReader::Create( - std::move(in_stream), options, - /*batch_size=*/1024, - /*file_metadata=*/nullptr, - /*storage_read_bytes=*/nullptr, arrow_pool_)); + ASSERT_OK_AND_ASSIGN(auto batch_reader, + ParquetFileBatchReader::Create(std::move(in_stream), options, + /*batch_size=*/1024, + /*file_metadata=*/nullptr, + /*storage_read_bytes=*/nullptr, arrow_pool_, + /*hints=*/std::nullopt)); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*paimon_schema_, c_schema.get()).ok()); ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, diff --git a/src/paimon/rest/rest_http_client.cpp b/src/paimon/rest/rest_http_client.cpp index b17b0a28..52e6417e 100644 --- a/src/paimon/rest/rest_http_client.cpp +++ b/src/paimon/rest/rest_http_client.cpp @@ -123,7 +123,9 @@ bool IsRetriableTransportError(CURLcode code) { case CURLE_RECV_ERROR: case CURLE_PARTIAL_FILE: case CURLE_HTTP2: +#if CURL_AT_LEAST_VERSION(7, 49, 0) case CURLE_HTTP2_STREAM: +#endif return true; default: return false; diff --git a/test/inte/data_evolution_table_test.cpp b/test/inte/data_evolution_table_test.cpp index 6232c1e7..13fede80 100644 --- a/test/inte/data_evolution_table_test.cpp +++ b/test/inte/data_evolution_table_test.cpp @@ -400,10 +400,11 @@ class DataEvolutionTableTest : public ::testing::Test, const std::shared_ptr& expected_array, const std::shared_ptr& predicate = nullptr, const std::vector& row_ranges = {}, - bool check_scan_plan_when_empty_result = true) const { + 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(predicate); + 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); @@ -1880,7 +1881,7 @@ TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) { expected_array)); } { - // first 4 records read with data evolution, ignore index + // The old file's f2 index does not contain 102, but the newer file owns f2. auto predicate = PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT, Literal(102)); auto expected_array = std::dynamic_pointer_cast( @@ -1892,51 +1893,45 @@ TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) { ])") .ValueOrDie()); ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(), - expected_array, predicate)); + expected_array, predicate, + /*row_ranges=*/{}, + /*check_scan_plan_when_empty_result=*/true, + /*apply_predicate_to_scan=*/false)); } { - // f2 has bitmap index, but data evolution scan and read ignore index + // 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)); - auto expected_array = std::dynamic_pointer_cast( - 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] - ])") - .ValueOrDie()); ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(), - expected_array, predicate)); + /*expected_array=*/nullptr, predicate, + /*row_ranges=*/{}, + /*check_scan_plan_when_empty_result=*/false, + /*apply_predicate_to_scan=*/false)); } { - // f2 has bitmap index, data evolution scan will ignore index => not empty plan - // data evolution split read will also ignore index => not empty read batch + // 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)); - auto expected_array = std::dynamic_pointer_cast( - 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, + /*expected_array=*/nullptr, predicate, /*row_ranges=*/{}, - /*check_scan_plan_when_empty_result=*/true)); + /*check_scan_plan_when_empty_result=*/false, + /*apply_predicate_to_scan=*/false)); } { - // f2 has bitmap index, data evolution split read will ignore index + // 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::ipc::internal::json::ArrayFromJSON(arrow_data_type, R"([ - [null, null, 202, 6.1], - [null, null, 204, 7.1] + [null, null, 202, 6.1] ])") .ValueOrDie()); ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(), - expected_array, predicate)); + expected_array, predicate, + /*row_ranges=*/{}, + /*check_scan_plan_when_empty_result=*/true, + /*apply_predicate_to_scan=*/false)); } { auto predicate = @@ -1953,7 +1948,7 @@ TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) { { // test row id with predicate std::vector row_ranges = {Range(0l, 2l)}; - // row id = {0, 1, 2}, while data evolution split read will ignore index + // A merged group keeps all selected row ids to preserve 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, @@ -1967,26 +1962,179 @@ TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) { .ValueOrDie()); ASSERT_OK(ScanAndRead(table_path, arrow::schema(arrow_data_type->fields())->field_names(), expected_array, predicate, - /*row_ranges=*/row_ranges)); + /*row_ranges=*/row_ranges, + /*check_scan_plan_when_empty_result=*/true, + /*apply_predicate_to_scan=*/false)); } { // test row id with predicate std::vector row_ranges = {Range(4l, 5l)}; - // row id = {4, 5}, data evolution split read will ignore bitmap index + // 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::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, - /*row_ranges=*/row_ranges)); + /*row_ranges=*/row_ranges, + /*check_scan_plan_when_empty_result=*/true, + /*apply_predicate_to_scan=*/false)); + } +} + +TEST_P(DataEvolutionTableTest, TestDataEvolutionPredicatePushDownBoundaries) { + 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::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(99l)); + auto read_type = + arrow::struct_({arrow::field("f2", arrow::int32()), SpecialFields::RowId().field_}); + auto expected_array = std::dynamic_pointer_cast( + 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(0l)); + 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().field_}); + auto expected_array = std::dynamic_pointer_cast( + 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().field_}); + auto expected_array = std::dynamic_pointer_cast( + 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") { + return; + } + + CreateDataEvolutionTable( + /*deletion_vectors_enabled=*/false, {{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"}}); + std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + + auto input = std::dynamic_pointer_cast( + 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_messages, WriteArray(table_path, {"f0", "f1", "f2"}, input)); + 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::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) { diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 6d7e5562..1e0952e0 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -47,6 +47,7 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/read_ahead_cache.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" @@ -91,7 +92,7 @@ struct TestParam { bool enable_prefetch; std::string enable_adaptive_prefetch_strategy; std::string file_format; - PrefetchCacheMode cache_mode; + bool read_ahead_cache_enabled; }; // read_inte_test.cpp test mainly for raw file split read (pk+dv & append only) @@ -354,19 +355,16 @@ Result CountDataFiles(const std::vector>& splits std::vector PrepareTestParam() { std::vector values = { - TestParam{false, "false", "parquet", PrefetchCacheMode::ALWAYS}, - TestParam{true, "true", "parquet", PrefetchCacheMode::ALWAYS}, - TestParam{true, "false", "parquet", PrefetchCacheMode::ALWAYS}, - TestParam{true, "false", "parquet", PrefetchCacheMode::NEVER}, - TestParam{true, "false", "parquet", PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE}}; + TestParam{false, "false", "parquet", /*read_ahead_cache_enabled=*/true}, + TestParam{true, "true", "parquet", /*read_ahead_cache_enabled=*/true}, + TestParam{true, "false", "parquet", /*read_ahead_cache_enabled=*/true}, + TestParam{true, "false", "parquet", /*read_ahead_cache_enabled=*/false}}; #ifdef PAIMON_ENABLE_ORC - values.push_back(TestParam{false, "false", "orc", PrefetchCacheMode::ALWAYS}); - values.push_back(TestParam{true, "true", "orc", PrefetchCacheMode::ALWAYS}); - values.push_back(TestParam{true, "false", "orc", PrefetchCacheMode::ALWAYS}); - values.push_back(TestParam{true, "false", "orc", PrefetchCacheMode::NEVER}); - values.push_back( - TestParam{true, "false", "orc", PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE}); + values.push_back(TestParam{false, "false", "orc", /*read_ahead_cache_enabled=*/true}); + values.push_back(TestParam{true, "true", "orc", /*read_ahead_cache_enabled=*/true}); + values.push_back(TestParam{true, "false", "orc", /*read_ahead_cache_enabled=*/true}); + values.push_back(TestParam{true, "false", "orc", /*read_ahead_cache_enabled=*/false}); #endif return values; } @@ -390,7 +388,7 @@ TEST_P(ReadInteTest, TestAppendSimple) { context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", "false") .AddOption("orc.read.enable-metrics", "true"); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); if (specific_table_schema) { context_builder.SetTableSchema(specific_table_schema.value()); @@ -496,7 +494,7 @@ TEST_P(ReadInteTest, TestReadWithLimits) { context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption(Options::READ_BATCH_SIZE, "1"); context_builder.EnablePrefetch(param.enable_prefetch) - .SetPrefetchCacheMode(param.cache_mode) + .SetReadAheadCacheEnabled(param.read_ahead_cache_enabled) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) .AddOption("orc.read.enable-metrics", "true") @@ -546,6 +544,76 @@ TEST_P(ReadInteTest, TestReadWithLimits) { } } +TEST_P(ReadInteTest, TestReadAheadCacheMetrics) { + auto param = GetParam(); + std::string path = + paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; + ReadContextBuilder context_builder(path); + context_builder.AddOption(Options::FILE_FORMAT, param.file_format); + context_builder.EnablePrefetch(param.enable_prefetch) + .AddOption("test.enable-adaptive-prefetch-strategy", "false") + .SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); + + ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + + std::vector file_list; + if (param.file_format == "orc") { + file_list = {"data-db2b44c0-0d73-449d-82a0-4075bd2cb6e3-0.orc", + "data-b913a160-a4d1-4084-af2a-18333c35668e-0.orc"}; + } else if (param.file_format == "parquet") { + file_list = {"data-b446f78a-2cfb-4b3b-add8-31295d24a277-0.parquet", + "data-fd72a479-53ae-42f7-aec0-e982ee555928-0.parquet"}; + } + + DataSplitsSimple input_data_splits = {{paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=20/" + "bucket-0", + BinaryRowGenerator::GenerateRow({20}, pool_.get()), + file_list}}; + + auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/3); + ASSERT_EQ(data_splits.size(), 1); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_TRUE(result_array); + ASSERT_EQ(result_array->length(), 2); + + // Verify the read-ahead cache metrics are surfaced through the reader chain. The prefetch + // reader merges the cache counters into its reader metrics only when a cache is created, + // so the counters must be present and effective exactly in that case. + auto read_metrics = batch_reader->GetReaderMetrics(); + ASSERT_TRUE(read_metrics); + if (param.enable_prefetch && param.read_ahead_cache_enabled) { + ASSERT_OK_AND_ASSIGN(uint64_t read_count, + read_metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); + ASSERT_GT(read_count, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t read_bytes, + read_metrics->GetCounter(ReadAheadCacheMetrics::READ_BYTES)); + ASSERT_GT(read_bytes, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t hits, + read_metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_GT(hits, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t hit_bytes, + read_metrics->GetCounter(ReadAheadCacheMetrics::READ_HIT_BYTES)); + ASSERT_GT(hit_bytes, 0u); + ASSERT_OK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_OK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISS_BYTES)); + // Serving hits requires prefetch IOs issued to the underlying stream. + ASSERT_OK_AND_ASSIGN(uint64_t io_count, + read_metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + ASSERT_GT(io_count, 0u); + ASSERT_OK_AND_ASSIGN(uint64_t io_bytes, + read_metrics->GetCounter(ReadAheadCacheMetrics::IO_BYTES)); + ASSERT_GT(io_bytes, 0u); + } else { + ASSERT_NOK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_COUNT)); + ASSERT_NOK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_HITS)); + ASSERT_NOK(read_metrics->GetCounter(ReadAheadCacheMetrics::READ_MISSES)); + ASSERT_NOK(read_metrics->GetCounter(ReadAheadCacheMetrics::IO_COUNT)); + } +} + TEST_P(ReadInteTest, TestReadOnlyPartitionField) { auto param = GetParam(); std::string path = paimon::test::GetDataDir() + "/" + param.file_format + @@ -557,7 +625,7 @@ TEST_P(ReadInteTest, TestReadOnlyPartitionField) { ReadContextBuilder context_builder(path); context_builder.AddOption(Options::FILE_FORMAT, param.file_format); context_builder.SetReadFieldNames({"dt"}); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("test.enable-adaptive-prefetch-strategy", @@ -1879,7 +1947,7 @@ TEST_P(ReadInteTest, TestAppendReadWithMultipleBuckets) { paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; ReadContextBuilder context_builder(path); context_builder.SetReadFieldNames({"f3", "f0", "f1"}); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2") .AddOption("test.enable-adaptive-prefetch-strategy", @@ -1959,7 +2027,7 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicate) { ReadContextBuilder context_builder(path); context_builder.SetReadFieldNames({"f3", "f0", "f1"}); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .SetPredicate(predicate) .EnablePredicateFilter(true) @@ -2061,7 +2129,7 @@ TEST_P(ReadInteTest, TestAppendReadWithComplexTypePredicate) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/append_complex_data.db/append_complex_data"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"f6", "f2", "f4", "f3", "f5"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -2134,7 +2202,7 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateOnlyPushdown) { paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"f3", "f0", "f1"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2") @@ -2210,7 +2278,7 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateAllFiltered) { paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"f3", "f0", "f1"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2") @@ -2297,7 +2365,7 @@ TEST_P(ReadInteTest, TestAppendReadIOException) { ReadContextBuilder context_builder(paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09/"); context_builder.SetReadFieldNames({"f3", "f0", "f1"}); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2") .EnablePrefetch(param.enable_prefetch) @@ -2342,7 +2410,7 @@ TEST_P(ReadInteTest, TestPkTableWithDeletionVectorSimple) { ReadContextBuilder context_builder(path); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); @@ -2388,7 +2456,7 @@ TEST_P(ReadInteTest, TestPkTableWithDeletionVector) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/pk_09.db/pk_09"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2") .EnablePrefetch(param.enable_prefetch) @@ -2454,7 +2522,7 @@ TEST_P(ReadInteTest, TestPkTableWithSnapshot6) { FieldType::DOUBLE, Literal(15.0)); std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/pk_09.db/pk_09"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); context_builder.SetPredicate(predicate); @@ -2539,7 +2607,7 @@ TEST_P(ReadInteTest, TestPkTableWithSnapshot8) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/pk_09.db/pk_09"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"f0", "f3", "f1"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -2617,7 +2685,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolution) { DataField(8, arrow::field("e", arrow::int32()))}; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); context_builder.EnablePrefetch(param.enable_prefetch) @@ -2713,13 +2781,13 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithPredicateFilter) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/append_table_with_alter_table.db/append_table_with_alter_table/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"a", "k", "key1", "d", "key0", "c"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); context_builder.SetPredicate(predicate); context_builder.EnablePredicateFilter(true); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); @@ -2792,7 +2860,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithPredicateOnlyPushDown) "/append_table_with_alter_table.db/" "append_table_with_alter_table/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"a", "k", "key1", "d", "key0", "c"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -2865,7 +2933,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot5WithSchemaEvolution) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/pk_table_with_alter_table.db/pk_table_with_alter_table/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"key1", "k", "key_2", "c", "d", "a", "key0", "e"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -2950,7 +3018,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolution) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/pk_table_with_alter_table.db/pk_table_with_alter_table/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"key1", "k", "key_2", "c", "d", "a", "key0", "e"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -3038,7 +3106,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush context_builder.SetReadFieldNames({{"key1", "k", "key_2", "c", "d", "a", "key0", "e"}}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetPredicate(predicate); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", @@ -3117,7 +3185,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateFilter) ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({equal, less_than})); ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"key1", "k", "key_2", "c", "d", "a", "key0", "e"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -3209,7 +3277,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithBuildInFieldId) { } ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"key0", "key1", "k", "c", "d", "a", "e"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -3272,7 +3340,7 @@ TEST_P(ReadInteTest, TestAppendReadNestedType) { std::string path = paimon::test::GetDataDir() + "/" + param.file_format + "/append_complex_build_in_fieldid.db/append_complex_build_in_fieldid/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); context_builder.EnablePrefetch(param.enable_prefetch) @@ -3327,7 +3395,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithCast) { "/append_table_alter_table_with_cast.db/" "append_table_alter_table_with_cast/"; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetReadFieldNames({"f4", "key0", "key1", "f3", "f1", "f2", "f0", "f6"}); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); @@ -3410,7 +3478,7 @@ TEST_P(ReadInteTest, TestAppendReadWithSchemaEvolutionWithCastWithPredicatePushD "append_table_alter_table_with_cast/"; ReadContextBuilder context_builder(path); context_builder.SetReadFieldNames({"f4", "key0", "key1", "f3", "f1", "f2", "f0", "f6"}); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("read.batch-size", "2"); context_builder.SetPredicate(predicate); @@ -3488,7 +3556,7 @@ TEST_P(ReadInteTest, TestReadWithPKFallBackBranch) { }; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption(Options::FILE_FORMAT, param.file_format) .AddOption("test.enable-adaptive-prefetch-strategy", @@ -3545,7 +3613,7 @@ TEST_P(ReadInteTest, TestReadWithAppendFallBackBranch) { ReadContextBuilder context_builder(path); context_builder.EnablePrefetch(param.enable_prefetch); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); ASSERT_OK_AND_ASSIGN(auto read_context, 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(data_splits)); @@ -3587,7 +3655,7 @@ TEST_P(ReadInteTest, TestFallBackBranchStreamRead) { DataField(1, arrow::field("name", arrow::utf8())), DataField(2, arrow::field("amount", arrow::int32()))}; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); @@ -3632,7 +3700,7 @@ TEST_P(ReadInteTest, TestReadWithPKRtBranch) { }; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) @@ -3689,7 +3757,7 @@ TEST_P(ReadInteTest, TestReadWithAppendPtBranch) { }; ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) @@ -3820,7 +3888,7 @@ TEST_P(ReadInteTest, TestSpecificFs) { auto countable_fs = std::make_shared(std::make_shared(), &io_count); ReadContextBuilder context_builder(path); - context_builder.SetPrefetchCacheMode(param.cache_mode); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.AddOption(Options::FILE_FORMAT, param.file_format); context_builder.EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", "false") diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 67cc21d0..78e8bade 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -23,10 +23,14 @@ #include #include #include +#include #include +#include #include #include #include +#include +#include #include #include #include @@ -40,7 +44,11 @@ #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/commit_context.h" +#include "paimon/common/factories/io_hook.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/realtime/realtime_context_impl.h" @@ -50,11 +58,15 @@ #include "paimon/defs.h" #include "paimon/file_store_commit.h" #include "paimon/file_store_write.h" +#include "paimon/fs/file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/orphan_files_cleaner.h" #include "paimon/predicate/function.h" #include "paimon/predicate/predicate.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" +#include "paimon/reader/count_reader.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" #include "paimon/record_batch.h" @@ -62,10 +74,587 @@ #include "paimon/table/source/table_read.h" #include "paimon/table/source/table_scan.h" #include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/test_helper.h" #include "paimon/testing/utils/testharness.h" #include "paimon/write_context.h" namespace paimon::test { +namespace { + +bool HasSuffix(const std::string& value, const std::string& suffix) { + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +Result> ListPhysicalArtifacts(const std::shared_ptr& file_system, + const std::string& root) { + std::set artifacts; + std::vector directories = {root}; + while (!directories.empty()) { + std::string directory = std::move(directories.back()); + directories.pop_back(); + std::vector statuses; + PAIMON_RETURN_NOT_OK(file_system->ListDir(directory, &statuses)); + for (const BasicFileStatus& status : statuses) { + if (status.IsDir()) { + directories.push_back(status.GetPath()); + } else if (HasSuffix(status.GetPath(), ".orc") || + HasSuffix(status.GetPath(), ".index") || + HasSuffix(status.GetPath(), ".channel")) { + artifacts.insert(status.GetPath()); + } + } + } + return artifacts; +} + +class FailAllocationMemoryPool final : public MemoryPool { + public: + explicit FailAllocationMemoryPool(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + void FailAfterAllocations(int64_t successful_allocations) { + allocations_before_failure_.store(successful_allocations, std::memory_order_release); + } + + void* Malloc(uint64_t size, uint64_t alignment = 0) override { + if (ShouldFail()) { + throw std::bad_alloc(); + } + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* p, size_t old_size, size_t new_size, uint64_t alignment = 0) override { + if (ShouldFail()) { + throw std::bad_alloc(); + } + return delegate_->Realloc(p, old_size, new_size, alignment); + } + + void Free(void* p, uint64_t size) override { + delegate_->Free(p, size); + } + + void Free(void* p, uint64_t size, uint64_t alignment) override { + delegate_->Free(p, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + private: + bool ShouldFail() { + int64_t remaining = allocations_before_failure_.load(std::memory_order_acquire); + while (remaining >= 0) { + if (allocations_before_failure_.compare_exchange_weak(remaining, remaining - 1, + std::memory_order_acq_rel)) { + return remaining == 0; + } + } + return false; + } + + std::shared_ptr delegate_; + std::atomic allocations_before_failure_{-1}; +}; + +class TrackingRealtimeReadView final : public RealtimeReadView { + public: + explicit TrackingRealtimeReadView(std::shared_ptr delegate) + : delegate_(std::move(delegate)) {} + + std::optional GetOffsetRange() const override { + return delegate_->GetOffsetRange(); + } + + const std::shared_ptr& Delegate() const { + return delegate_; + } + + private: + std::shared_ptr delegate_; +}; + +class ReadViewCheckingBatchReader final : public BatchReader { + public: + ReadViewCheckingBatchReader(std::unique_ptr delegate, + std::weak_ptr read_view) + : delegate_(std::move(delegate)), read_view_(std::move(read_view)) {} + + Result NextBatch() override { + if (read_view_.expired()) { + return Status::Invalid("real-time read view was released before reader completion"); + } + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::weak_ptr read_view_; +}; + +class DelegatingRealtimeStore : public RealtimeStore { + public: + explicit DelegatingRealtimeStore(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + protected: + std::shared_ptr delegate_; +}; + +class DecoratingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + using Decorator = + std::function(const std::shared_ptr&)>; + + explicit DecoratingRealtimeStoreFactory(Decorator decorator) + : decorator_(std::move(decorator)) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return decorator_(delegate); + } + + private: + ArrowRealtimeStoreFactory delegate_; + Decorator decorator_; +}; + +template +std::shared_ptr MakeDecoratingFactory(Args... args) { + return std::make_shared( + [=](const std::shared_ptr& delegate) -> std::shared_ptr { + return std::make_shared(delegate, args...); + }); +} + +class QueryTrackingRealtimeStore final : public DelegatingRealtimeStore { + public: + QueryTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : DelegatingRealtimeStore(delegate), + saw_query_predicate_(saw_query_predicate), + query_view_(query_view) {} + + Result> AcquireReadView() override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate_view, + delegate_->AcquireReadView()); + return std::shared_ptr( + std::make_shared(delegate_view)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + if (context.predicate) { + saw_query_predicate_->store(true, std::memory_order_release); + } + *query_view_ = view; + std::shared_ptr tracking_view = + std::dynamic_pointer_cast(view); + if (!tracking_view) { + return Status::Invalid("query tracking store received an unexpected read view"); + } + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + delegate_->CreateQueryReaders(tracking_view->Delegate(), offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), view); + } + return readers; + } + + private: + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +class CloseTrackingBatchReader final : public BatchReader { + public: + CloseTrackingBatchReader(std::unique_ptr delegate, + const std::shared_ptr>& close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + close_count_->fetch_add(1, std::memory_order_release); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr> close_count_; +}; + +struct CloseTrackingReaderState { + std::shared_ptr> query_close_count = + std::make_shared>(0); + std::shared_ptr> commit_close_count = + std::make_shared>(0); + int32_t query_null_index = -1; + int32_t commit_null_index = -1; +}; + +class CloseTrackingRealtimeStore final : public DelegatingRealtimeStore { + public: + CloseTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr& state) + : DelegatingRealtimeStore(delegate), state_(state) {} + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), + state_->commit_close_count); + } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->commit_null_index, &readers)); + return readers; + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateQueryReaders(view, offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), + state_->query_close_count); + } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->query_null_index, &readers)); + return readers; + } + + private: + static Status InsertNullReader(int32_t index, + std::vector>* readers) { + if (index < 0) { + return Status::OK(); + } + if (index > static_cast(readers->size())) { + return Status::Invalid("null reader index exceeds reader count"); + } + readers->insert(readers->begin() + index, nullptr); + return Status::OK(); + } + + std::shared_ptr state_; +}; + +class SplitBatchReader final : public BatchReader { + public: + explicit SplitBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + while (!current_batch_ || next_row_ == current_batch_->length()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("split batch reader received a non-struct batch"); + } + current_batch_ = std::dynamic_pointer_cast(array); + next_row_ = 0; + } + std::shared_ptr slice = current_batch_->Slice(next_row_, /*length=*/1); + ++next_row_; + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*slice, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + current_batch_.reset(); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr current_batch_; + int64_t next_row_ = 0; +}; + +class SplitCommitReaderRealtimeStore final : public DelegatingRealtimeStore { + public: + explicit SplitCommitReaderRealtimeStore(const std::shared_ptr& delegate) + : DelegatingRealtimeStore(delegate) {} + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader)); + } + return readers; + } +}; + +class FailAfterPhysicalFileBatchReader final : public BatchReader { + public: + FailAfterPhysicalFileBatchReader(std::unique_ptr delegate, + const std::shared_ptr& file_system, + std::string root, size_t baseline_artifact_count, + const std::shared_ptr>& saw_artifacts) + : delegate_(std::move(delegate)), + file_system_(file_system), + root_(std::move(root)), + baseline_artifact_count_(baseline_artifact_count), + saw_artifacts_(saw_artifacts) {} + + Result NextBatch() override { + if (returned_batch_count_ < 4) { + ++returned_batch_count_; + return delegate_->NextBatch(); + } + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (std::chrono::steady_clock::now() < deadline) { + PAIMON_ASSIGN_OR_RAISE(std::set artifacts, + ListPhysicalArtifacts(file_system_, root_)); + bool has_data = false; + for (const std::string& artifact : artifacts) { + has_data = has_data || HasSuffix(artifact, ".orc"); + } + if (artifacts.size() > baseline_artifact_count_ && has_data) { + saw_artifacts_->store(true, std::memory_order_release); + return Status::IOError( + "injected commit reader failure after physical file creation"); + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return Status::IOError("timed out waiting for partial physical files"); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr file_system_; + std::string root_; + size_t baseline_artifact_count_; + std::shared_ptr> saw_artifacts_; + int32_t returned_batch_count_ = 0; +}; + +class FailAfterPhysicalFileRealtimeStore final : public DelegatingRealtimeStore { + public: + FailAfterPhysicalFileRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr& file_system, + const std::string& root, size_t baseline_artifact_count, + const std::shared_ptr>& saw_artifacts) + : DelegatingRealtimeStore(delegate), + file_system_(file_system), + root_(root), + baseline_artifact_count_(baseline_artifact_count), + saw_artifacts_(saw_artifacts) {} + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + if (readers.empty()) { + return Status::Invalid("commit reader failure test requires a reader"); + } + readers[0] = std::make_unique( + std::make_unique(std::move(readers[0])), file_system_, root_, + baseline_artifact_count_, saw_artifacts_); + return readers; + } + + private: + std::shared_ptr file_system_; + std::string root_; + size_t baseline_artifact_count_; + std::shared_ptr> saw_artifacts_; +}; + +enum class CommitReaderMalformation { DROP_LAST, DUPLICATE_OFFSET, OUT_OF_RANGE_OFFSET }; + +class CorruptingBatchReader final : public BatchReader { + public: + CorruptingBatchReader(std::unique_ptr delegate, + CommitReaderMalformation malformation) + : delegate_(std::move(delegate)), malformation_(malformation) {} + + Result NextBatch() override { + switch (malformation_) { + case CommitReaderMalformation::DROP_LAST: + return DropLast(); + case CommitReaderMalformation::DUPLICATE_OFFSET: + return SubstituteOffset(/*offset=*/0); + case CommitReaderMalformation::OUT_OF_RANGE_OFFSET: + return SubstituteOffset(/*offset=*/-1); + } + return Status::Invalid("unknown commit reader malformation"); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + buffered_.reset(); + delegate_->Close(); + } + + private: + Result DropLast() { + if (!buffered_.has_value()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch first, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(first)) { + return MakeEofBatch(); + } + buffered_ = std::move(first); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch next, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(next)) { + buffered_.reset(); + return MakeEofBatch(); + } + ReadBatch result = std::move(buffered_.value()); + buffered_ = std::move(next); + return result; + } + + Result SubstituteOffset(int64_t offset) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT || array->length() == 0) { + return Status::Invalid("offset substitution requires a non-empty struct batch"); + } + std::shared_ptr struct_array = + std::dynamic_pointer_cast(array); + std::shared_ptr offsets = + std::dynamic_pointer_cast(struct_array->field(2)); + if (!offsets) { + return Status::Invalid("offset substitution requires an int64 REALTIME_OFFSET"); + } + arrow::Int64Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(offsets->length())); + for (int64_t row = 0; row < offsets->length(); ++row) { + builder.UnsafeAppend(offset); + } + std::shared_ptr substituted_offsets; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&substituted_offsets)); + std::shared_ptr substituted_data = struct_array->data()->Copy(); + substituted_data->child_data[2] = substituted_offsets->data(); + std::shared_ptr substituted = arrow::MakeArray(std::move(substituted_data)); + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*substituted, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::unique_ptr delegate_; + CommitReaderMalformation malformation_; + std::optional buffered_; +}; + +class MalformedCoverageRealtimeStore final : public DelegatingRealtimeStore { + public: + MalformedCoverageRealtimeStore(const std::shared_ptr& delegate, + CommitReaderMalformation malformation) + : DelegatingRealtimeStore(delegate), malformation_(malformation) {} + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), malformation_); + } + return readers; + } + + private: + CommitReaderMalformation malformation_; +}; + +} // namespace + +namespace { + +constexpr char kDropPartitionCommitUser[] = "drop_partition_commit_user"; +constexpr char kRollbackCommitUser[] = "rollback_commit_user"; +constexpr char kTruncateCommitUser[] = "truncate_commit_user"; + +} // namespace class UnsupportedFunction : public Function { public: @@ -184,9 +773,10 @@ class RealtimeWriteInteTest : public ::testing::Test { arrow::field("pt", arrow::utf8())}; schema_ = arrow::schema(fields_); options_ = { - {Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, - {Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "1"}, - {Options::BUCKET_KEY, "id"}, {Options::TARGET_FILE_SIZE, "1048576"}, + {Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, + {Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, {Options::TARGET_FILE_SIZE, "1048576"}, + {Options::REALTIME_ENABLED, "true"}, }; } @@ -205,6 +795,20 @@ class RealtimeWriteInteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } + void CreatePkTable(const std::vector& partition_keys = {}, + const std::vector& primary_keys = {"id"}) const { + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema_, c_schema.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, + Catalog::Create(dir_->Str(), options_)); + ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false)); + std::vector table_primary_keys = partition_keys; + table_primary_keys.insert(table_primary_keys.end(), primary_keys.begin(), + primary_keys.end()); + ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), c_schema.get(), partition_keys, + table_primary_keys, options_, /*ignore_if_exists=*/false)); + } + Result> CreateRealtimeWriter( const std::shared_ptr& realtime_context) const { WriteContextBuilder builder(table_path_, commit_user_); @@ -226,6 +830,12 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> MakeBatch(const std::vector& rows, bool partitioned, int32_t bucket) const { + return MakeBatch(rows, partitioned, bucket, /*row_kinds=*/{}); + } + + Result> MakeBatch( + const std::vector& rows, bool partitioned, int32_t bucket, + const std::vector& row_kinds) const { if (rows.empty()) { return Status::Invalid("cannot create an empty test batch"); } @@ -233,7 +843,7 @@ class RealtimeWriteInteTest : public ::testing::Test { std::string json = "["; for (size_t i = 0; i < rows.size(); ++i) { const auto& [id, payload, pt] = rows[i]; - if (pt != partition) { + if (partitioned && pt != partition) { return Status::Invalid("one test batch must contain only one partition"); } if (i > 0) { @@ -249,12 +859,50 @@ class RealtimeWriteInteTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); RecordBatchBuilder builder(&c_array); + builder.SetRowKinds(row_kinds); if (partitioned) { builder.SetPartition({{"pt", partition}}); } return builder.SetBucket(bucket).Finish(); } + Result> MakeDatePartitionBatch( + int64_t first_id, int64_t count, int32_t date, const std::string& partition) const { + if (count <= 0) { + return Status::Invalid("cannot create an empty test batch"); + } + std::string json = "["; + for (int64_t i = 0; i < count; ++i) { + if (i > 0) { + json += ","; + } + int64_t id = first_id + i; + json += "[" + std::to_string(id) + ",\"value-" + std::to_string(id) + "\"," + + std::to_string(date) + "]"; + } + json += "]"; + + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + return RecordBatchBuilder(&c_array) + .SetPartition({{"pt", partition}}) + .SetBucket(/*bucket=*/0) + .Finish(); + } + + Result> MakeUnpartitionedBatchFromJson( + const std::string& json) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + return RecordBatchBuilder(&c_array).SetBucket(/*bucket=*/0).Finish(); + } + static std::vector MakeRows(int64_t first_id, int64_t count, const std::string& partition) { std::vector rows; @@ -277,26 +925,145 @@ class RealtimeWriteInteTest : public ::testing::Test { /*watermark=*/std::nullopt); } - Result> CreatePlan( - const std::shared_ptr& realtime_context, - const std::shared_ptr& predicate) const { - ScanContextBuilder scan_builder(table_path_); - if (realtime_context) { - scan_builder.WithRealtimeContext(realtime_context); + Result DropPartition(const std::map& partition, + int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, kDropPartitionCommitUser); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + PAIMON_RETURN_NOT_OK(commit->DropPartition({partition}, commit_identifier)); + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + if (!latest_snapshot) { + return Status::Invalid("drop partition did not produce a snapshot"); } - scan_builder.SetPredicate(predicate).WithMemoryPool(pool_); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, - scan_builder.SetOptions(options_).Finish()); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan, - TableScan::Create(std::move(scan_context))); - return scan->CreatePlan(); + return latest_snapshot->Id(); } - Result ReadPlan(const std::shared_ptr& plan, - const std::shared_ptr& realtime_context, - const std::vector& read_fields, - const std::shared_ptr& predicate, - bool enable_predicate_filter) const { + Result TruncateTable(int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, kTruncateCommitUser); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + PAIMON_RETURN_NOT_OK(commit->TruncateTable(commit_identifier)); + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + if (!latest_snapshot) { + return Status::Invalid("truncate did not produce a snapshot"); + } + return latest_snapshot->Id(); + } + + Result RollbackToAsLatest(int64_t target_snapshot_id) const { + CommitContextBuilder builder(table_path_, kRollbackCommitUser); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + PAIMON_ASSIGN_OR_RAISE(bool rolled_back, commit->RollbackToAsLatest(target_snapshot_id)); + if (!rolled_back) { + return Status::Invalid("failed to commit rollback snapshot"); + } + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + if (!latest_snapshot) { + return Status::Invalid("rollback did not produce a snapshot"); + } + return latest_snapshot->Id(); + } + + Result CompactAndCommit(const std::map& partition, + int32_t bucket, int64_t commit_identifier) const { + WriteContextBuilder write_builder(table_path_, commit_user_); + write_builder.SetOptions(options_).WithStreamingMode(true); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr write_context, write_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr compaction_writer, + FileStoreWrite::Create(std::move(write_context))); + PAIMON_RETURN_NOT_OK(compaction_writer->Compact(partition, bucket, + /*full_compaction=*/true)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> compaction_messages, + compaction_writer->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + if (compaction_messages.empty()) { + return Status::Invalid("compaction did not produce a commit message"); + } + + CommitContextBuilder commit_builder(table_path_, commit_user_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr compaction_commit, + FileStoreCommit::Create(std::move(commit_context))); + PAIMON_RETURN_NOT_OK(compaction_commit->Commit(compaction_messages, commit_identifier)); + PAIMON_RETURN_NOT_OK(compaction_writer->Close()); + + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(core_options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional compact_snapshot, + snapshot_manager.LatestSnapshot()); + if (!compact_snapshot) { + return Status::Invalid("compaction did not produce a snapshot"); + } + return compact_snapshot.value(); + } + + Result ExpireSnapshots() const { + CommitContextBuilder builder(table_path_, commit_user_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + return commit->Expire(); + } + + Result> CreatePlan( + const std::shared_ptr& realtime_context, + const std::shared_ptr& predicate) const { + ScanContextBuilder scan_builder(table_path_); + if (realtime_context) { + scan_builder.WithRealtimeContext(realtime_context); + } + scan_builder.SetPredicate(predicate).WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, + scan_builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan, + TableScan::Create(std::move(scan_context))); + return scan->CreatePlan(); + } + + Result> CreateQueryReader( + const std::shared_ptr& plan, + const std::shared_ptr& realtime_context) const { + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + return table_read->CreateReader(plan->Splits()); + } + + Result> CreateQueryReader( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + return CreateQueryReader(plan, realtime_context); + } + + Result ReadPlan(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context, + const std::vector& read_fields, + const std::shared_ptr& predicate, + bool enable_predicate_filter) const { ReadContextBuilder read_builder(table_path_); read_builder.SetOptions(options_) .SetReadFieldNames(read_fields) @@ -314,6 +1081,36 @@ class RealtimeWriteInteTest : public ::testing::Test { return CollectedReadResult{std::move(reader), std::move(result)}; } + void ReadPlanWithSchemaAndCheck(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context, + const std::shared_ptr& read_schema, + const std::string& expected_json) const { + std::unique_ptr c_read_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_read_schema.get()).ok()); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadSchema(std::move(c_read_schema)) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(reader.get())); + + arrow::FieldVector result_fields = {arrow::field("_VALUE_KIND", arrow::int8())}; + result_fields.insert(result_fields.end(), read_schema->fields().begin(), + read_schema->fields().end()); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), expected_json) + .ValueOrDie(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) + << result->ToString(); + } + Result> ReadRows( const std::shared_ptr& realtime_context) const { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, @@ -367,6 +1164,20 @@ class RealtimeWriteInteTest : public ::testing::Test { return ReadRows(/*realtime_context=*/nullptr); } + Result CountRows(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context) const { + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr count_reader, + table_read->CreateCountReader(plan->Splits())); + return count_reader->CountRows(); + } + Result GetRealtimeMemoryUsage( const std::shared_ptr& realtime_context) const { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, @@ -380,6 +1191,75 @@ class RealtimeWriteInteTest : public ::testing::Test { return memory_usage; } + Result> ReadPkSequences( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + realtime_context_impl->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected one PK real-time read view"); + } + auto read_schema = std::make_unique(); + arrow::FieldVector requested_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; + requested_fields.insert(requested_fields.end(), schema_->fields().begin(), + schema_->fields().end()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(requested_fields), read_schema.get())); + ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + views[0].store->CreateQueryReaders(views[0].read_view, + /*offset_begin=*/0, query_context)); + std::vector sequences; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr imported, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(imported); + if (!values) { + return Status::Invalid("PK query reader did not return a StructArray"); + } + std::shared_ptr sequence_array = + std::dynamic_pointer_cast( + values->GetFieldByName(SpecialFields::SequenceNumber().Name())); + if (!sequence_array) { + return Status::Invalid("PK query reader did not return sequence numbers"); + } + for (int64_t row = 0; row < sequence_array->length(); ++row) { + sequences.push_back(sequence_array->Value(row)); + } + } + reader->Close(); + } + return sequences; + } + + static std::vector> NewFiles( + const std::vector& progresses) { + std::vector> files; + for (const RealtimeCommitProgress& progress : progresses) { + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + if (!message) { + continue; + } + const std::vector>& new_files = + message->GetNewFilesIncrement().NewFiles(); + files.insert(files.end(), new_files.begin(), new_files.end()); + } + return files; + } + static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { std::vector seen(static_cast(total_rows), false); int64_t max_id = -1; @@ -405,6 +1285,8 @@ class RealtimeWriteInteTest : public ::testing::Test { return Status::OK(); } + void RunConcurrencyTest(bool primary_key); + Result ReadCommittedOffsets() const { PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); @@ -427,6 +1309,146 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_EQ(expected_rows, actual_rows); } + void ReplayPkWalAndCommit(const std::vector& wal, + const std::vector& row_kinds, + int64_t commit_identifier, + const std::vector& expected_rows) const { + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + } + + void CheckDropDatePartitionRemovesOffset(bool legacy_partition_name_enabled) { + fields_ = {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), + arrow::field("pt", arrow::date32())}; + schema_ = arrow::schema(fields_); + options_[Options::PARTITION_GENERATE_LEGACY_NAME] = + legacy_partition_name_enabled ? "true" : "false"; + CreateTable(/*partition_keys=*/{"pt"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + constexpr int32_t kDate = 19723; + constexpr int64_t kRowCount = 3; + const std::string partition = "2024-01-01"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeDatePartitionBatch(/*first_id=*/0, kRowCount, kDate, partition)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); + + const std::string normalized_partition = + legacy_partition_name_enabled ? std::to_string(kDate) : partition; + RealtimePartitionBucket partition_bucket({{"pt", normalized_partition}}, /*bucket=*/0); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_before_drop, ReadCommittedOffsets()); + ASSERT_EQ(1, offsets_before_drop.size()); + ASSERT_EQ(kRowCount, offsets_before_drop.at(partition_bucket)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_before_drop, + CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t rows_before_drop, + CountRows(plan_before_drop, /*realtime_context=*/nullptr)); + ASSERT_EQ(kRowCount, rows_before_drop); + + ASSERT_OK(DropPartition({{"pt", partition}}, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_after_drop, ReadCommittedOffsets()); + ASSERT_TRUE(offsets_after_drop.empty()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_after_drop, + CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t rows_after_drop, + CountRows(plan_after_drop, /*realtime_context=*/nullptr)); + ASSERT_EQ(0, rows_after_drop); + ASSERT_OK(writer->Close()); + } + + void CheckVectorReaderRetry(bool primary_key) { + if (primary_key) { + CreatePkTable(/*partition_keys=*/{"pt"}); + } else { + CreateTable(/*partition_keys=*/{"pt"}); + } + auto close_state = std::make_shared(); + auto factory = MakeDecoratingFactory(close_state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, + MakeBatch(p0_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p0_batch))); + std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, + MakeBatch(p1_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p1_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(2, plan->Splits().size()); + + std::vector> invalid_splits = plan->Splits(); + std::shared_ptr second_split = + std::dynamic_pointer_cast(invalid_splits[1]); + ASSERT_NE(nullptr, second_split); + std::vector> second_disk_splits = second_split->DiskSplits(); + invalid_splits[1] = std::make_shared( + RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), + second_split->Partition(), second_split->Bucket(), std::move(second_disk_splits), + second_split->CommittedEndOffset(), second_split->MemoryEndOffset(), + second_split->OpaqueTicket()); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), + "unsupported real-time split version"); + ASSERT_EQ(1, close_state->query_close_count->load(std::memory_order_acquire)); + + std::vector expected_rows = p0_rows; + expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); + } + + void CheckPkRejectsCommitReaderMalformation(CommitReaderMalformation malformation, + const std::string& expected_error) { + CreatePkTable(); + auto factory = MakeDecoratingFactory(malformation); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + expected_error); + ASSERT_OK(writer->Close()); + } + std::unique_ptr dir_; std::string table_path_; std::string commit_user_ = "realtime_commit_user"; @@ -436,6 +1458,44 @@ class RealtimeWriteInteTest : public ::testing::Test { std::shared_ptr pool_; }; +TEST_F(RealtimeWriteInteTest, TestRealtimeOperationsRequireEnabledOption) { + CreateTable(/*partition_keys=*/{}); + std::map disabled_options = options_; + disabled_options[Options::REALTIME_ENABLED] = "false"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + + WriteContextBuilder write_builder(table_path_, commit_user_); + write_builder.SetOptions(disabled_options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, write_builder.Finish()); + ASSERT_NOK_WITH_MSG(FileStoreWrite::Create(std::move(write_context)), + "real-time write requires realtime.enabled=true"); + + ScanContextBuilder scan_builder(table_path_); + scan_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "real-time scan requires realtime.enabled=true"); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), + "real-time read requires realtime.enabled=true"); + + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(disabled_options).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_NOK_WITH_MSG(commit->CommitWithProgress(/*realtime_commits=*/{}, + /*commit_identifier=*/0, + /*watermark=*/std::nullopt), + "CommitWithProgress requires realtime.enabled=true"); +} + TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); @@ -446,39 +1506,1266 @@ TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } -TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { - options_[Options::TARGET_FILE_ROW_NUM] = "10"; +TEST_F(RealtimeWriteInteTest, TestPkRead) { + CreatePkTable(); + auto saw_query_predicate = std::make_shared>(false); + auto query_view = std::make_shared>(); + auto factory = + MakeDecoratingFactory(saw_query_predicate, query_view); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}, {1, "new-in-run", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + MakeBatch({Row{1, "new", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(update_batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "new", "p0"}, {2, "two", "p0"}}), memory_rows); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + + std::vector second_rows = {{1, "latest", "p0"}, {2, "gone", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector union_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "latest", "p0"}, {3, "three", "p0"}}), union_rows); + + const std::string expected_payload = "new"; + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, expected_payload.data(), expected_payload.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_plan, + CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult filtered_result, + ReadPlan(filtered_plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/true)); + ASSERT_EQ(nullptr, filtered_result.data); + ASSERT_FALSE(saw_query_predicate->load(std::memory_order_acquire)); + filtered_result.reader->Close(); + filtered_result.reader.reset(); + ASSERT_OK(writer->Close()); + writer.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr lifetime_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + table_read->CreateReader(lifetime_plan->Splits())); + ASSERT_FALSE(query_view->expired()); + + std::weak_ptr weak_context = realtime_context; + table_read.reset(); + lifetime_plan.reset(); + realtime_context.reset(); + ASSERT_TRUE(weak_context.expired()); + ASSERT_FALSE(query_view->expired()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch read_batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(read_batch)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_array, + ReadResultCollector::GetArray(std::move(read_batch))); + ASSERT_NE(nullptr, read_array); + read_array.reset(); + reader->Close(); + reader.reset(); + ASSERT_TRUE(query_view->expired()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRealtimeReadOptimizedScanUnsupported) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + const std::vector rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector disk_rows, ReadRows()); + ASSERT_EQ(rows, disk_rows); + + ScanContextBuilder scan_builder(table_path_ + "$ro"); + scan_builder.SetOptions(options_).WithRealtimeContext(realtime_context).WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + Result> scan = TableScan::Create(std::move(scan_context)); + ASSERT_TRUE(scan.status().IsNotImplemented()) << scan.status().ToString(); + ASSERT_NE(std::string::npos, scan.status().ToString().find( + "PK real-time union read does not support read-optimized")); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkDeleteInsertAndPinnedReadsAcrossRefresh) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr delete_batch, + MakeBatch({Row{1, "deleted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(delete_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr insert_batch, + MakeBatch({Row{1, "inserted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(insert_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr pinned_reader, + table_read->CreateReader(reader_plan->Splits())); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector plan_rows, ReadRows(pinned_plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "inserted", "p0"}}), plan_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_rows, + ReadResultCollector::CollectResult(pinned_reader.get())); + ASSERT_EQ(1, reader_rows->length()); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { + options_[Options::READ_BATCH_SIZE] = "2"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}, {3, "disk-3", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + }; + int64_t commit_identifier = 0; + for (const std::vector& disk_rows : disk_batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ++commit_identifier; + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "sealed-1", "p0"}, Row{2, "deleted-2", "p0"}, Row{4, "sealed-4", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "active-1", "p0"}, Row{4, "deleted-4", "p0"}, Row{5, "active-5", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"payload", "id"}, /*predicate=*/nullptr, + /*enable_predicate_filter=*/false)); + ASSERT_NE(nullptr, result.data); + ASSERT_GT(result.data->num_chunks(), 1); + for (const std::shared_ptr& chunk : result.data->chunks()) { + ASSERT_LE(chunk->length(), 2); + } + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()), + arrow::field("id", arrow::int64())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "active-1", 1], + [0, "disk-3", 3], + [0, "active-5", 5], + [0, "disk-10", 10], + [0, "disk-11", 11] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkMergeAllDiskSplitsWithMemory) { + options_[Options::SOURCE_SPLIT_OPEN_FILE_COST] = "1"; + options_[Options::SOURCE_SPLIT_TARGET_SIZE] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + {{20, "disk-20", "p0"}, {21, "disk-21", "p0"}}, + }; + for (int64_t commit_identifier = 0; + commit_identifier < static_cast(disk_batches.size()); ++commit_identifier) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_batches[commit_identifier], /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr memory_batch, + MakeBatch({Row{1, "memory-1", "p0"}, Row{10, "deleted-10", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, realtime_split); + ASSERT_EQ(3, realtime_split->DiskSplits().size()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "memory-1", "p0"}, + {2, "disk-2", "p0"}, + {11, "disk-11", "p0"}, + {20, "disk-20", "p0"}, + {21, "disk-21", "p0"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { + const std::shared_ptr projected_b = arrow::field("b", arrow::int64()); + fields_ = { + arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::struct_({arrow::field("a", arrow::int64()), projected_b})), + arrow::field("pt", arrow::utf8()), + }; + schema_ = arrow::schema(fields_); + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + auto make_batch = [&](const std::string& json) -> Result> { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetBucket(0).Finish(); + }; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + make_batch(R"([[1, [101, 1001], "p0"], [2, [102, 1002], "p0"]])")); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr sealed_batch, + make_batch(R"([[1, [201, 2001], "p0"], [3, [203, 2003], "p0"]])")); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr active_batch, + make_batch(R"([[1, [301, 3001], "p0"], [4, [304, null], "p0"]])")); + ASSERT_OK(writer->Write(std::move(active_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + auto projected_schema = arrow::schema({ + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadSchema(std::move(c_schema)) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(reader.get())); + const std::shared_ptr result_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + const std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, [3001], 1], + [0, [1002], 2], + [0, [2003], 3], + [0, [null], 4] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) + << actual->ToString(); + reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkMemoryReadAfterSchemaEvolution) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "old", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + std::shared_ptr renamed_payload = arrow::field("renamed_payload", arrow::utf8()); + std::shared_ptr added = arrow::field("added", arrow::int32()); + ASSERT_OK(TestHelper::WriteNextSchema(dir_->GetFileSystem(), table_path_, + {DataField(0, fields_[0]), DataField(1, renamed_payload), + DataField(2, fields_[2]), DataField(3, added)}, + /*highest_field_id=*/3, options_)); + fields_[1] = renamed_payload; + fields_.push_back(added); + schema_ = arrow::schema(fields_); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "renamed_payload", "pt", "added"}, + /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + ASSERT_EQ(1, result.data->num_chunks()); + std::shared_ptr row = + std::dynamic_pointer_cast(result.data->chunk(0)); + ASSERT_NE(nullptr, row); + ASSERT_EQ(1, row->length()); + std::shared_ptr renamed_values = + std::dynamic_pointer_cast(row->field(2)); + ASSERT_NE(nullptr, renamed_values); + ASSERT_EQ("old", renamed_values->GetString(0)); + ASSERT_TRUE(row->field(4)->IsNull(0)); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { + CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({Row{1, "a", "disk-1a"}, Row{1, "b", "disk-1b"}, + Row{2, "a", "disk-2a"}, Row{3, "c", "disk-3c"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, disk_progress.size()); + ASSERT_EQ(OffsetRange(0, 4), disk_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(disk_progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "a", "sealed-1a"}, Row{1, "b", "deleted-1b"}, Row{2, "b", "sealed-2b"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + ASSERT_EQ(OffsetRange(4, 7), sealed_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(sealed_progress).size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "a", "active-1a"}, Row{1, "c", "active-1c"}, Row{2, "a", "active-2a"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + ASSERT_FALSE(split->DiskSplits().empty()); + ASSERT_EQ(4, split->CommittedEndOffset()); + ASSERT_EQ(10, split->MemoryEndOffset()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "a", "active-1a"}, + {1, "c", "active-1c"}, + {2, "a", "active-2a"}, + {2, "b", "sealed-2b"}, + {3, "c", "disk-3c"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector first_rows = { + {0, "value-0", "p0"}, {1, "value-1", "p0"}, {2, "value-2", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, first_progress.size()); + ASSERT_EQ(OffsetRange(0, 3), first_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(first_progress).size()); + ASSERT_EQ(0, NewFiles(first_progress)[0]->min_sequence_number); + ASSERT_EQ(2, NewFiles(first_progress)[0]->max_sequence_number); + ASSERT_OK(first_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector second_rows = {{0, "updated-0", "p0"}, {3, "value-3", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_progress.size()); + ASSERT_EQ(OffsetRange(3, 5), second_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(second_progress).size()); + ASSERT_EQ(3, NewFiles(second_progress)[0]->min_sequence_number); + ASSERT_EQ(4, NewFiles(second_progress)[0]->max_sequence_number); + + first_progress.push_back(std::move(second_progress[0])); + ASSERT_OK(Commit(first_progress, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{0, "updated-0", "p0"}, + {1, "value-1", "p0"}, + {2, "value-2", "p0"}, + {3, "value-3", "p0"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkPartitionBucketRecovery) { + options_[Options::BUCKET] = "2"; + CreatePkTable(/*partition_keys=*/{"pt"}); + const RealtimePartitionBucket p0b0({{"pt", "p0"}}, /*bucket=*/0); + const RealtimePartitionBucket p1b1({{"pt", "p1"}}, /*bucket=*/1); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_first_batch, + MakeBatch({Row{0, "p0-zero", "p0"}, Row{1, "p0-one", "p0"}}, + /*partitioned=*/true, /*bucket=*/0)); + ASSERT_OK(first_writer->Write(std::move(p0_first_batch))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p1_first_batch, + MakeBatch({Row{10, "p1-ten", "p1"}, Row{11, "p1-eleven", "p1"}, Row{12, "p1-twelve", "p1"}}, + /*partitioned=*/true, /*bucket=*/1)); + ASSERT_OK(first_writer->Write(std::move(p1_first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, first_progress.size()); + std::map first_ranges; + std::map> first_sequences; + for (const RealtimeCommitProgress& progress : first_progress) { + first_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + first_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(0, 2), first_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(0, 3), first_ranges.at(p1b1)); + ASSERT_EQ((std::make_pair(0, 1)), first_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(0, 2)), first_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, + Commit(first_progress, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->RefreshCommittedSnapshot(first_snapshot_id)); + ASSERT_OK(first_writer->Close()); + first_writer.reset(); + first_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p0_second_batch, + MakeBatch({Row{0, "p0-zero-new", "p0"}, Row{2, "p0-two", "p0"}}, + /*partitioned=*/true, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p0_second_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_second_batch, + MakeBatch({Row{10, "p1-ten-deleted", "p1"}, Row{13, "p1-thirteen", "p1"}}, + /*partitioned=*/true, /*bucket=*/1, + {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p1_second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, second_progress.size()); + std::map second_ranges; + std::map> second_sequences; + for (const RealtimeCommitProgress& progress : second_progress) { + second_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + second_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(2, 4), second_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(3, 5), second_ranges.at(p1b1)); + ASSERT_EQ((std::make_pair(2, 3)), second_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(3, 4)), second_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_progress, /*commit_identifier=*/1)); + ASSERT_OK(second_writer->RefreshCommittedSnapshot(second_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(second_context)); + std::sort(actual_rows.begin(), actual_rows.end()); + ASSERT_EQ((std::vector{{0, "p0-zero-new", "p0"}, + {1, "p0-one", "p0"}, + {2, "p0-two", "p0"}, + {11, "p1-eleven", "p1"}, + {12, "p1-twelve", "p1"}, + {13, "p1-thirteen", "p1"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); + + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(2, offsets.size()); + ASSERT_EQ(4, offsets.at(p0b0)); + ASSERT_EQ(5, offsets.at(p1b1)); +} + +TEST_F(RealtimeWriteInteTest, TestPkRecovery) { + CreatePkTable(); + + WriteContextBuilder seed_builder(table_path_, commit_user_); + seed_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_context, seed_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_writer, + FileStoreWrite::Create(std::move(seed_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_batch, + MakeBatch({Row{99, "seed", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(seed_writer->Write(std::move(seed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> seed_messages, + seed_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/0)); + CommitContextBuilder seed_commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit_context, + seed_commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit, + FileStoreCommit::Create(std::move(seed_commit_context))); + ASSERT_OK(seed_commit->Commit(seed_messages, /*commit_identifier=*/0)); + ASSERT_OK(seed_writer->Close()); + const std::vector mutations = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector mutation_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result> failed_prepare = + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1); + io_hook->Clear(); + ASSERT_TRUE(failed_prepare.status().IsIOError()) << failed_prepare.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failed_prepare, ReadRows()); + ASSERT_EQ((std::vector{{99, "seed", "p0"}}), rows_after_failed_prepare); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(first_writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_sequences, ReadPkSequences(first_context)); + ASSERT_EQ((std::vector{1, 2, 3, 4}), memory_sequences); + ASSERT_OK_AND_ASSIGN(std::vector progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_EQ(2, NewFiles(progress)[0]->min_sequence_number); + ASSERT_EQ(memory_sequences.back(), NewFiles(progress)[0]->max_sequence_number); + ASSERT_EQ(1, NewFiles(progress)[0]->delete_row_count); + ASSERT_OK(Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK(first_writer->Close()); + first_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_replay, ReadRows()); + ASSERT_EQ((std::vector{{1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}), + rows_after_replay); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr restart_batch, + MakeBatch({Row{4, "four", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(restart_batch))); + ASSERT_OK_AND_ASSIGN(std::vector restart_sequences, ReadPkSequences(second_context)); + ASSERT_EQ((std::vector{5}), restart_sequences); + ASSERT_OK_AND_ASSIGN(std::vector restart_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, restart_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), restart_progress[0].offset_range); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->min_sequence_number); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->max_sequence_number); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompaction) { + options_[Options::NUM_SORTED_RUNS_COMPACTION_TRIGGER] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + int64_t latest_snapshot_id = -1; + constexpr int64_t kCommitRoundsBeforeCompaction = 4; + std::set committed_file_names; + for (int64_t round = 0; round < kCommitRoundsBeforeCompaction; ++round) { + const bool delete_latest_live_row = round == kCommitRoundsBeforeCompaction - 1; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeBatch( + {Row{delete_latest_live_row ? round - 1 : round, + delete_latest_live_row ? "deleted" : "value-" + std::to_string(round), "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + delete_latest_live_row + ? std::vector{RecordBatch::RowKind::DELETE} + : std::vector{})); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(round)); + ASSERT_EQ(1, progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(progress[0].commit_message); + ASSERT_NE(nullptr, message); + ASSERT_TRUE(message->GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, NewFiles(progress).size()); + committed_file_names.insert(NewFiles(progress)[0]->file_name); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(progress, round)); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + } + WriteContextBuilder compact_builder(table_path_, commit_user_); + compact_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_context, compact_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_writer, + FileStoreWrite::Create(std::move(compact_context))); + ASSERT_OK(compact_writer->Compact(/*partition=*/{}, /*bucket=*/0, + /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> compact_messages, + compact_writer->PrepareCommit(/*wait_compaction=*/true, /*commit_identifier=*/4)); + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact_message = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact_message); + ASSERT_TRUE(compact_message->GetNewFilesIncrement().IsEmpty()); + ASSERT_EQ(kCommitRoundsBeforeCompaction, + compact_message->GetCompactIncrement().CompactBefore().size()); + std::set compacted_file_names; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactBefore()) { + compacted_file_names.insert(file->file_name); + } + ASSERT_EQ(committed_file_names, compacted_file_names); + ASSERT_FALSE(compact_message->GetCompactIncrement().CompactAfter().empty()); + constexpr int64_t kHistoricalMaxSequenceNumber = kCommitRoundsBeforeCompaction - 1; + int64_t compacted_live_max_sequence_number = -1; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactAfter()) { + compacted_live_max_sequence_number = + std::max(compacted_live_max_sequence_number, file->max_sequence_number); + } + ASSERT_LT(compacted_live_max_sequence_number, kHistoricalMaxSequenceNumber); + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK(commit->Commit(compact_messages, /*commit_identifier=*/4)); + ASSERT_OK(compact_writer->Close()); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional compact_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(compact_snapshot); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot->GetCommitKind()); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(4, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector compacted_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}}), compacted_rows); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr fresh_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_writer, + CreateRealtimeWriter(fresh_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_batch, + MakeBatch({Row{4, "value-4", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(fresh_writer->Write(std::move(fresh_batch))); + ASSERT_OK_AND_ASSIGN(std::vector fresh_sequences, ReadPkSequences(fresh_context)); + ASSERT_EQ((std::vector{compacted_live_max_sequence_number + 1}), fresh_sequences); + ASSERT_LT(fresh_sequences.front(), kHistoricalMaxSequenceNumber); + ASSERT_OK_AND_ASSIGN(std::vector fresh_progress, + fresh_writer->PrepareCommitWithProgress(/*commit_identifier=*/5)); + ASSERT_EQ(1, fresh_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), fresh_progress[0].offset_range); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->min_sequence_number); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->max_sequence_number); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(fresh_progress, /*commit_identifier=*/5)); + ASSERT_OK(fresh_writer->Close()); + + ASSERT_OK_AND_ASSIGN(offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}, {4, "value-4", "p0"}}), + final_rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkMultipleStoredBatchesMergeForQueryAndCommit) { + CreatePkTable(); + auto factory = MakeDecoratingFactory(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{4, "four", "p0"}, Row{2, "two", "p0"}, Row{1, "one", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr second_batch, + MakeBatch({Row{3, "three", "p0"}, Row{2, "deleted", "p0"}, Row{1, "one-new", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(second_batch))); + + const std::vector expected = {{1, "one-new", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector query_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected, query_rows); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 6), progress[0].offset_range); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector rows, ReadRows()); + ASSERT_EQ(expected, rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsMalformedCoverage) { + CheckPkRejectsCommitReaderMalformation(CommitReaderMalformation::DROP_LAST, + "commit readers did not cover the sealed range"); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsDuplicateOffset) { + CheckPkRejectsCommitReaderMalformation(CommitReaderMalformation::DUPLICATE_OFFSET, + "commit readers did not cover the sealed range"); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsOutOfRangeOffset) { + CheckPkRejectsCommitReaderMalformation(CommitReaderMalformation::OUT_OF_RANGE_OFFSET, + "offset is outside the sealed range"); +} + +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderClose) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = MakeDecoratingFactory(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + auto release_reader = [&](bool explicit_close) -> Status { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + CreateQueryReader(realtime_context)); + if (explicit_close) { + reader->Close(); + } + return Status::OK(); + }; + + ASSERT_OK(release_reader(/*explicit_close=*/true)); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); + + ASSERT_OK(release_reader(/*explicit_close=*/false)); + ASSERT_EQ(2, state->query_close_count->load(std::memory_order_acquire)); + + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = MakeDecoratingFactory(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch({Row{2, "two", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + + for (int32_t null_index = 0; null_index <= 1; ++null_index) { + state->query_null_index = null_index; + ASSERT_NOK_WITH_MSG(CreateQueryReader(realtime_context), + "PK real-time store returned a null query reader"); + ASSERT_EQ((null_index + 1) * 2, state->query_close_count->load(std::memory_order_acquire)); + } + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestAppendQueryReaderCloseFailure) { + CreateTable(/*partition_keys=*/{}); + auto state = std::make_shared(); + state->query_null_index = 1; + auto factory = MakeDecoratingFactory(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + ASSERT_NOK_WITH_MSG(CreateQueryReader(plan, realtime_context), + "append-only real-time store returned a null query reader"); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); + + state->query_null_index = -1; + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCommitReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + state->commit_null_index = 1; + auto factory = MakeDecoratingFactory(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PK real-time store returned a null commit reader"); + ASSERT_EQ(1, state->commit_close_count->load(std::memory_order_acquire)); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkPrepareFailureCleansPartialPhysicalFiles) { + options_[Options::WRITE_BATCH_SIZE] = "1"; + options_[Options::TARGET_FILE_ROW_NUM] = "1"; + options_["file-index.bitmap.columns"] = "payload"; + options_[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; + CreatePkTable(); + std::shared_ptr file_system = dir_->GetFileSystem(); + ASSERT_OK_AND_ASSIGN(std::set baseline_artifacts, + ListPhysicalArtifacts(file_system, dir_->Str())); + auto saw_artifacts = std::make_shared>(false); + auto factory = MakeDecoratingFactory( + file_system, dir_->Str(), baseline_artifacts.size(), saw_artifacts); + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create(factory)); + WriteContextBuilder failed_builder(table_path_, commit_user_); + failed_builder.SetOptions(options_) + .WithStreamingMode(true) + .WithRealtimeContext(failed_context) + .WithTempDirectory(PathUtil::JoinPath(dir_->Str(), "tmp")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_write_context, + failed_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + FileStoreWrite::Create(std::move(failed_write_context))); + + const std::vector wal = { + {1, "old", "p0"}, {1, "new", "p0"}, {2, "two", "p0"}, {2, "gone", "p0"}}; + const std::vector row_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + Result> failed_prepare = + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/0); + ASSERT_TRUE(failed_prepare.status().IsIOError()) << failed_prepare.status().ToString(); + ASSERT_NE(std::string::npos, + failed_prepare.status().ToString().find( + "injected commit reader failure after physical file creation")); + ASSERT_TRUE(saw_artifacts->load(std::memory_order_acquire)); + ASSERT_OK_AND_ASSIGN(std::set artifacts_after_abort, + ListPhysicalArtifacts(file_system, dir_->Str())); + ASSERT_EQ(baseline_artifacts, artifacts_after_abort); + + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + + const std::vector expected_rows = {{1, "new", "p0"}}; + ReplayPkWalAndCommit(wal, row_kinds, /*commit_identifier=*/0, expected_rows); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); + ASSERT_EQ(4, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); +} + +TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { + options_[Options::TARGET_FILE_ROW_NUM] = "10"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + + std::vector expected_rows; + constexpr int64_t kBatchCount = 3; + constexpr int64_t kRowsPerBatch = 10; + for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + std::vector rows = + MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + } + + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(OffsetRange(0, kBatchCount * kRowsPerBatch), commits[0].offset_range); + std::shared_ptr commit_message = + std::dynamic_pointer_cast(commits[0].commit_message); + ASSERT_NE(nullptr, commit_message); + ASSERT_EQ(3, commit_message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); +} + +TEST_F(RealtimeWriteInteTest, TestAppendScanKeepsDiskSplitsIndependent) { + options_[Options::TARGET_FILE_ROW_NUM] = "2"; + options_[Options::SOURCE_SPLIT_OPEN_FILE_COST] = "1"; + options_[Options::SOURCE_SPLIT_TARGET_SIZE] = "1"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector expected_rows; + for (int64_t first_id = 0; first_id < 6; first_id += 2) { + std::vector rows = MakeRows(first_id, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + } + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + std::vector memory_rows = MakeRows(/*first_id=*/6, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(3, plan->Splits().size()); + ASSERT_EQ(nullptr, std::dynamic_pointer_cast(plan->Splits()[0])); + ASSERT_EQ(nullptr, std::dynamic_pointer_cast(plan->Splits()[1])); + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(plan->Splits()[2]); + ASSERT_NE(nullptr, realtime_split); + ASSERT_EQ(1, realtime_split->DiskSplits().size()); + ASSERT_EQ(6, realtime_split->CommittedEndOffset()); + ASSERT_EQ(8, realtime_split->MemoryEndOffset()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + + std::vector first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(OffsetRange(0, 3), commits[0].offset_range); + + std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_commits.size()); + ASSERT_EQ(OffsetRange(3, 5), second_commits[0].offset_range); + + commits.push_back(std::move(second_commits[0])); + std::reverse(commits.begin(), commits.end()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK(writer->Close()); + + std::vector expected_rows = first_rows; + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); +} + +TEST_F(RealtimeWriteInteTest, TestCommitWithProgressRetryReturnsLatestSnapshot) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector expected_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(expected_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t retry_snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_EQ(first_snapshot_id, retry_snapshot_id); + + ASSERT_OK(writer->RefreshCommittedSnapshot(first_snapshot_id)); + std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_commits, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(retry_snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_EQ(second_snapshot_id, retry_snapshot_id); + ASSERT_NE(first_snapshot_id, retry_snapshot_id); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestCommitWithProgressRejectsCoveredRangesFromAnotherUser) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector expected_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(expected_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); + + CommitContextBuilder builder(table_path_, "another_realtime_commit_user"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + ASSERT_NOK_WITH_MSG(commit->CommitWithProgress(commits, /*commit_identifier=*/0, + /*watermark=*/std::nullopt), + "another commit user or identifier"); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestRealtimeWriteAcrossAppendCompaction) { + options_[Options::TARGET_FILE_ROW_NUM] = "2"; + options_[Options::COMPACTION_MIN_FILE_NUM] = "2"; CreateTable(/*partition_keys=*/{}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); - std::vector expected_rows; - constexpr int64_t kBatchCount = 3; - constexpr int64_t kRowsPerBatch = 10; - for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + std::vector first_rows; + for (int64_t first_id = 0; first_id < 5; first_id += 2) { std::vector rows = - MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); + MakeRows(first_id, std::min(2, 5 - first_id), /*partition=*/"p0"); ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(rows, /*partitioned=*/false)); ASSERT_OK(writer->Write(std::move(batch))); - expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + first_rows.insert(first_rows.end(), rows.begin(), rows.end()); } - - ASSERT_OK_AND_ASSIGN(std::vector commits, + ASSERT_OK_AND_ASSIGN(std::vector first_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); - ASSERT_EQ(1, commits.size()); - ASSERT_EQ(OffsetRange(0, kBatchCount * kRowsPerBatch), commits[0].offset_range); - std::shared_ptr commit_message = - std::dynamic_pointer_cast(commits[0].commit_message); - ASSERT_NE(nullptr, commit_message); - ASSERT_EQ(3, commit_message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); - ASSERT_OK(writer->Close()); + ASSERT_EQ(1, first_commits.size()); + ASSERT_EQ(OffsetRange(0, 5), first_commits[0].offset_range); + std::shared_ptr first_commit_message = + std::dynamic_pointer_cast(first_commits[0].commit_message); + ASSERT_NE(nullptr, first_commit_message); + ASSERT_EQ(3, first_commit_message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, Commit(first_commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(first_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(Snapshot compact_snapshot, CompactAndCommit(/*partition=*/{}, /*bucket=*/0, + /*commit_identifier=*/1)); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot.GetCommitKind()); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap compacted_offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, compacted_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + + ASSERT_OK(writer->RefreshCommittedSnapshot(compact_snapshot.Id())); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_compaction, ReadRows(realtime_context)); + ASSERT_EQ(first_rows, rows_after_compaction); + + std::vector second_rows = MakeRows(/*first_id=*/5, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + std::vector expected_rows = first_rows; + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector rows_with_building_memory, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows_with_building_memory); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); - ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, second_commits.size()); + ASSERT_EQ(OffsetRange(5, 7), second_commits[0].offset_range); + ASSERT_OK_AND_ASSIGN(int64_t final_snapshot_id, + Commit(second_commits, /*commit_identifier=*/2)); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap final_offsets, ReadCommittedOffsets()); + ASSERT_EQ(7, final_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector final_rows_before_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, final_rows_before_refresh); + ASSERT_OK(writer->RefreshCommittedSnapshot(final_snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector final_rows_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, final_rows_after_refresh); + ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { +TEST_F(RealtimeWriteInteTest, TestRealtimeOffsetFileLifecycle) { + options_[Options::SNAPSHOT_NUM_RETAINED_MIN] = "1"; + options_[Options::SNAPSHOT_NUM_RETAINED_MAX] = "1"; + options_[Options::SNAPSHOT_TIME_RETAINED] = "1ms"; CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); @@ -486,10 +2773,9 @@ TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, MakeBatch(first_rows, /*partitioned=*/false)); ASSERT_OK(writer->Write(std::move(first_batch))); - ASSERT_OK_AND_ASSIGN(std::vector commits, + ASSERT_OK_AND_ASSIGN(std::vector first_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); - ASSERT_EQ(1, commits.size()); - ASSERT_EQ(OffsetRange(0, 3), commits[0].offset_range); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, Commit(first_commits, /*commit_identifier=*/0)); std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, @@ -497,20 +2783,99 @@ TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { ASSERT_OK(writer->Write(std::move(second_batch))); ASSERT_OK_AND_ASSIGN(std::vector second_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_EQ(1, second_commits.size()); - ASSERT_EQ(OffsetRange(3, 5), second_commits[0].offset_range); - - commits.push_back(std::move(second_commits[0])); - std::reverse(commits.begin(), commits.end()); - ASSERT_OK(Commit(commits, /*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_commits, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options_)); + std::shared_ptr file_system = core_options.GetFileSystem(); + SnapshotManager snapshot_manager(file_system, table_path_); + ASSERT_OK_AND_ASSIGN(Snapshot first_snapshot, snapshot_manager.LoadSnapshot(first_snapshot_id)); + ASSERT_OK_AND_ASSIGN(Snapshot second_snapshot, + snapshot_manager.LoadSnapshot(second_snapshot_id)); + std::optional first_offsets_path = + RealtimeCommitProperties::GetOffsetsPath(first_snapshot); + std::optional second_offsets_path = + RealtimeCommitProperties::GetOffsetsPath(second_snapshot); + ASSERT_TRUE(first_offsets_path); + ASSERT_TRUE(second_offsets_path); + ASSERT_NE(first_offsets_path, second_offsets_path); + + std::string orphan_offsets_path = PathUtil::JoinPath( + RealtimeCommitProperties::OffsetsDirectory(table_path_, core_options.GetBranch()), + "orphan.offsets"); + ASSERT_OK(file_system->WriteFile(orphan_offsets_path, "orphan", /*overwrite=*/false)); + CleanContextBuilder clean_builder(table_path_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr clean_context, + clean_builder.WithFileSystem(file_system) + .WithOlderThanMs(std::numeric_limits::max()) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr cleaner, + OrphanFilesCleaner::Create(std::move(clean_context))); + ASSERT_OK_AND_ASSIGN(std::set cleaned_paths, cleaner->Clean()); + ASSERT_EQ(std::set({orphan_offsets_path}), cleaned_paths); + ASSERT_OK_AND_ASSIGN(bool first_offsets_exist, file_system->Exists(first_offsets_path.value())); + ASSERT_TRUE(first_offsets_exist); + ASSERT_OK_AND_ASSIGN(bool second_offsets_exist, + file_system->Exists(second_offsets_path.value())); + ASSERT_TRUE(second_offsets_exist); + + ASSERT_OK_AND_ASSIGN(int32_t expired_snapshots, ExpireSnapshots()); + ASSERT_EQ(1, expired_snapshots); + ASSERT_OK_AND_ASSIGN(first_offsets_exist, file_system->Exists(first_offsets_path.value())); + ASSERT_FALSE(first_offsets_exist); + ASSERT_OK_AND_ASSIGN(second_offsets_exist, file_system->Exists(second_offsets_path.value())); + ASSERT_TRUE(second_offsets_exist); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(second_snapshot, file_system)); ASSERT_EQ(5, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); ASSERT_OK(writer->Close()); +} - std::vector expected_rows = first_rows; - expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); - ASSERT_EQ(expected_rows, actual_rows); +TEST_F(RealtimeWriteInteTest, TestCompactionSnapshotRetainsSharedOffsetFile) { + options_[Options::TARGET_FILE_ROW_NUM] = "2"; + options_[Options::COMPACTION_MIN_FILE_NUM] = "2"; + options_[Options::SNAPSHOT_NUM_RETAINED_MIN] = "1"; + options_[Options::SNAPSHOT_NUM_RETAINED_MAX] = "1"; + options_[Options::SNAPSHOT_TIME_RETAINED] = "1ms"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + + for (int64_t first_id = 0; first_id < 5; first_id += 2) { + std::vector rows = + MakeRows(first_id, std::min(2, 5 - first_id), /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + } + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t realtime_snapshot_id, Commit(commits, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(Snapshot compact_snapshot, CompactAndCommit(/*partition=*/{}, /*bucket=*/0, + /*commit_identifier=*/1)); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot.GetCommitKind()); + + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options_)); + std::shared_ptr file_system = core_options.GetFileSystem(); + SnapshotManager snapshot_manager(file_system, table_path_); + ASSERT_OK_AND_ASSIGN(Snapshot realtime_snapshot, + snapshot_manager.LoadSnapshot(realtime_snapshot_id)); + std::optional realtime_offsets_path = + RealtimeCommitProperties::GetOffsetsPath(realtime_snapshot); + std::optional compact_offsets_path = + RealtimeCommitProperties::GetOffsetsPath(compact_snapshot); + ASSERT_TRUE(realtime_offsets_path); + ASSERT_TRUE(compact_offsets_path); + ASSERT_EQ(realtime_offsets_path, compact_offsets_path); + + ASSERT_OK_AND_ASSIGN(int32_t expired_snapshots, ExpireSnapshots()); + ASSERT_EQ(1, expired_snapshots); + ASSERT_OK_AND_ASSIGN(bool offsets_exist, file_system->Exists(compact_offsets_path.value())); + ASSERT_TRUE(offsets_exist); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(compact_snapshot, file_system)); + ASSERT_EQ(5, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK(writer->Close()); } TEST_F(RealtimeWriteInteTest, TestReadMemoryBeforePrepareCommit) { @@ -590,6 +2955,44 @@ TEST_F(RealtimeWriteInteTest, TestPlanExcludesRowsWrittenAfterMemoryEndOffset) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkRejectsReversedVisibleOffsets) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + std::vector> disk_splits = split->DiskSplits(); + std::vector> invalid_splits = {std::make_shared( + split->Version(), split->SnapshotId(), split->Partition(), split->Bucket(), + std::move(disk_splits), split->MemoryEndOffset() + 1, split->MemoryEndOffset(), + split->OpaqueTicket())}; + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), + "memory end offset precedes committed end offset"); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestReadFailsAfterRealtimeSplitTicketExpires) { options_[Options::REALTIME_READ_VIEW_TTL] = "10 ms"; CreateTable(/*partition_keys=*/{}); @@ -671,50 +3074,12 @@ TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTick ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestVectorReaderFailurePreservesEarlierSplitTicket) { - CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, - MakeBatch(p0_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p0_batch))); - std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, - MakeBatch(p1_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p1_batch))); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, - CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_EQ(2, plan->Splits().size()); - - std::vector> invalid_splits = plan->Splits(); - std::shared_ptr second_split = - std::dynamic_pointer_cast(invalid_splits[1]); - ASSERT_NE(nullptr, second_split); - std::vector> second_disk_splits = second_split->DiskSplits(); - invalid_splits[1] = std::make_shared( - RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), second_split->Partition(), - second_split->Bucket(), std::move(second_disk_splits), second_split->CommittedEndOffset(), - second_split->MemoryEndOffset(), second_split->OpaqueTicket()); - - ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(options_) - .SetReadFieldNames({"id", "payload", "pt"}) - .WithRealtimeContext(realtime_context) - .WithMemoryPool(pool_); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, - TableRead::Create(std::move(read_context))); - ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), - "unsupported real-time split version"); +TEST_F(RealtimeWriteInteTest, TestAppendVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/false); +} - std::vector expected_rows = p0_rows; - expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); - ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); +TEST_F(RealtimeWriteInteTest, TestPkVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/true); } TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { @@ -816,6 +3181,44 @@ TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestCountMemoryAndDiskAcrossRefresh) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr memory_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t memory_count, CountRows(memory_plan, realtime_context)); + ASSERT_EQ(3, memory_count); + + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id, + Commit(disk_commits, /*commit_identifier=*/0)); + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr union_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t union_count, CountRows(union_plan, realtime_context)); + ASSERT_EQ(5, union_count); + + ASSERT_OK(writer->RefreshCommittedSnapshot(committed_snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr refreshed_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t refreshed_count, CountRows(refreshed_plan, realtime_context)); + ASSERT_EQ(union_count, refreshed_count); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -855,76 +3258,380 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { ASSERT_OK_AND_ASSIGN(std::vector disk_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); - std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, - MakeBatch(memory_rows, /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(memory_batch))); + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr union_plan, + CreatePlan(realtime_context, scan_predicate)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult union_result, + ReadPlan(union_plan, realtime_context, read_fields, read_predicate, + /*enable_predicate_filter=*/true)); + std::shared_ptr expected_union = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "value-2", 2], + [0, "value-3", 3], + [0, "value-4", 4], + [0, "value-5", 5] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, union_result.data); + ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result.data)); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { + options_[Options::FILE_FORMAT] = "parquet"; + options_[Options::WRITE_BATCH_SIZE] = "1"; + options_["parquet.page.size"] = "1"; + options_["parquet.enable-dictionary"] = "false"; + options_["parquet.write.enable-page-index"] = "true"; + options_["parquet.read.enable-page-index-filter"] = "true"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::shared_ptr predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(1))); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/false)); + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, 1, "value-1", "p0"], + [0, 3, "value-3", "p0"], + [0, 4, "value-4", "p0"], + [0, 5, "value-5", "p0"] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, result.data); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdown) { + CreateTable(/*partition_keys=*/{}); + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); + auto make_expected = [&](const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(result_type, json).ValueOrDie(); + return std::make_shared(array); + }; + std::vector> predicates = { + PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(100))), + PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(5))), + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(10))), + }; + + auto check_candidates = + [&](const std::string& statistics_mode, + const std::vector>& expected) -> Status { + if (expected.size() != predicates.size()) { + return Status::Invalid("unexpected real-time candidate result count"); + } + options_[Options::REALTIME_STORE_STATS_MODE] = statistics_mode; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context, + RealtimeContext::Create()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + for (int64_t first_id : {0, 10, 20}) { + std::vector rows = MakeRows(first_id, /*count=*/3, /*partition=*/"p0"); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + PAIMON_RETURN_NOT_OK(writer->Write(std::move(batch))); + } + + for (size_t i = 0; i < predicates.size(); ++i) { + const std::shared_ptr& predicate = predicates[i]; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, predicate)); + PAIMON_ASSIGN_OR_RAISE( + CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/false)); + std::shared_ptr actual = + result.data ? result.data : make_expected("[]"); + if (!expected[i]->Equals(*actual)) { + return Status::Invalid("unexpected real-time candidate rows: " + + actual->ToString()); + } + } + PAIMON_RETURN_NOT_OK(writer->Close()); + return Status::OK(); + }; + + std::shared_ptr all_rows = make_expected(R"([ + [0, 0, "value-0", "p0"], + [0, 1, "value-1", "p0"], + [0, 2, "value-2", "p0"], + [0, 10, "value-10", "p0"], + [0, 11, "value-11", "p0"], + [0, 12, "value-12", "p0"], + [0, 20, "value-20", "p0"], + [0, 21, "value-21", "p0"], + [0, 22, "value-22", "p0"] + ])"); + ASSERT_OK(check_candidates("none", {all_rows, all_rows, all_rows})); + + std::shared_ptr partially_filtered = make_expected(R"([ + [0, 10, "value-10", "p0"], + [0, 11, "value-11", "p0"], + [0, 12, "value-12", "p0"], + [0, 20, "value-20", "p0"], + [0, 21, "value-21", "p0"], + [0, 22, "value-22", "p0"] + ])"); + std::shared_ptr matching_batch = make_expected(R"([ + [0, 10, "value-10", "p0"], + [0, 11, "value-11", "p0"], + [0, 12, "value-12", "p0"] + ])"); + ASSERT_OK(check_candidates("full", {make_expected("[]"), partially_filtered, matching_batch})); +} + +TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdownWithDisk) { + options_[Options::FILE_FORMAT] = "parquet"; + options_[Options::WRITE_BATCH_SIZE] = "1"; + options_[Options::REALTIME_STORE_STATS_MODE] = "full"; + options_["parquet.page.size"] = "1"; + options_["parquet.enable-dictionary"] = "false"; + options_["parquet.write.enable-page-index"] = "true"; + options_["parquet.read.enable-page-index-filter"] = "true"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/6, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr union_plan, - CreatePlan(realtime_context, scan_predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult union_result, - ReadPlan(union_plan, realtime_context, read_fields, read_predicate, - /*enable_predicate_filter=*/true)); - std::shared_ptr expected_union = + for (int64_t first_id : {0, 10}) { + std::vector rows = MakeRows(first_id, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + } + std::shared_ptr predicate = + PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/false)); + + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); + std::shared_ptr expected = arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ - [0, "value-2", 2], - [0, "value-3", 3], - [0, "value-4", 4], - [0, "value-5", 5] + [0, 4, "value-4", "p0"], + [0, 5, "value-5", "p0"], + [0, 10, "value-10", "p0"], + [0, 11, "value-11", "p0"], + [0, 12, "value-12", "p0"] ])") .ValueOrDie(); - ASSERT_NE(nullptr, union_result.data); - ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result.data)); + ASSERT_NE(nullptr, result.data); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { +TEST_F(RealtimeWriteInteTest, TestNullPredicateForMemoryAndDisk) { options_[Options::FILE_FORMAT] = "parquet"; options_[Options::WRITE_BATCH_SIZE] = "1"; + options_[Options::REALTIME_STORE_STATS_MODE] = "full"; options_["parquet.page.size"] = "1"; options_["parquet.enable-dictionary"] = "false"; options_["parquet.write.enable-page-index"] = "true"; + options_["parquet.write.max-row-group-length"] = "1"; options_["parquet.read.enable-page-index-filter"] = "true"; CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter(realtime_context)); - std::shared_ptr predicate = - PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, - Literal(static_cast(1))); - std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, - MakeBatch(disk_rows, /*partitioned=*/false)); + MakeUnpartitionedBatchFromJson(R"([ + [0, null, "p0"], + [1, "disk-value", "p0"] + ])")); ASSERT_OK(writer->Write(std::move(disk_batch))); ASSERT_OK_AND_ASSIGN(std::vector disk_commits, writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); - std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, - MakeBatch(memory_rows, /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(memory_batch))); - + ASSERT_OK_AND_ASSIGN(std::unique_ptr non_null_memory_batch, + MakeUnpartitionedBatchFromJson(R"([ + [2, "memory-value-2", "p0"], + [3, "memory-value-3", "p0"] + ])")); + ASSERT_OK(writer->Write(std::move(non_null_memory_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr nullable_memory_batch, + MakeUnpartitionedBatchFromJson(R"([ + [4, null, "p0"], + [5, "memory-value-5", "p0"] + ])")); + ASSERT_OK(writer->Write(std::move(nullable_memory_batch))); + + std::shared_ptr predicate = PredicateBuilder::IsNull( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); ASSERT_OK_AND_ASSIGN(CollectedReadResult result, ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, /*enable_predicate_filter=*/false)); + std::shared_ptr result_type = arrow::struct_( {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); std::shared_ptr expected = arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, 0, null, "p0"], + [0, 4, null, "p0"], + [0, 5, "memory-value-5", "p0"] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, result.data); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestUnionReadAfterColumnRename) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->Close()); + + std::shared_ptr renamed_payload = arrow::field("renamed_payload", arrow::utf8()); + ASSERT_OK(TestHelper::WriteNextSchema( + dir_->GetFileSystem(), table_path_, + {DataField(0, fields_[0]), DataField(1, renamed_payload), DataField(2, fields_[2])}, + /*highest_field_id=*/2, options_)); + fields_[1] = renamed_payload; + schema_ = arrow::schema(fields_); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(second_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, second_context, {"id", "renamed_payload", "pt"}, + /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), + renamed_payload, arrow::field("pt", arrow::utf8())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, 0, "value-0", "p0"], [0, 1, "value-1", "p0"], + [0, 2, "value-2", "p0"], [0, 3, "value-3", "p0"], - [0, 4, "value-4", "p0"], - [0, 5, "value-5", "p0"] + [0, 4, "value-4", "p0"] ])") .ValueOrDie(); ASSERT_NE(nullptr, result.data); ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) << result.data->ToString(); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestUnionReadWithNestedStructProjection) { + std::shared_ptr address_type = + arrow::struct_({arrow::field("city", arrow::utf8()), arrow::field("zip", arrow::int64())}); + std::shared_ptr profile_type = arrow::struct_( + {arrow::field("name", arrow::utf8()), arrow::field("address", address_type)}); + fields_ = {arrow::field("id", arrow::int64()), arrow::field("profile", profile_type), + arrow::field("pt", arrow::utf8())}; + schema_ = arrow::schema(fields_); + CreateTable(/*partition_keys=*/{}); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeUnpartitionedBatchFromJson(R"([ + [0, ["disk-0", ["hangzhou", 310000]], "p0"], + [1, ["disk-1", ["shanghai", 200000]], "p0"] + ])")); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t disk_snapshot_id, Commit(disk_commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(disk_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeUnpartitionedBatchFromJson(R"([ + [2, ["memory-2", ["beijing", 100000]], "p0"], + [3, ["memory-3", ["shenzhen", 518000]], "p0"] + ])")); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + std::shared_ptr projected_address_type = + arrow::struct_({arrow::field("city", arrow::utf8())}); + std::shared_ptr projected_profile_type = + arrow::struct_({arrow::field("address", projected_address_type)}); + std::shared_ptr projected_schema = arrow::schema( + {arrow::field("id", arrow::int64()), arrow::field("profile", projected_profile_type), + arrow::field("pt", arrow::utf8())}); + ReadPlanWithSchemaAndCheck(plan, realtime_context, projected_schema, R"([ + [0, 0, [["hangzhou"]], "p0"], + [0, 1, [["shanghai"]], "p0"], + [0, 2, [["beijing"]], "p0"], + [0, 3, [["shenzhen"]], "p0"] + ])"); ASSERT_OK(writer->Close()); } @@ -1109,8 +3816,167 @@ TEST_F(RealtimeWriteInteTest, TestRepeatedCommitReadAndRefresh) { ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { +TEST_F(RealtimeWriteInteTest, TestRefreshLatestSnapshotReclaimsMultipleCommittedSegments) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + constexpr int64_t kSnapshotCount = 3; + constexpr int64_t kRowsPerSnapshot = 2; + std::vector expected_rows; + int64_t latest_snapshot_id = -1; + for (int64_t snapshot_index = 0; snapshot_index < kSnapshotCount; ++snapshot_index) { + std::vector rows = + MakeRows(snapshot_index * kRowsPerSnapshot, kRowsPerSnapshot, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN( + std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/snapshot_index)); + ASSERT_EQ(1, commits.size()); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, + Commit(commits, /*commit_identifier=*/snapshot_index)); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + } + + ASSERT_OK_AND_ASSIGN(std::vector rows_before_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows_before_refresh); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_GT(memory_usage_before_refresh, 0); + + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::vector rows_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(rows_before_refresh, rows_after_refresh); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage_after_refresh); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestOverwriteRequiresReopenRealtimeContext) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector committed_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committed_batch, + MakeBatch(committed_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(committed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(committed_snapshot_id)); + + std::vector building_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr building_batch, + MakeBatch(building_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(building_batch))); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_overwrite, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_GT(memory_usage_before_overwrite, 0); + + ASSERT_OK_AND_ASSIGN(int64_t overwrite_snapshot_id, TruncateTable(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(core_options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(Snapshot overwrite_snapshot, + snapshot_manager.LoadSnapshot(overwrite_snapshot_id)); + ASSERT_EQ(Snapshot::CommitKind::Overwrite(), overwrite_snapshot.GetCommitKind()); + ASSERT_FALSE(RealtimeCommitProperties::GetOffsetsPath(overwrite_snapshot)); + + ASSERT_NOK_WITH_MSG(writer->RefreshCommittedSnapshot(overwrite_snapshot_id), + "recreate RealtimeContext"); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_failed_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(memory_usage_before_overwrite, memory_usage_after_failed_refresh); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(realtime_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(writer, CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(building_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(replay_batch))); + ASSERT_OK_AND_ASSIGN(std::vector replay_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, replay_commits.size()); + ASSERT_EQ(OffsetRange(0, 2), replay_commits[0].offset_range); + ASSERT_OK_AND_ASSIGN(int64_t replay_snapshot_id, + Commit(replay_commits, /*commit_identifier=*/2)); + ASSERT_OK(writer->RefreshCommittedSnapshot(replay_snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector replayed_rows, ReadRows(realtime_context)); + ASSERT_EQ(building_rows, replayed_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestReopenRealtimeContextAfterRollback) { CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, Commit(first_commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(first_snapshot_id)); + + std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_commits, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(second_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(int64_t rollback_snapshot_id, RollbackToAsLatest(first_snapshot_id)); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap rollback_offsets, ReadCommittedOffsets()); + ASSERT_EQ(1, rollback_offsets.size()); + ASSERT_EQ(3, rollback_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_NOK_WITH_MSG(writer->RefreshCommittedSnapshot(rollback_snapshot_id), + "recreate RealtimeContext"); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + // Reopen the same real-time writer identity from the target snapshot's progress and replay + // input after that restored boundary. + ASSERT_OK_AND_ASSIGN(realtime_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(writer, CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(replay_batch))); + ASSERT_OK_AND_ASSIGN(std::vector replay_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, replay_commits.size()); + ASSERT_EQ(OffsetRange(3, 5), replay_commits[0].offset_range); + ASSERT_OK(Commit(replay_commits, /*commit_identifier=*/2)); + + std::vector expected_rows = first_rows; + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); +} + +void RealtimeWriteInteTest::RunConcurrencyTest(bool primary_key) { + if (primary_key) { + CreatePkTable(); + } else { + CreateTable(/*partition_keys=*/{}); + } ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -1119,7 +3985,45 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { constexpr int32_t kReadThreadCount = 4; constexpr int64_t kBatchCount = 12; constexpr int64_t kRowsPerBatch = 2; - constexpr int64_t kTotalRows = kBatchCount * kRowsPerBatch; + const int64_t total_rows = kBatchCount * (primary_key ? 3 : kRowsPerBatch); + + std::vector> pk_batches; + std::vector> pk_row_kinds; + std::vector> pk_expected_states(1); + if (primary_key) { + std::map current_rows; + for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + const int64_t key = batch_index % 4; + const int64_t deleted_key = (key + 2) % 4; + std::vector rows = {{key, "update-" + std::to_string(batch_index), "p0"}, + {key, "latest-" + std::to_string(batch_index), "p0"}, + {deleted_key, "deleted-" + std::to_string(batch_index), "p0"}}; + pk_batches.push_back(rows); + pk_row_kinds.push_back({batch_index < 4 ? RecordBatch::RowKind::INSERT + : RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE}); + current_rows[key] = rows[1]; + current_rows.erase(deleted_key); + std::vector expected; + for (const auto& [id, row] : current_rows) { + static_cast(id); + expected.push_back(row); + } + pk_expected_states.push_back(std::move(expected)); + } + } + + auto validate_read = [&](const std::vector& rows) { + if (!primary_key) { + return ValidateReadPrefix(rows, total_rows); + } + if (std::find(pk_expected_states.begin(), pk_expected_states.end(), rows) == + pk_expected_states.end()) { + return Status::Invalid("PK real-time read does not match any completed write"); + } + return Status::OK(); + }; std::atomic writer_done{false}; std::atomic prepare_done{false}; @@ -1156,10 +4060,14 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { state.WaitForStart(); for (int64_t batch_index = 0; batch_index < kBatchCount && !state.ShouldStop(); ++batch_index) { - std::vector rows = - MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); + std::vector rows = primary_key + ? pk_batches[static_cast(batch_index)] + : MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, + /*partition=*/"p0"); Result> batch_result = - MakeBatch(rows, /*partitioned=*/false); + primary_key ? MakeBatch(rows, /*partitioned=*/false, /*bucket=*/0, + pk_row_kinds[static_cast(batch_index)]) + : MakeBatch(rows, /*partitioned=*/false); if (state.RecordErrorIfNotOk(batch_result)) { break; } @@ -1294,7 +4202,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { if (state.RecordErrorIfNotOk(result)) { break; } - Status status = ValidateReadPrefix(result.value(), kTotalRows); + Status status = validate_read(result.value()); if (state.RecordErrorIfNotOk(status)) { break; } @@ -1336,16 +4244,28 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { ASSERT_GE(commit_count.load(), 2); ASSERT_GE(refresh_count.load(), 2); ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); - ASSERT_EQ(kTotalRows, static_cast(final_rows.size())); - ASSERT_OK(ValidateReadPrefix(final_rows, kTotalRows)); + if (primary_key) { + ASSERT_EQ(pk_expected_states.back(), final_rows); + } else { + ASSERT_EQ(total_rows, static_cast(final_rows.size())); + ASSERT_OK(ValidateReadPrefix(final_rows, total_rows)); + } ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); - ASSERT_EQ(kTotalRows, + ASSERT_EQ(total_rows, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); ASSERT_EQ(0, memory_usage); ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { + RunConcurrencyTest(/*primary_key=*/false); +} + +TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { + RunConcurrencyTest(/*primary_key=*/true); +} + TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { CreateTable(/*partition_keys=*/{"pt"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -1392,9 +4312,168 @@ TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { } ASSERT_EQ(committed_offsets.end(), committed_offsets.find(RealtimePartitionBucket({{"pt", "p2"}}, /*bucket=*/0))); + + ASSERT_OK_AND_ASSIGN(std::vector final_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, final_commits.size()); + ASSERT_OK_AND_ASSIGN(int64_t final_snapshot_id, Commit(final_commits, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(final_snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector committed_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, committed_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestDropPartitionRequiresReopenRealtimeContext) { + CreateTable(/*partition_keys=*/{"pt"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + std::vector retained_disk_rows = + MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr retained_disk_batch, + MakeBatch(retained_disk_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(retained_disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t disk_snapshot_id, Commit(disk_commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(disk_snapshot_id)); + + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + std::vector rows_before_drop = disk_rows; + rows_before_drop.insert(rows_before_drop.end(), memory_rows.begin(), memory_rows.end()); + rows_before_drop.insert(rows_before_drop.end(), retained_disk_rows.begin(), + retained_disk_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows_before_drop, ReadRows(realtime_context)); + ASSERT_EQ(rows_before_drop, actual_rows_before_drop); + + ASSERT_OK_AND_ASSIGN(int64_t drop_snapshot_id, + DropPartition({{"pt", "p0"}}, /*commit_identifier=*/1)); + RealtimePartitionBucket partition_bucket({{"pt", "p0"}}, /*bucket=*/0); + RealtimePartitionBucket retained_partition_bucket({{"pt", "p1"}}, /*bucket=*/0); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_after_drop, ReadCommittedOffsets()); + ASSERT_EQ(1, offsets_after_drop.size()); + ASSERT_EQ(3, offsets_after_drop.at(retained_partition_bucket)); + ASSERT_EQ(offsets_after_drop.end(), offsets_after_drop.find(partition_bucket)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_NOK_WITH_MSG(writer->RefreshCommittedSnapshot(drop_snapshot_id), + "recreate RealtimeContext"); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(memory_usage_before_refresh, memory_usage_after_refresh); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + // Reopen with p1's retained progress and replay p0 input that existed only in the old context. + ASSERT_OK_AND_ASSIGN(realtime_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(writer, CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(memory_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(replay_batch))); + ASSERT_OK_AND_ASSIGN(std::vector replayed_rows, ReadRows(realtime_context)); + // The retained p1 disk split is read before the tail real-time split containing p0. + std::vector expected_replayed_rows = retained_disk_rows; + expected_replayed_rows.insert(expected_replayed_rows.end(), memory_rows.begin(), + memory_rows.end()); + ASSERT_EQ(expected_replayed_rows, replayed_rows); + + ASSERT_OK_AND_ASSIGN(std::vector memory_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, memory_commits.size()); + ASSERT_EQ(OffsetRange(0, 2), memory_commits[0].offset_range); + ASSERT_OK_AND_ASSIGN(int64_t memory_snapshot_id, + Commit(memory_commits, /*commit_identifier=*/2)); + ASSERT_OK(writer->RefreshCommittedSnapshot(memory_snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_memory_commit, ReadRows(realtime_context)); + std::vector expected_committed_rows = memory_rows; + expected_committed_rows.insert(expected_committed_rows.end(), retained_disk_rows.begin(), + retained_disk_rows.end()); + ASSERT_EQ(expected_committed_rows, rows_after_memory_commit); + ASSERT_OK_AND_ASSIGN(uint64_t final_memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, final_memory_usage); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap final_offsets, ReadCommittedOffsets()); + ASSERT_EQ(2, final_offsets.size()); + ASSERT_EQ(2, final_offsets.at(partition_bucket)); + ASSERT_EQ(3, final_offsets.at(retained_partition_bucket)); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestDropInactivePartitionDoesNotRequireReopenRealtimeContext) { + CreateTable(/*partition_keys=*/{"pt"}); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr seed_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_writer, + CreateRealtimeWriter(seed_context)); + for (int64_t partition_index = 0; partition_index < 2; ++partition_index) { + std::string partition = "p" + std::to_string(partition_index); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(MakeRows(partition_index * 10, /*count=*/3, partition), + /*partitioned=*/true)); + ASSERT_OK(seed_writer->Write(std::move(batch))); + } + ASSERT_OK_AND_ASSIGN(std::vector seed_commits, + seed_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, seed_commits.size()); + ASSERT_OK(Commit(seed_commits, /*commit_identifier=*/0)); + ASSERT_OK(seed_writer->Close()); + seed_writer.reset(); + seed_context.reset(); + + // The new context loads offsets for both partitions, but creates a store only for p0. + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, + MakeBatch(MakeRows(/*first_id=*/20, /*count=*/1, /*partition=*/"p0"), + /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p0_batch))); + + ASSERT_OK_AND_ASSIGN(int64_t drop_snapshot_id, + DropPartition({{"pt", "p1"}}, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(drop_snapshot_id)); + const RealtimePartitionBucket p0_partition_bucket({{"pt", "p0"}}, /*bucket=*/0); + const RealtimePartitionBucket p1_partition_bucket({{"pt", "p1"}}, /*bucket=*/0); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_after_drop, ReadCommittedOffsets()); + ASSERT_EQ(1, offsets_after_drop.size()); + ASSERT_EQ(3, offsets_after_drop.at(p0_partition_bucket)); + ASSERT_EQ(offsets_after_drop.end(), offsets_after_drop.find(p1_partition_bucket)); + + // Since p1 was never active in this context, writing it after the drop starts from zero. + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, + MakeBatch(MakeRows(/*first_id=*/30, /*count=*/2, /*partition=*/"p1"), + /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p1_batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(2, commits.size()); + auto p1_commit = + std::find_if(commits.begin(), commits.end(), [&](const RealtimeCommitProgress& commit) { + return commit.partition_bucket == p1_partition_bucket; + }); + ASSERT_NE(commits.end(), p1_commit); + ASSERT_EQ(OffsetRange(0, 2), p1_commit->offset_range); ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestDropDatePartitionRemovesOffsetWithLegacyPartitionName) { + CheckDropDatePartitionRemovesOffset(/*legacy_partition_name_enabled=*/true); +} + +TEST_F(RealtimeWriteInteTest, TestDropDatePartitionRemovesOffsetWithoutLegacyPartitionName) { + CheckDropDatePartitionRemovesOffset(/*legacy_partition_name_enabled=*/false); +} + TEST_F(RealtimeWriteInteTest, TestMultipleBucketsRestoreIndependentOffsets) { options_[Options::BUCKET] = "2"; CreateTable(/*partition_keys=*/{}); @@ -1491,4 +4570,167 @@ TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { ASSERT_EQ(5, second_committed_offsets.at(partition_bucket)); } +TEST_F(RealtimeWriteInteTest, TestPkWriteFailureRecovery) { + CreatePkTable(); + const std::vector seed_rows = {{99, "seed", "p0"}}; + ReplayPkWalAndCommit(seed_rows, /*row_kinds=*/{}, /*commit_identifier=*/0, seed_rows); + + const std::vector wal = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector row_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + std::shared_ptr failing_pool = + std::make_shared(pool_); + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + WriteContextBuilder failed_builder(table_path_, commit_user_); + failed_builder.SetOptions(options_) + .WithStreamingMode(true) + .WithRealtimeContext(failed_context) + .WithMemoryPool(failing_pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_write_context, + failed_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + FileStoreWrite::Create(std::move(failed_write_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr empty_batch, + MakeUnpartitionedBatchFromJson("[]")); + ASSERT_OK(failed_writer->Write(std::move(empty_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + failing_pool->FailAfterAllocations(1); + Status failed_write = failed_writer->Write(std::move(failed_batch)); + ASSERT_TRUE(failed_write.IsOutOfMemory()) << failed_write.ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(seed_rows, rows_after_failure); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr replay_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_writer, + CreateRealtimeWriter(replay_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(replay_writer->Write(std::move(replay_batch))); + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector replayed_rows, ReadRows(replay_context)); + ASSERT_EQ(expected_rows, replayed_rows); + ASSERT_OK_AND_ASSIGN(std::vector progress, + replay_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(1, 5), progress[0].offset_range); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_EQ(1, NewFiles(progress)[0]->delete_row_count); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK(replay_writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK(replay_writer->Close()); + replay_writer.reset(); + replay_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::vector persisted_rows, ReadRows()); + ASSERT_EQ(expected_rows, persisted_rows); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); +} + +TEST_F(RealtimeWriteInteTest, TestPkExternalCommitRecovery) { + CreatePkTable(); + const std::vector seed_rows = {{99, "seed", "p0"}}; + ReplayPkWalAndCommit(seed_rows, /*row_kinds=*/{}, /*commit_identifier=*/0, seed_rows); + + const std::vector wal = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector row_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector failed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, failed_progress.size()); + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result failed_commit = + commit->CommitWithProgress(failed_progress, /*commit_identifier=*/1, + /*watermark=*/std::nullopt); + io_hook->Clear(); + ASSERT_TRUE(failed_commit.status().IsIOError()) << failed_commit.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(seed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}; + ReplayPkWalAndCommit(wal, row_kinds, /*commit_identifier=*/1, expected_rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkRefreshRecovery) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + + const std::vector base_rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(base_rows, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t base_snapshot_id, Commit(base_progress, /*commit_identifier=*/0)); + ASSERT_OK(failed_writer->RefreshCommittedSnapshot(base_snapshot_id)); + + const std::vector committed_wal = { + {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector committed_kinds = {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr committed_batch, + MakeBatch(committed_wal, /*partitioned=*/false, /*bucket=*/0, committed_kinds)); + ASSERT_OK(failed_writer->Write(std::move(committed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector committed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id, + Commit(committed_progress, /*commit_identifier=*/1)); + + const std::vector replay_wal = {{4, "four", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(replay_wal, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(replay_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Status failed_refresh = failed_writer->RefreshCommittedSnapshot(committed_snapshot_id); + io_hook->Clear(); + ASSERT_TRUE(failed_refresh.IsIOError()) << failed_refresh.ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + const std::vector committed_rows = {{1, "one-new", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(committed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}; + ReplayPkWalAndCommit(replay_wal, /*row_kinds=*/{}, /*commit_identifier=*/2, expected_rows); +} + } // namespace paimon::test diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index d41b1a4a..5916fc97 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -309,6 +309,266 @@ TEST_P(WriteAndReadInteTest, TestAppendSimple) { ASSERT_TRUE(success); } +TEST_P(WriteAndReadInteTest, TestAppendVector) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("embedding", vector_type)}; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "-1"}, + {Options::FILE_SYSTEM, file_system}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + const std::string data_json = R"([ + [1, [1.0, 2.0, 3.0]], + [2, null], + [3, [4.0, 5.0, 6.0]] + ])"; + auto data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), data_json).ValueOrDie(); + auto c_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*data, c_array.get()).ok()); + RecordBatchBuilder batch_builder(c_array.get()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, batch_builder.SetBucket(0).Finish()); + ASSERT_OK_AND_ASSIGN(auto commit_messages, + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + (void)commit_messages; + + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + helper->ReadResult(data_splits)); + const std::string expected_json = R"([ + [0, 1, [1.0, 2.0, 3.0]], + [0, 2, null], + [0, 3, [4.0, 5.0, 6.0]] + ])"; + auto expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(result_fields), expected_json) + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(actual)); +} + +TEST_P(WriteAndReadInteTest, TestAppendNestedVector) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 2); + arrow::FieldVector fields = { + arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::struct_({arrow::field("embedding", vector_type)})), + arrow::field("history", arrow::list(vector_type)), + arrow::field("by_name", arrow::map(arrow::utf8(), vector_type)), + }; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "-1"}, + {Options::FILE_SYSTEM, file_system}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + const std::string data_json = R"([ + [1, [[1.0, 2.0]], [[3.0, 4.0], null], [["a", [5.0, 6.0]], ["b", null]]], + [2, [null], null, []], + [3, null, [], [["c", [7.0, 8.0]]]] + ])"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), data_json, + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + arrow::FieldVector result_fields = fields; + result_fields.insert(result_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); + const std::string expected_json = R"([ + [0, 1, [[1.0, 2.0]], [[3.0, 4.0], null], [["a", [5.0, 6.0]], ["b", null]]], + [0, 2, [null], null, []], + [0, 3, null, [], [["c", [7.0, 8.0]]]] + ])"; + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult(arrow::struct_(result_fields), + data_splits, expected_json)); + ASSERT_TRUE(success); +} + +// Pushing a predicate down on a non-vector column must not disturb the VECTOR column, whose +// read schema differs from the type stored in the data file. +TEST_P(WriteAndReadInteTest, TestAppendVectorWithPredicate) { + auto [file_format, file_system] = GetParam(); + if (file_format != "parquet") { + return; + } + + auto vector_type = + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 3); + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("embedding", vector_type)}; + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1048576"}, + {Options::BUCKET, "-1"}, + {Options::FILE_SYSTEM, file_system}, + // One row per row group, so the predicate prunes row groups instead of rows. + {"parquet.write.max-row-group-length", "1"}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + const std::string data_json = R"([ + [1, [1.0, 2.0, 3.0]], + [2, null], + [3, [4.0, 5.0, 6.0]], + [4, [7.0, 8.0, 9.0]] + ])"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), data_json, + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", + FieldType::INT, Literal(2)); + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.SetOptions(options) + .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()) + .SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, scan_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); + ASSERT_FALSE(result_plan->Splits().empty()); + + ReadContextBuilder read_context_builder(table_path); + read_context_builder.SetOptions(options).SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto 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(batch_reader.get())); + + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_with_row_kind), R"([ + [0, 3, [4.0, 5.0, 6.0]], + [0, 4, [7.0, 8.0, 9.0]] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); +} + +// TODO(jinli.zjw): move to a single file for a file index inte test +TEST_P(WriteAndReadInteTest, TestAppendWithExternalBitmapAndRangeBitmapIndexes) { + arrow::FieldVector fields = {arrow::field("name", arrow::utf8()), + arrow::field("score", arrow::int32())}; + auto [file_format, file_system] = GetParam(); + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, file_format}, + {Options::TARGET_FILE_SIZE, "1MB"}, + {Options::BUCKET, "-1"}, + {Options::FILE_SYSTEM, file_system}, + {"file-index.bitmap.columns", "name"}, + {"file-index.range-bitmap.columns", "score"}, + {"file-index.range-bitmap.score.chunk-size", "1KB"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1B"}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + + auto schema = arrow::schema(fields); + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(test_dir_, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([ + ["Alice", 10], + ["Bob", 20], + ["Alice", 30], + ["Lucy", 40] + ])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(auto data_files, CurrentDataFiles(options)); + ASSERT_EQ(1, data_files.size()); + const auto& [bucket_path, data_file] = data_files[0]; + ASSERT_FALSE(data_file->embedded_index); + ASSERT_EQ(1, data_file->extra_files.size()); + ASSERT_TRUE(data_file->extra_files[0]); + ASSERT_EQ(data_file->file_name + ".index", data_file->extra_files[0].value()); + std::string index_path = PathUtil::JoinPath(bucket_path, data_file->extra_files[0].value()); + ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); + ASSERT_TRUE(index_exists); + + std::string indexed_name = "Alice"; + auto name_predicate = PredicateBuilder::Equal( + /*field_index=*/0, /*field_name=*/"name", FieldType::STRING, + Literal(FieldType::STRING, indexed_name.data(), indexed_name.size())); + auto score_predicate = PredicateBuilder::GreaterThan( + /*field_index=*/1, /*field_name=*/"score", FieldType::INT, Literal(20)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({name_predicate, score_predicate})); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.SetOptions(options) + .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()) + .SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, scan_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(auto plan, table_scan->CreatePlan()); + ASSERT_EQ(1, plan->Splits().size()); + + // Keep precise post-read filtering disabled. The exact result therefore verifies that the + // bitmap and range-bitmap indexes produced by the write path are consumed by the read path. + ReadContextBuilder read_context_builder(table_path); + read_context_builder.SetOptions(options).SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto 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(plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(batch_reader.get())); + + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + auto expected_result = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields_with_row_kind), R"([[0, "Alice", 30]])"); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + auto expected = std::make_shared(expected_result.ValueOrDie()); + ASSERT_TRUE(expected->Equals(actual)) << actual->ToString(); +} + TEST_P(WriteAndReadInteTest, TestPKSimple) { arrow::FieldVector fields = { arrow::field("pk", arrow::utf8()), diff --git a/test/test_data/parquet/vector_compatibility/README.md b/test/test_data/parquet/vector_compatibility/README.md new file mode 100644 index 00000000..8a2fe25c --- /dev/null +++ b/test/test_data/parquet/vector_compatibility/README.md @@ -0,0 +1,36 @@ +# VECTOR Parquet compatibility fixtures + +These files pin the two physical Arrow schemas produced by Java and Rust writers for Paimon +VECTOR columns, with and without null vectors. + +- `java_vector.parquet` was copied from Apache Paimon Rust commit + `403a2b2e9bfc4ea66cd7e633619f1460efd18bc8`, path + `crates/paimon/testdata/pkvector/pk_vector_ivf_flat/bucket-0/data-932a1249-f7e0-4a03-8e1f-ab8c85cbb76f-0.parquet`. + The fixture documentation records Apache Paimon Java commit `7234e4c34` and + `PkVectorFixtureGenerator` as its source. Its VECTOR column is exposed as Arrow `list`. +- `java_vector_nullable.parquet` was generated with parquet-mr 1.15.1 (`parquet-avro` + `AvroParquetWriter` with `parquet.avro.write-old-list-structure=false`, so the column uses the + standard 3-level `list` / `element` layout Paimon Java writes). The rows are `(1, [1, 2, 3])`, + `(2, null)` and `(3, [4, 5, 6])`. The file carries no `ARROW:schema` key, so the VECTOR column + is exposed as Arrow `list`. +- `rust_vector.parquet` was generated with Apache Arrow Rust 58.4.0 using + `FixedSizeListBuilder` and `parquet::arrow::ArrowWriter`, the same Arrow and + Parquet representation used by Apache Paimon Rust. Its VECTOR column is exposed as Arrow + `fixed_size_list[3]`. The rows are `(1, [1, 2, 3])`, `(2, [7, 8, 9])`, and + `(3, [4, 5, 6])`. +- `rust_vector_nullable.parquet` was generated the same way, with the rows `(1, [1, 2, 3])`, + `(2, null)` and `(3, [4, 5, 6])`. + +A file that stores the Arrow schema, as the Rust writer does, is read back as +`fixed_size_list`. The bundled Arrow 17 patch backports the Arrow community fix that pads the +decoded child array for null fixed-size-list slots, so `rust_vector_nullable.parquet` is readable +as a nullable VECTOR. + +SHA-256 checksums: + +```text +2b2325cc2266301beaa2c78ec666cb5e0ee62283049de2a7231e3c9ae07bf3ca java_vector.parquet +42352e11daf5a291e8a8c4cfc8d0f0f6f8c9099cf7dcf24c28d9e159a29e0d8a java_vector_nullable.parquet +b5ba47e766ad72fca9c8485aa718ad27709c1fb4d34fb3670aa35e2001cbdbb0 rust_vector.parquet +f86058b1bc6cf803003446ca0abb7923e928d455fd39a7288c6d3acdd5fd10e9 rust_vector_nullable.parquet +``` diff --git a/test/test_data/parquet/vector_compatibility/java_vector.parquet b/test/test_data/parquet/vector_compatibility/java_vector.parquet new file mode 100644 index 00000000..5184c7a9 Binary files /dev/null and b/test/test_data/parquet/vector_compatibility/java_vector.parquet differ diff --git a/test/test_data/parquet/vector_compatibility/java_vector_nullable.parquet b/test/test_data/parquet/vector_compatibility/java_vector_nullable.parquet new file mode 100644 index 00000000..fa900ac0 Binary files /dev/null and b/test/test_data/parquet/vector_compatibility/java_vector_nullable.parquet differ diff --git a/test/test_data/parquet/vector_compatibility/rust_vector.parquet b/test/test_data/parquet/vector_compatibility/rust_vector.parquet new file mode 100644 index 00000000..761ee734 Binary files /dev/null and b/test/test_data/parquet/vector_compatibility/rust_vector.parquet differ diff --git a/test/test_data/parquet/vector_compatibility/rust_vector_nullable.parquet b/test/test_data/parquet/vector_compatibility/rust_vector_nullable.parquet new file mode 100644 index 00000000..8610e46e Binary files /dev/null and b/test/test_data/parquet/vector_compatibility/rust_vector_nullable.parquet differ