diff --git a/.asf.yaml b/.asf.yaml index 0cc39ed40..e4846c164 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/.github/workflows/build_and_test.yaml b/.github/workflows/build_and_test.yaml index b19e49069..0c8baf82a 100644 --- a/.github/workflows/build_and_test.yaml +++ b/.github/workflows/build_and_test.yaml @@ -75,7 +75,6 @@ jobs: - name: asan-ubsan-x86_64 build_args: --enable_asan --enable_ubsan - name: tsan-x86_64 - skip_rust: true build_args: --enable_tsan - name: gcc-debug-aarch64 runner: ubuntu-24.04-arm @@ -96,7 +95,6 @@ jobs: build_args: --build_type Release - name: tsan-aarch64 runner: ubuntu-24.04-arm - skip_rust: true build_args: --enable_tsan steps: - name: Checkout paimon-cpp @@ -108,8 +106,7 @@ jobs: uses: ./.github/actions/setup-ccache with: cache-key-prefix: ccache-${{ matrix.name }} - - name: Install Rust toolchain (tantivy-fts) - if: ${{ !matrix.skip_rust }} + - name: Install Rust toolchain (Mosaic and tantivy-fts) shell: bash run: ci/scripts/setup_rust.sh - name: Install HTTP and TLS development dependencies diff --git a/.github/workflows/gcc8_test.yaml b/.github/workflows/gcc8_test.yaml index aae2ce98b..c0b4c1593 100644 --- a/.github/workflows/gcc8_test.yaml +++ b/.github/workflows/gcc8_test.yaml @@ -60,6 +60,9 @@ jobs: ls -la - name: Checkout paimon-cpp uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + - name: Install Rust toolchain (Mosaic) + shell: bash + run: ci/scripts/setup_rust.sh - name: Setup ccache uses: ./.github/actions/setup-ccache with: diff --git a/.github/workflows/release_candidate.yaml b/.github/workflows/release_candidate.yaml index 6b5b1ebda..fdf9f5391 100644 --- a/.github/workflows/release_candidate.yaml +++ b/.github/workflows/release_candidate.yaml @@ -131,7 +131,7 @@ jobs: name: source-archive path: release/ci - - name: Install Rust toolchain (tantivy-fts) + - name: Install Rust toolchain (Mosaic and tantivy-fts) shell: bash run: ci/scripts/setup_rust.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index e99034dcc..2fe2df8e7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -60,7 +60,9 @@ option(PAIMON_USE_UBSAN "Use Undefined Behavior Sanitizer" OFF) option(PAIMON_USE_CXX11_ABI "Use C++11 ABI" ON) option(PAIMON_ENABLE_AVRO "Whether to enable avro file format" ON) option(PAIMON_ENABLE_ORC "Whether to enable orc file format" ON) +option(PAIMON_ENABLE_MOSAIC "Whether to enable mosaic file format (Rust FFI)" OFF) option(PAIMON_ENABLE_JINDO "Whether to enable jindo file system" OFF) +option(PAIMON_ENABLE_OSS "Whether to enable OSS SDK v2 file system" OFF) option(PAIMON_ENABLE_S3 "Whether to enable S3 file system" OFF) option(PAIMON_ENABLE_NETWORK_TESTS "Whether to enable tests that access real remote services over the network" OFF) @@ -68,23 +70,35 @@ option(PAIMON_ENABLE_LUMINA "Whether to enable lumina vector index" OFF) option(PAIMON_ENABLE_LUCENE "Whether to enable lucene index" OFF) option(PAIMON_ENABLE_TANTIVY "Whether to enable tantivy-fulltext global index (Rust FFI, experimental)" OFF) -option(PAIMON_ENABLE_REST "Whether to enable the rest catalog (requires libcurl)" OFF) +option(PAIMON_ENABLE_REST + "Whether to enable the rest catalog (requires libcurl and OpenSSL)" OFF) if(PAIMON_ENABLE_ORC) add_definitions(-DPAIMON_ENABLE_ORC) endif() if(PAIMON_ENABLE_REST) add_definitions(-DPAIMON_ENABLE_REST) endif() -# libcurl backs the HTTP client shared by the S3 file system and the rest catalog. -if(PAIMON_ENABLE_S3 OR PAIMON_ENABLE_REST) +# libcurl supports the REST and S3 HTTP clients, the OSS SDK transport, and object store timestamps. +if(PAIMON_ENABLE_OSS + OR PAIMON_ENABLE_S3 + OR PAIMON_ENABLE_REST) find_package(CURL REQUIRED) endif() +if(PAIMON_ENABLE_REST) + find_package(OpenSSL 1.1.0 REQUIRED) +endif() if(PAIMON_ENABLE_AVRO) add_definitions(-DPAIMON_ENABLE_AVRO) endif() +if(PAIMON_ENABLE_MOSAIC) + add_definitions(-DPAIMON_ENABLE_MOSAIC) +endif() if(PAIMON_ENABLE_JINDO) add_definitions(-DPAIMON_ENABLE_JINDO) endif() +if(PAIMON_ENABLE_OSS) + add_definitions(-DPAIMON_ENABLE_OSS) +endif() if(PAIMON_ENABLE_S3) add_definitions(-DPAIMON_ENABLE_S3) endif() @@ -384,6 +398,11 @@ if(PAIMON_BUILD_TESTS OR PAIMON_BUILD_BENCHMARKS) paimon_link_libraries_whole_archive(PAIMON_PARQUET_FILE_FORMAT_STATIC_LINK_LIBS paimon_parquet_file_format_static) + if(PAIMON_ENABLE_MOSAIC) + paimon_link_libraries_whole_archive(PAIMON_MOSAIC_FILE_FORMAT_STATIC_LINK_LIBS + paimon_mosaic_file_format_static) + endif() + if(PAIMON_ENABLE_ORC) paimon_link_libraries_whole_archive(PAIMON_ORC_FILE_FORMAT_STATIC_LINK_LIBS paimon_orc_file_format_static) @@ -435,6 +454,14 @@ if(PAIMON_BUILD_TESTS) paimon_link_libraries_whole_archive(PAIMON_PARQUET_FILE_FORMAT_STATIC_LINK_LIBS paimon_parquet_file_format_static) + if(PAIMON_ENABLE_MOSAIC) + paimon_link_libraries_no_as_needed(TEST_PLUGIN_LINK_LIBS + paimon_mosaic_file_format_shared) + list(APPEND TEST_STATIC_LINK_LIBS ${TEST_PLUGIN_LINK_LIBS}) + paimon_link_libraries_whole_archive(PAIMON_MOSAIC_FILE_FORMAT_STATIC_LINK_LIBS + paimon_mosaic_file_format_static) + endif() + if(PAIMON_ENABLE_ORC) paimon_link_libraries_no_as_needed(TEST_PLUGIN_LINK_LIBS paimon_orc_file_format_shared) @@ -452,6 +479,13 @@ if(PAIMON_BUILD_TESTS) paimon_jindo_file_system_shared) list(APPEND TEST_STATIC_LINK_LIBS ${TEST_PLUGIN_LINK_LIBS}) endif() + if(PAIMON_ENABLE_OSS) + paimon_link_libraries_whole_archive(PAIMON_OSS_FILE_SYSTEM_STATIC_LINK_LIBS + paimon_oss_file_system_static) + paimon_link_libraries_no_as_needed(TEST_PLUGIN_LINK_LIBS + paimon_oss_file_system_shared) + list(APPEND TEST_STATIC_LINK_LIBS ${TEST_PLUGIN_LINK_LIBS}) + endif() if(PAIMON_ENABLE_S3) paimon_link_libraries_whole_archive(PAIMON_S3_FILE_SYSTEM_STATIC_LINK_LIBS paimon_s3_file_system_static) @@ -511,11 +545,15 @@ add_subdirectory(src/paimon/fs/local) if(PAIMON_ENABLE_JINDO) add_subdirectory(src/paimon/fs/jindo) endif() +add_subdirectory(src/paimon/fs/oss) add_subdirectory(src/paimon/fs/s3) add_subdirectory(src/paimon/format/blob) add_subdirectory(src/paimon/format/orc) add_subdirectory(src/paimon/format/parquet) add_subdirectory(src/paimon/format/avro) +if(PAIMON_ENABLE_MOSAIC) + add_subdirectory(src/paimon/format/mosaic) +endif() if(PAIMON_ENABLE_LUMINA) add_subdirectory(src/paimon/global_index/lumina) endif() diff --git a/README.md b/README.md index 00a391c03..da928e713 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ Paimon C++ currently provides: - **Scan**: batch and stream scan for append tables and primary key tables without changelog. - **Read**: append table read, primary key table read with deletion vector, and primary key table merge-on-read. - **Arrow integration**: batch read and write interfaces based on the [Arrow Columnar In-Memory Format](https://arrow.apache.org). -- **File systems**: file system abstraction with built-in local and Jindo file system support. +- **File systems**: file system abstraction with built-in local, Jindo, OSS, and S3 file system support. - **File formats**: file format abstraction with built-in ORC, Parquet, and Avro support. - **Runtime utilities**: memory pool and thread pool abstractions with default implementations. - **AI-Oriented Features**: supports RowTracking and DataEvolution mode and provides Global Index diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 375e98961..0d4dc28cf 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -63,6 +63,20 @@ if(PAIMON_BUILD_BENCHMARKS) ${PAIMON_BENCHMARK_LINK_TOOLCHAIN} EXTRA_INCLUDES ${CMAKE_SOURCE_DIR}) + + add_paimon_benchmark(parquet_format_benchmark + SOURCES + parquet_format_benchmark.cpp + STATIC_LINK_LIBS + arrow + parquet + ${PAIMON_BENCHMARK_STATIC_LINK_LIBS} + test_utils_static + Threads::Threads + ${PAIMON_BENCHMARK_PLATFORM_LINK_LIBS} + ${PAIMON_BENCHMARK_LINK_TOOLCHAIN} + EXTRA_INCLUDES + ${CMAKE_SOURCE_DIR}) endif() if(PAIMON_BUILD_TESTS) @@ -74,4 +88,21 @@ if(PAIMON_BUILD_TESTS) STATIC_LINK_LIBS paimon_shared ${GTEST_LINK_TOOLCHAIN}) + + # Guards the format-layer assumptions parquet_format_benchmark.cpp is built on. The benchmark + # itself is only compiled under PAIMON_BUILD_BENCHMARKS, which CI does not set, so this test is + # what keeps those assumptions covered. + add_paimon_test(parquet_format_benchmark_test + SOURCES + parquet_format_benchmark_test.cpp + EXTRA_INCLUDES + ${CMAKE_SOURCE_DIR} + STATIC_LINK_LIBS + arrow + parquet + paimon_shared + ${PAIMON_LOCAL_FILE_SYSTEM_SHARED_LINK_LIBS} + ${PAIMON_PARQUET_FILE_FORMAT_STATIC_LINK_LIBS} + test_utils_static + ${GTEST_LINK_TOOLCHAIN}) endif() diff --git a/benchmark/parquet_format_benchmark.cpp b/benchmark/parquet_format_benchmark.cpp new file mode 100644 index 000000000..e8d3e0026 --- /dev/null +++ b/benchmark/parquet_format_benchmark.cpp @@ -0,0 +1,1444 @@ +/* + * 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. + */ + +// Format-level micro-benchmarks for the Parquet reader and writer. These drive +// ParquetWriterBuilder / ParquetFileBatchReader directly, so a format-layer change can be +// attributed without the catalog lookup, split planning, merge/sort and commit that the +// table-level read_write_benchmark includes. Each case comment says what that case answers. +// +// Every axis - type, cardinality, null density, batch size, encoding, selectivity - is swept on +// its own rather than as a combined matrix, because the point is attributing one change rather +// than describing a workload. +// +// Data is generated outside the timed region and on Arrow's default pool, because the writer +// cuts a new row group once its own pool crosses parquet.writer.max.memory.use. +// +// IO goes through the local FileSystem into a temporary directory - the project has no in-memory +// FileSystem - so absolute numbers are only meaningful relative to each other. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "arrow/util/bit_util.h" +#include "benchmark/benchmark.h" +#include "paimon/common/utils/arrow/arrow_input_stream_adapter.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/path_util.h" +#include "paimon/defs.h" +#include "paimon/format/format_writer.h" +#include "paimon/format/parquet/parquet_field_id_converter.h" +#include "paimon/format/parquet/parquet_file_batch_reader.h" +#include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/parquet/parquet_writer_builder.h" +#include "paimon/fs/file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/metrics.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/result.h" +#include "paimon/status.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" + +namespace { + +using ::paimon::ArrowInputStreamAdapter; +using ::paimon::BatchReader; +using ::paimon::FieldType; +using ::paimon::FileStatus; +using ::paimon::FileSystem; +using ::paimon::FormatWriter; +using ::paimon::InputStream; +using ::paimon::Literal; +using ::paimon::OutputStream; +using ::paimon::PathUtil; +using ::paimon::Predicate; +using ::paimon::PredicateBuilder; +using ::paimon::Result; +using ::paimon::RoaringBitmap32; +using ::paimon::Status; +using ::paimon::parquet::ParquetFieldIdConverter; +using ::paimon::parquet::ParquetFileBatchReader; +using ::paimon::parquet::ParquetWriterBuilder; + +constexpr int64_t kRowsPerFile = 100'000; +constexpr int64_t kRowsPerBatch = 10'000; +constexpr int32_t kReadBatchSize = 4096; +constexpr int32_t kWriteBatchSize = 1024; +// Small enough that a 100K-row file spans many pages, so page-level pruning has something to +// prune. Arrow's page limit is byte-based; a row-count limit is not available yet. +constexpr int64_t kPageSizeBytes = 64 * 1024; +// Four row groups per read fixture, so row-group pruning and page pruning are both in play. +constexpr int64_t kRowGroupLength = 25'000; +constexpr int64_t kStringCardinality = 1'000; +// Few enough distinct values that arrow keeps the column dictionary-encoded for the whole file, +// which is the shape the wide-schema case wants: per-column work small, per-batch cost visible. +constexpr int64_t kLowStringCardinality = 10; +constexpr int32_t kVectorDimension = 16; +constexpr int32_t kListLength = 4; +// Same as kListLength, so the MAP and LIST cases differ only in the extra key leaf. +constexpr int32_t kMapEntries = kListLength; +constexpr char kDefaultCompression[] = "zstd"; +// Nulls are placed over a 100-row window, so a requested density is exact, not statistical. +constexpr int64_t kNullWindow = 100; + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id) { + return arrow::field(name, type, + arrow::KeyValueMetadata::Make({ParquetFieldIdConverter::PARQUET_FIELD_ID}, + {std::to_string(field_id)})); +} + +std::shared_ptr StructColumnType() { + return arrow::struct_({arrow::field("a", arrow::int64()), arrow::field("b", arrow::utf8())}); +} + +std::shared_ptr ListColumnType() { + return arrow::list(arrow::int64()); +} + +std::shared_ptr VectorColumnType() { + return arrow::fixed_size_list(arrow::field("element", arrow::float32(), /*nullable=*/false), + kVectorDimension); +} + +std::shared_ptr MapColumnType() { + return arrow::map(arrow::utf8(), arrow::int64()); +} + +std::shared_ptr DictionaryStringType() { + return arrow::dictionary(arrow::int32(), arrow::utf8()); +} + +std::shared_ptr DictionaryInt32Type() { + return arrow::dictionary(arrow::int32(), arrow::int32()); +} + +// The three-column file the flat read cases scan. `id` is ordered so a range predicate maps +// onto a contiguous row range, which is what makes page-index pruning measurable. +std::shared_ptr FlatSchema() { + return arrow::schema({MakeField("id", arrow::int64(), 0), MakeField("name", arrow::utf8(), 1), + MakeField("amount", arrow::decimal128(18, 4), 2)}); +} + +std::shared_ptr DecimalSchema(int32_t precision) { + return arrow::schema({MakeField("amount", arrow::decimal128(precision, 4), 0)}); +} + +std::shared_ptr DoubleSchema() { + return arrow::schema({MakeField("value", arrow::float64(), 0)}); +} + +std::shared_ptr NestedSchema() { + return arrow::schema( + {MakeField("id", arrow::int64(), 0), MakeField("info", StructColumnType(), 1), + MakeField("tags", ListColumnType(), 2), MakeField("embedding", VectorColumnType(), 3), + MakeField("attrs", MapColumnType(), 4)}); +} + +// A VECTOR is stored as a Parquet LIST and the reader hands back the file's own types, so a +// format-level read asks for the physical type; VectorFileBatchReader restores the view above. +std::shared_ptr NestedReadSchema() { + return arrow::schema({MakeField("id", arrow::int64(), 0), + MakeField("info", StructColumnType(), 1), + MakeField("tags", ListColumnType(), 2), + MakeField("embedding", + arrow::list(arrow::field("element", arrow::float32(), + /*nullable=*/false)), + 3), + MakeField("attrs", MapColumnType(), 4)}); +} + +Result> MakeInt64Column(int64_t num_rows, int64_t offset) { + arrow::Int64Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(offset + i); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +Result> MakeDoubleColumn(int64_t num_rows, int64_t offset) { + arrow::DoubleBuilder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(offset + i) * 1.5); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +Result> MakeBooleanColumn(int64_t num_rows, int64_t offset) { + arrow::BooleanBuilder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(((offset + i) & 1) == 0); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// Low cardinality is what arrow dictionary-encodes; high cardinality falls back to plain. +Result> MakeStringColumn(int64_t num_rows, int64_t offset, + int64_t cardinality) { + arrow::StringBuilder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + const int64_t value = (offset + i) % cardinality; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append("value_" + std::to_string(value))); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// INT32 cycling through `cardinality` distinct values, the flat control for +// BM_ParquetWrite_DictionaryInt32: same logical values, same cardinality, same width, so the +// delta between the two is the dictionary materialization alone. +Result> MakeInt32Column(int64_t num_rows, int64_t offset, + int64_t cardinality) { + arrow::Int32Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(((offset + i) % cardinality) * 7)); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// Dictionary-encoded input: `values->length()` distinct values behind an int32 index array. Arrow +// hands the indices straight to Parquet when the value type is binary-like +// (DictionaryDirectWriteSupported) and densifies them otherwise, so the value type alone decides +// whether the writer materializes anything. +Result> MakeDictionaryColumn( + const std::shared_ptr& values, int64_t num_rows, int64_t offset) { + const int64_t cardinality = values->length(); + arrow::Int32Builder index_builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + index_builder.UnsafeAppend(static_cast((offset + i) % cardinality)); + } + std::shared_ptr indices; + PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Finish(&indices)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::DictionaryArray::FromArrays(indices, values)); + return array; +} + +// STRING values: binary-like, so arrow can write the indices directly. +Result> MakeDictionaryStringColumn(int64_t num_rows, int64_t offset, + int64_t cardinality) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, + MakeStringColumn(cardinality, /*offset=*/0, cardinality)); + return MakeDictionaryColumn(values, num_rows, offset); +} + +// INT32 values: is_base_binary_like excludes them, so arrow densifies before writing. The +// dictionary holds the same `i * 7` values MakeInt32Column emits inline, so the two are directly +// comparable. +Result> MakeDictionaryInt32Column(int64_t num_rows, int64_t offset, + int64_t cardinality) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, + MakeInt32Column(cardinality, /*offset=*/0, cardinality)); + return MakeDictionaryColumn(values, num_rows, offset); +} + +// Precision drives the Parquet physical type: <= 9 is INT32, <= 18 is INT64, larger is +// FIXED_LEN_BYTE_ARRAY, and the three take different transfer paths on read. +Result> MakeDecimalColumn(int64_t num_rows, int64_t offset, + int32_t precision, int32_t scale) { + arrow::Decimal128Builder builder(arrow::decimal128(precision, scale)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + for (int64_t i = 0; i < num_rows; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(arrow::Decimal128(offset + i))); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +Result> MakeStructArray( + const arrow::FieldVector& fields, const std::vector>& columns) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::StructArray::Make(columns, fields)); + return paimon::checked_pointer_cast(array); +} + +Result> MakeStructColumn(int64_t num_rows, int64_t offset) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr a, MakeInt64Column(num_rows, offset)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr b, + MakeStringColumn(num_rows, offset, kStringCardinality)); + return MakeStructArray(StructColumnType()->fields(), {a, b}); +} + +// Constant element count per row, so the delta against a flat BIGINT column is the levels. +Result> MakeListColumn(int64_t num_rows, int64_t offset) { + auto value_builder = std::make_shared(); + arrow::ListBuilder builder(arrow::default_memory_pool(), value_builder, ListColumnType()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Reserve(num_rows * kListLength)); + for (int64_t i = 0; i < num_rows; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append()); + for (int32_t j = 0; j < kListLength; ++j) { + value_builder->UnsafeAppend(offset + i + j); + } + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// The writer converts VECTOR to a Parquet LIST, so this also covers ParquetVectorConverter. +Result> MakeVectorColumn(int64_t num_rows, int64_t offset) { + auto value_builder = std::make_shared(); + arrow::FixedSizeListBuilder builder(arrow::default_memory_pool(), value_builder, + VectorColumnType()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Reserve(num_rows * kVectorDimension)); + for (int64_t i = 0; i < num_rows; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append()); + for (int32_t j = 0; j < kVectorDimension; ++j) { + value_builder->UnsafeAppend(static_cast((offset + i) % 1024) + + static_cast(j)); + } + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// A Parquet MAP is a LIST of a two-field STRUCT, so against LIST the difference is the key leaf. +Result> MakeMapColumn(int64_t num_rows, int64_t offset, + int32_t entries) { + auto key_builder = std::make_shared(); + auto item_builder = std::make_shared(); + arrow::MapBuilder builder(arrow::default_memory_pool(), key_builder, item_builder, + MapColumnType()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(num_rows)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(key_builder->Reserve(num_rows * entries)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(item_builder->Reserve(num_rows * entries)); + for (int64_t i = 0; i < num_rows; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append()); + for (int32_t j = 0; j < entries; ++j) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(key_builder->Append("key_" + std::to_string(j))); + item_builder->UnsafeAppend(offset + i + j); + } + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} + +// Masks `null_pct` percent of the slots null. Doing it after the fact rather than in every +// generator keeps the values identical across densities, so the delta between two settings is the +// level machinery alone. Only the top level is masked: a STRUCT or LIST row, not its leaves. +Result> WithNulls(const std::shared_ptr& array, + int64_t offset, int64_t null_pct) { + if (null_pct <= 0) { + return array; + } + // Builder output starts at slot 0, so the bitmap below can be indexed by position. + if (array->data()->offset != 0) { + return Status::Invalid("WithNulls expects an unsliced array"); + } + const int64_t length = array->length(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr validity, + arrow::AllocateEmptyBitmap(length, arrow::default_memory_pool())); + int64_t null_count = 0; + for (int64_t i = 0; i < length; ++i) { + if ((offset + i) % kNullWindow < null_pct) { + ++null_count; + } else { + arrow::bit_util::SetBit(validity->mutable_data(), i); + } + } + std::shared_ptr data = array->data()->Copy(); + data->buffers[0] = std::move(validity); + data->SetNullCount(null_count); + return arrow::MakeArray(data); +} + +using BatchFactory = std::function>( + const std::shared_ptr& schema, int64_t offset, int64_t rows)>; + +using ColumnFactory = + std::function>(int64_t rows, int64_t offset)>; + +ColumnFactory NullableColumnFactory(const ColumnFactory& make_column, int64_t null_pct) { + return [make_column, null_pct](int64_t rows, + int64_t offset) -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr column, make_column(rows, offset)); + return WithNulls(column, offset, null_pct); + }; +} + +// Lifts a one-column generator into a batch factory for a one-field schema. +BatchFactory SingleColumnBatch(const ColumnFactory& make_column) { + return [make_column](const std::shared_ptr& schema, int64_t offset, + int64_t rows) -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr column, make_column(rows, offset)); + return MakeStructArray(schema->fields(), {column}); + }; +} + +Result> MakeNullableFlatBatch( + const std::shared_ptr& schema, int64_t offset, int64_t rows, int64_t null_pct) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr ids, MakeInt64Column(rows, offset)); + PAIMON_ASSIGN_OR_RAISE(ids, WithNulls(ids, offset, null_pct)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr names, + MakeStringColumn(rows, offset, kStringCardinality)); + PAIMON_ASSIGN_OR_RAISE(names, WithNulls(names, offset, null_pct)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr amounts, + MakeDecimalColumn(rows, offset, /*precision=*/18, /*scale=*/4)); + PAIMON_ASSIGN_OR_RAISE(amounts, WithNulls(amounts, offset, null_pct)); + return MakeStructArray(schema->fields(), {ids, names, amounts}); +} + +Result> MakeFlatBatch(const std::shared_ptr& schema, + int64_t offset, int64_t rows) { + return MakeNullableFlatBatch(schema, offset, rows, /*null_pct=*/0); +} + +// One low-cardinality VARCHAR per field, cheap enough that a wide schema leaves the per-column +// setup holding the measurement. The staggered offset keeps the columns from being identical. +Result> MakeWideBatch(const std::shared_ptr& schema, + int64_t offset, int64_t rows) { + std::vector> columns; + columns.reserve(schema->num_fields()); + for (int32_t i = 0; i < schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr column, + MakeStringColumn(rows, offset + i, kLowStringCardinality)); + columns.push_back(std::move(column)); + } + return MakeStructArray(schema->fields(), columns); +} + +Result> MakeNestedBatch(const std::shared_ptr& schema, + int64_t offset, int64_t rows) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr ids, MakeInt64Column(rows, offset)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr info, MakeStructColumn(rows, offset)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr tags, MakeListColumn(rows, offset)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr embedding, MakeVectorColumn(rows, offset)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr attrs, + MakeMapColumn(rows, offset, kMapEntries)); + return MakeStructArray(schema->fields(), {ids, info, tags, embedding, attrs}); +} + +// Whether every AddBatch call gets its own array or they all share one. It matters for dictionary +// input: arrow compares each batch's dictionary against the previous one, so N distinct-but-equal +// dictionaries and one dictionary written N times are not the same write. Only +// BM_ParquetWrite_MemoryThreshold reuses one batch; every other case generates each batch +// independently. That is about the arrays, not the values - generators that cycle with a period +// dividing the batch size, such as the boolean and the low-cardinality string ones, produce +// batches that are equal in value but separate objects. +enum class BatchReuse { kFresh, kReused }; + +Result>> MakeBatches( + const std::shared_ptr& schema, const BatchFactory& make_batch, + int64_t rows_per_batch, int64_t total_rows = kRowsPerFile, + BatchReuse reuse = BatchReuse::kFresh) { + if (reuse == BatchReuse::kReused) { + // One array written N times cannot express a short final batch, and silently writing a + // full one instead would put more rows in the file than the reported metrics divide by. + if (total_rows % rows_per_batch != 0) { + return Status::Invalid("BatchReuse::kReused needs total_rows divisible by " + + std::to_string(rows_per_batch) + ", got " + + std::to_string(total_rows)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr batch, + make_batch(schema, /*offset=*/0, rows_per_batch)); + return std::vector>(total_rows / rows_per_batch, + std::move(batch)); + } + std::vector> batches; + for (int64_t offset = 0; offset < total_rows; offset += rows_per_batch) { + const int64_t rows = std::min(rows_per_batch, total_rows - offset); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr batch, + make_batch(schema, offset, rows)); + batches.push_back(std::move(batch)); + } + return batches; +} + +Result WriteParquetFile(const std::shared_ptr& fs, const std::string& path, + const std::shared_ptr& schema, + const std::vector>& batches, + const std::map& options, + const std::string& compression) { + ParquetWriterBuilder writer_builder(schema, kWriteBatchSize, options); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr out, fs->Create(path, /*overwrite=*/true)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + writer_builder.Build(out, compression)); + for (const auto& batch : batches) { + // AddBatch imports - and so consumes - the C array, so each batch needs a fresh export. + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, &c_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array)); + } + PAIMON_RETURN_NOT_OK(writer->Finish()); + PAIMON_RETURN_NOT_OK(out->Close()); + PAIMON_ASSIGN_OR_RAISE(FileStatus file_status, fs->GetFileStatus(path)); + return file_status.GetLen(); +} + +struct ReadStats { + int64_t rows = 0; + // Reader-side counterpart of file size: makes pruning and skipping visible apart from CPU. + uint64_t storage_bytes = 0; + // Without these a filtered case shows only that it got faster, not that pruning is why. + uint64_t row_groups_total = 0; + uint64_t row_groups_after_filter = 0; + uint64_t batches = 0; +}; + +// Every counter read below is set unconditionally by the reader - the row-group pair in +// SetReadSchema, the batch count in NextBatch - so an absent one means the metric moved or stopped +// being recorded. Reporting that as a zero would hide the regression behind a plausible number, +// so the error propagates and fails the case instead. +Result ReadCounter(const std::shared_ptr& metrics, + const std::string& name) { + return metrics->GetCounter(name); +} + +Result ReadParquetFile(const std::shared_ptr& fs, const std::string& path, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::optional& selection_bitmap, + const std::map& options, + int32_t batch_size, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(FileStatus file_status, fs->GetFileStatus(path)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, fs->Open(path)); + auto in_stream = std::make_shared(input, file_status.GetLen(), pool); + // Held separately so the counter outlives the adapter the reader takes ownership of. + std::shared_ptr> storage_read_bytes = in_stream->StorageReadBytes(); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + ParquetFileBatchReader::Create(std::move(in_stream), options, batch_size, + /*file_metadata=*/nullptr, storage_read_bytes, pool, + /*hints=*/std::nullopt)); + + // SetReadSchema imports the C schema and takes ownership of it. + ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_schema)); + PAIMON_RETURN_NOT_OK(reader->SetReadSchema(&c_schema, predicate, selection_bitmap)); + + ReadStats stats; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + stats.rows += batch.first->length; + ArrowArrayRelease(batch.first.get()); + ArrowSchemaRelease(batch.second.get()); + } + + std::shared_ptr metrics = reader->GetReaderMetrics(); + PAIMON_ASSIGN_OR_RAISE( + stats.row_groups_total, + ReadCounter(metrics, paimon::parquet::ParquetMetrics::READ_ROW_GROUPS_TOTAL)); + PAIMON_ASSIGN_OR_RAISE( + stats.row_groups_after_filter, + ReadCounter(metrics, paimon::parquet::ParquetMetrics::READ_ROW_GROUPS_AFTER_FILTER)); + PAIMON_ASSIGN_OR_RAISE(stats.batches, + ReadCounter(metrics, paimon::parquet::ParquetMetrics::READ_BATCH_COUNT)); + reader->Close(); + stats.storage_bytes = storage_read_bytes->load(); + return stats; +} + +// Row groups the written file actually ended up with. The writer decides that itself - by row +// count, or by its pool crossing parquet.writer.max.memory.use - so the footer is the only +// reliable source. +// +// The write schema is deliberately not used to read back, because for two cases it does not +// describe the file. A VECTOR is stored as a LIST, and CollectLeafIndices branches on the file +// type, so a FixedSizeList read type against a file LIST is rejected outright. A dictionary column +// is stored as its value type; that one survives leaf collection, which compares nothing for +// atomic fields, and would fail later against read_data_type_ once a batch is read. Create() +// resolves the file's own schema and sets the row-group counters while doing so, which is all this +// needs - no read schema, no batch read, and neither trap. +Result ReadRowGroupCount(const std::shared_ptr& fs, const std::string& path, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(FileStatus file_status, fs->GetFileStatus(path)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, fs->Open(path)); + auto in_stream = std::make_shared(input, file_status.GetLen(), pool); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + ParquetFileBatchReader::Create(std::move(in_stream), /*options=*/{}, kReadBatchSize, + /*file_metadata=*/nullptr, /*storage_read_bytes=*/nullptr, + pool, /*hints=*/std::nullopt)); + PAIMON_ASSIGN_OR_RAISE(uint64_t row_groups, + ReadCounter(reader->GetReaderMetrics(), + paimon::parquet::ParquetMetrics::READ_ROW_GROUPS_TOTAL)); + reader->Close(); + return row_groups; +} + +// Owns a temporary directory plus the shared FileSystem / MemoryPool, removed on destruction. +class BenchmarkEnv { + public: + static Result> Create() { + std::unique_ptr dir = + paimon::test::UniqueTestDirectory::Create(); + if (!dir) { + return Status::IOError("failed to create a temporary benchmark directory"); + } + return std::unique_ptr(new BenchmarkEnv(std::move(dir))); + } + + const std::shared_ptr& fs() const { + return fs_; + } + + // The reader takes an arrow pool; the writer builder allocates its own from the paimon pool. + const std::shared_ptr& arrow_pool() const { + return arrow_pool_; + } + + std::string PathOf(const std::string& file_name) const { + return PathUtil::JoinPath(dir_->Str(), file_name); + } + + private: + explicit BenchmarkEnv(std::unique_ptr dir) + : dir_(std::move(dir)), + fs_(dir_->GetFileSystem()), + arrow_pool_(paimon::GetArrowPool(paimon::GetDefaultPool())) {} + + std::unique_ptr dir_; + std::shared_ptr fs_; + std::shared_ptr arrow_pool_; +}; + +// A file the read cases scan, built once per configuration; rebuilding it would dominate. +class ReadFixture { + public: + ReadFixture(const std::string& file_name, const std::shared_ptr& schema, + const BatchFactory& make_batch, + const std::map& write_options = {}) { + status_ = Build(file_name, schema, make_batch, write_options); + } + + const Status& status() const { + return status_; + } + + const std::string& path() const { + return path_; + } + + const std::shared_ptr& fs() const { + return env_->fs(); + } + + const std::shared_ptr& arrow_pool() const { + return env_->arrow_pool(); + } + + int64_t file_bytes() const { + return file_bytes_; + } + + private: + Status Build(const std::string& file_name, const std::shared_ptr& schema, + const BatchFactory& make_batch, + const std::map& write_options) { + PAIMON_ASSIGN_OR_RAISE(env_, BenchmarkEnv::Create()); + path_ = env_->PathOf(file_name); + PAIMON_ASSIGN_OR_RAISE(std::vector> batches, + MakeBatches(schema, make_batch, kRowsPerBatch)); + std::map options = write_options; + options[paimon::parquet::PARQUET_PAGE_SIZE] = std::to_string(kPageSizeBytes); + options[paimon::parquet::PARQUET_WRITE_MAX_ROW_GROUP_LENGTH] = + std::to_string(kRowGroupLength); + PAIMON_ASSIGN_OR_RAISE(file_bytes_, WriteParquetFile(env_->fs(), path_, schema, batches, + options, kDefaultCompression)); + return Status::OK(); + } + + std::unique_ptr env_; + std::string path_; + int64_t file_bytes_ = 0; + Status status_; +}; + +// Built on first use, so a case excluded by --benchmark_filter never pays to write its file. +// google/benchmark may report from a different thread than it registered on, hence the lock. +const ReadFixture& GetFixture(const std::string& key, + const std::function()>& build) { + static std::mutex mutex; + static std::map> fixtures; + std::lock_guard guard(mutex); + std::unique_ptr& fixture = fixtures[key]; + if (!fixture) { + fixture = build(); + } + return *fixture; +} + +const ReadFixture& FlatFixture() { + return GetFixture("flat", [] { + return std::make_unique("flat.parquet", FlatSchema(), &MakeFlatBatch); + }); +} + +const ReadFixture& NestedFixture() { + return GetFixture("nested", [] { + return std::make_unique("nested.parquet", NestedSchema(), &MakeNestedBatch); + }); +} + +const ReadFixture& NullableFlatFixture(int64_t null_pct) { + const std::string key = "flat_nulls_" + std::to_string(null_pct); + return GetFixture(key, [key, null_pct] { + return std::make_unique( + key + ".parquet", FlatSchema(), + [null_pct](const std::shared_ptr& schema, int64_t offset, int64_t rows) { + return MakeNullableFlatBatch(schema, offset, rows, null_pct); + }); + }); +} + +// A one-column file, for read cases that isolate a single decoder. +const ReadFixture& ColumnFixture(const std::string& key, + const std::shared_ptr& schema, + const ColumnFactory& make_column) { + return GetFixture(key, [key, schema, make_column] { + return std::make_unique(key + ".parquet", schema, + SingleColumnBatch(make_column)); + }); +} + +// The flat fixture only carries precision 18, so without this the FIXED_LEN_BYTE_ARRAY path that +// precision 38 takes is written but never read. +const ReadFixture& DecimalFixture(int32_t precision) { + return ColumnFixture("decimal_" + std::to_string(precision), DecimalSchema(precision), + [precision](int64_t rows, int64_t offset) { + return MakeDecimalColumn(rows, offset, precision, /*scale=*/4); + }); +} + +// The flat schema carries no floating-point column, so a DOUBLE read needs a file of its own. +const ReadFixture& DoubleFixture() { + return ColumnFixture("double", DoubleSchema(), &MakeDoubleColumn); +} + +// The same data with dictionary encoding off, giving the read side a plain baseline. +const ReadFixture& PlainFlatFixture() { + return GetFixture("flat_plain", [] { + std::map options; + options[paimon::parquet::PARQUET_ENABLE_DICTIONARY] = "false"; + return std::make_unique("flat_plain.parquet", FlatSchema(), &MakeFlatBatch, + options); + }); +} + +// google/benchmark's own main exits 0 whatever happened, so a case that called SkipWithError +// prints as skipped while `ctest -L benchmark` still passes. Nothing here skips on purpose, so +// the flag every error sets becomes the process exit code. It is recorded here rather than in a +// custom reporter because passing one to RunSpecifiedBenchmarks would bypass +// CreateDefaultDisplayReporter and with it --benchmark_format, --benchmark_color and +// --benchmark_counters_tabular. +std::atomic g_failed{false}; + +bool FailBenchmark(::benchmark::State& state, const Status& status) { + if (status.ok()) { + return false; + } + g_failed.store(true); + state.SkipWithError(status.ToString().c_str()); + return true; +} + +class Timer { + public: + double ElapsedNanos() const { + return std::chrono::duration(std::chrono::steady_clock::now() - started_) + .count(); + } + + private: + std::chrono::steady_clock::time_point started_ = std::chrono::steady_clock::now(); +}; + +// google/benchmark reports items/s; ns per row is what the issue asks for, so the timed region +// is also measured directly and divided by the row count. +void ReportRowRate(::benchmark::State& state, int64_t rows, double elapsed_ns) { + state.SetItemsProcessed(static_cast(state.iterations()) * rows); + const double total_rows = static_cast(state.iterations()) * static_cast(rows); + if (total_rows > 0) { + state.counters["ns_per_row"] = ::benchmark::Counter(elapsed_ns / total_rows); + } +} + +// Bytes per iteration, plus the per-row form that shows a CPU-for-size trade next to ns_per_row. +void ReportBytes(::benchmark::State& state, const std::string& name, int64_t bytes, int64_t rows) { + state.counters[name] = ::benchmark::Counter(static_cast(bytes)); + if (rows > 0) { + state.counters["bytes_per_row"] = + ::benchmark::Counter(static_cast(bytes) / static_cast(rows)); + } +} + +// Everything the writer needs per file - properties, output stream, schema conversion, footer - +// is inside the timed region, because that is what a caller pays per data file. +void RunWriteBenchmark(::benchmark::State& state, const std::shared_ptr& schema, + const BatchFactory& make_batch, int64_t rows_per_batch, + const std::map& options, + const std::string& compression, int64_t total_rows = kRowsPerFile, + BatchReuse reuse = BatchReuse::kFresh) { + Result> env = BenchmarkEnv::Create(); + if (FailBenchmark(state, env.status())) { + return; + } + Result>> batches = + MakeBatches(schema, make_batch, rows_per_batch, total_rows, reuse); + if (FailBenchmark(state, batches.status())) { + return; + } + const std::string path = env.value()->PathOf("write_case.parquet"); + + int64_t file_bytes = 0; + Timer timer; + for (auto _ : state) { + Result written = WriteParquetFile(env.value()->fs(), path, schema, batches.value(), + options, compression); + if (FailBenchmark(state, written.status())) { + return; + } + file_bytes = written.value(); + } + // Captured before the read-back below, which reopens the file and parses its footer. Reading + // the timer after would put that inside ns_per_row but not inside google/benchmark's own + // real_time, leaving the two disagreeing by a fixed amount that matters at low iteration + // counts. + const double elapsed_ns = timer.ElapsedNanos(); + + // Read back rather than computed: with a byte-triggered flush the count is not predictable + // from the arguments, and where a row-count limit does make it predictable, reporting the real + // number is what would catch the prediction being wrong. + Result row_groups = + ReadRowGroupCount(env.value()->fs(), path, env.value()->arrow_pool()); + if (FailBenchmark(state, row_groups.status())) { + return; + } + + ReportRowRate(state, total_rows, elapsed_ns); + ReportBytes(state, "file_bytes", file_bytes, total_rows); + state.counters["batches"] = ::benchmark::Counter( + static_cast((total_rows + rows_per_batch - 1) / rows_per_batch)); + state.counters["row_groups"] = ::benchmark::Counter(static_cast(row_groups.value())); +} + +// Single-column variant: one column isolates one encoder. +void RunColumnWriteBenchmark(::benchmark::State& state, const std::shared_ptr& field, + const ColumnFactory& make_column, int64_t rows_per_batch, + const std::map& options, + int64_t total_rows = kRowsPerFile, + BatchReuse reuse = BatchReuse::kFresh) { + RunWriteBenchmark(state, arrow::schema({field}), SingleColumnBatch(make_column), rows_per_batch, + options, kDefaultCompression, total_rows, reuse); +} + +void BM_ParquetWrite_Int64(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("id", arrow::int64(), 0), &MakeInt64Column, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_Double(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("value", arrow::float64(), 0), &MakeDoubleColumn, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_Boolean(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("flag", arrow::boolean(), 0), &MakeBooleanColumn, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_String(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunColumnWriteBenchmark(state, MakeField("name", arrow::utf8(), 0), + [cardinality](int64_t rows, int64_t offset) { + return MakeStringColumn(rows, offset, cardinality); + }, + kRowsPerBatch, /*options=*/{}); +} + +// arg: number of distinct values. The flat control for BM_ParquetWrite_DictionaryInt32. +void BM_ParquetWrite_FlatInt32(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunColumnWriteBenchmark(state, MakeField("value", arrow::int32(), 0), + [cardinality](int64_t rows, int64_t offset) { + return MakeInt32Column(rows, offset, cardinality); + }, + kRowsPerBatch, /*options=*/{}); +} + +// arg: dictionary cardinality. Same values as BM_ParquetWrite_String at the same cardinality, but +// handed to the writer already dictionary-encoded - nothing in ParquetWriterBuilder rejects a +// dictionary arrow type, so the writer does reach this path - and the pair isolates what it saves +// when it can pass indices through instead of materializing every value. +void BM_ParquetWrite_DictionaryString(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunColumnWriteBenchmark(state, MakeField("name", DictionaryStringType(), 0), + [cardinality](int64_t rows, int64_t offset) { + return MakeDictionaryStringColumn(rows, offset, cardinality); + }, + kRowsPerBatch, /*options=*/{}); +} + +// The same axis on an INTEGER dictionary, which arrow cannot direct-write - is_base_binary_like +// excludes int32, so it densifies first. Its baseline is BM_ParquetWrite_FlatInt32 at the same +// cardinality, not the String case: only the flat INT32 control holds value, width and encoding +// fixed, so only that delta is the materialization cost. +void BM_ParquetWrite_DictionaryInt32(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunColumnWriteBenchmark(state, MakeField("value", DictionaryInt32Type(), 0), + [cardinality](int64_t rows, int64_t offset) { + return MakeDictionaryInt32Column(rows, offset, cardinality); + }, + kRowsPerBatch, /*options=*/{}); +} + +// Same data with dictionary encoding turned off, as an encoding baseline. +void BM_ParquetWrite_StringNoDictionary(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + std::map options; + options[paimon::parquet::PARQUET_ENABLE_DICTIONARY] = "false"; + RunColumnWriteBenchmark( + state, MakeField("name", arrow::utf8(), 0), + [cardinality](int64_t rows, int64_t offset) { + return MakeStringColumn(rows, offset, cardinality); + }, + kRowsPerBatch, options); +} + +void BM_ParquetWrite_Decimal(::benchmark::State& state) { + const auto precision = static_cast(state.range(0)); + RunColumnWriteBenchmark(state, MakeField("amount", arrow::decimal128(precision, 4), 0), + [precision](int64_t rows, int64_t offset) { + return MakeDecimalColumn(rows, offset, precision, /*scale=*/4); + }, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_Struct(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("info", StructColumnType(), 0), &MakeStructColumn, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_List(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("tags", ListColumnType(), 0), &MakeListColumn, + kRowsPerBatch, /*options=*/{}); +} + +void BM_ParquetWrite_Vector(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("embedding", VectorColumnType(), 0), &MakeVectorColumn, + kRowsPerBatch, /*options=*/{}); +} + +// arg: entries per map row. +void BM_ParquetWrite_Map(::benchmark::State& state) { + const auto entries = static_cast(state.range(0)); + RunColumnWriteBenchmark( + state, MakeField("attrs", MapColumnType(), 0), + [entries](int64_t rows, int64_t offset) { return MakeMapColumn(rows, offset, entries); }, + kRowsPerBatch, /*options=*/{}); +} + +// arg: percentage of nulls. Every field is nullable, so the column is `optional` and definition +// levels are written at every density including zero; the sweep moves what arrow's null-free fast +// path is worth. BIGINT is the narrowest column, so levels are the largest share of what is left. +void BM_ParquetWrite_Nulls(::benchmark::State& state) { + const int64_t null_pct = state.range(0); + RunColumnWriteBenchmark(state, MakeField("id", arrow::int64(), 0), + NullableColumnFactory(&MakeInt64Column, null_pct), kRowsPerBatch, + /*options=*/{}); +} + +// arg: rows per AddBatch call, at a fixed total row count. Smaller batches mean the same data +// carries more per-batch fixed cost, and the file-level setup is amortized over more calls. +void BM_ParquetWrite_BatchSize(::benchmark::State& state) { + RunColumnWriteBenchmark(state, MakeField("id", arrow::int64(), 0), &MakeInt64Column, + state.range(0), /*options=*/{}); +} + +// arg: number of columns at a fixed row count. With BM_ParquetWrite_BatchSize this separates the +// two ways per-AddBatch work grows - more calls, or more columns per call - where that work is +// AddBatch importing the C array plus arrow walking the schema, neither of which encodes a value. +void BM_ParquetWrite_ColumnCount(::benchmark::State& state) { + const auto columns = static_cast(state.range(0)); + arrow::FieldVector fields; + fields.reserve(columns); + for (int32_t i = 0; i < columns; ++i) { + fields.push_back(MakeField("c" + std::to_string(i), arrow::utf8(), i)); + } + RunWriteBenchmark(state, arrow::schema(fields), &MakeWideBatch, kRowsPerBatch, /*options=*/{}, + kDefaultCompression); + state.counters["columns"] = ::benchmark::Counter(static_cast(columns)); +} + +// arg: maximum rows per row group, the write-side half of BM_ParquetRead_Filtered - a row group +// is the unit the reader prunes, so more of them buys finer pruning at the cost of footer +// metadata, flush work and codec material. The byte-triggered counterpart is +// BM_ParquetWrite_MemoryThreshold. +void BM_ParquetWrite_RowGroupSize(::benchmark::State& state) { + const int64_t row_group_length = state.range(0); + std::map options; + options[paimon::parquet::PARQUET_WRITE_MAX_ROW_GROUP_LENGTH] = std::to_string(row_group_length); + RunWriteBenchmark(state, FlatSchema(), &MakeFlatBatch, kRowsPerBatch, options, + kDefaultCompression); +} + +// args: the parquet.writer.max.memory.use threshold in KiB, and the number of AddBatch calls. +// Unlike every other case here the file is not a fixed kRowsPerFile: batches are a fixed 10'000 +// rows and the total grows with the batch count, because the question is whether retained size +// accumulates across batches until the writer cuts a row group. Holding the total fixed and +// shrinking the batch would measure AddBatch granularity instead, which is what +// BM_ParquetWrite_BatchSize already does. The input is low-cardinality dictionary data, where +// retained and flat size differ most, written as one batch object repeatedly rather than a fresh +// one each time so arrow sees the same dictionary rather than N equal ones. +// +// Nothing here sets a row-group row limit, so the byte threshold is the only thing that can +// trigger a flush. row_groups reports whether it fired at all; if it reads 1 everywhere the sweep +// measured nothing, and the unit test only covers the mechanism, not this particular setting. +void BM_ParquetWrite_MemoryThreshold(::benchmark::State& state) { + constexpr int64_t kFlushRowsPerBatch = 10'000; + std::map options; + options[paimon::parquet::PARQUET_WRITER_MAX_MEMORY_USE] = std::to_string(state.range(0) * 1024); + RunColumnWriteBenchmark( + state, MakeField("name", DictionaryStringType(), 0), + [](int64_t rows, int64_t offset) { + return MakeDictionaryStringColumn(rows, offset, kLowStringCardinality); + }, + kFlushRowsPerBatch, options, kFlushRowsPerBatch * state.range(1), BatchReuse::kReused); +} + +// Codec sweep over the mixed flat schema, so the CPU-versus-size trade shows on data that is not +// uniformly one type. Levels stay at the ParquetWriterBuilder defaults. The names are the ones +// Parquet accepts, not the ones arrow does: "lz4" resolves to arrow's LZ4_FRAME, which +// parquet::IsCodecSupported rejects, so the two framings Parquet defines are spelled out. +void BM_ParquetWrite_Compression(::benchmark::State& state, const char* compression) { + RunWriteBenchmark(state, FlatSchema(), &MakeFlatBatch, kRowsPerBatch, /*options=*/{}, + compression); +} + +// Rows a read case has to materialize for its number to mean anything. Pruning is not precise, so +// a filtered case gets a range rather than an exact count. Without the bound, a fixture that +// silently stopped producing rows would just look fast. +struct RowExpectation { + int64_t min = kRowsPerFile; + int64_t max = kRowsPerFile; +}; + +void RunReadBenchmark(::benchmark::State& state, const ReadFixture& fixture, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::optional& selection_bitmap, + const std::map& options, int32_t batch_size, + const RowExpectation& expected = {}) { + if (FailBenchmark(state, fixture.status())) { + return; + } + + ReadStats stats; + Timer timer; + for (auto _ : state) { + Result result = + ReadParquetFile(fixture.fs(), fixture.path(), read_schema, predicate, selection_bitmap, + options, batch_size, fixture.arrow_pool()); + if (FailBenchmark(state, result.status())) { + return; + } + stats = result.value(); + } + if (stats.rows < expected.min || stats.rows > expected.max) { + FailBenchmark(state, Status::Invalid("read " + std::to_string(stats.rows) + + " rows, expected " + std::to_string(expected.min) + + " to " + std::to_string(expected.max))); + return; + } + + const double elapsed_ns = timer.ElapsedNanos(); + ReportRowRate(state, stats.rows, elapsed_ns); + // Normalized by what the file holds, not what this case materialized: pruning drops the + // per-materialized-row numerator and denominator together, so ns_per_row and bytes_per_row can + // rise while the run gets faster. Only a denominator every setting shares is comparable across + // settings that prune by different amounts. + const double total_input_rows = + static_cast(state.iterations()) * static_cast(kRowsPerFile); + if (total_input_rows > 0) { + state.counters["ns_per_input_row"] = ::benchmark::Counter(elapsed_ns / total_input_rows); + state.counters["bytes_per_input_row"] = ::benchmark::Counter( + static_cast(stats.storage_bytes) / static_cast(kRowsPerFile)); + } + ReportBytes(state, "read_bytes", static_cast(stats.storage_bytes), stats.rows); + state.counters["rows_read"] = ::benchmark::Counter(static_cast(stats.rows)); + state.counters["file_bytes"] = ::benchmark::Counter(static_cast(fixture.file_bytes())); + state.counters["row_groups"] = + ::benchmark::Counter(static_cast(stats.row_groups_total)); + state.counters["row_groups_after_filter"] = + ::benchmark::Counter(static_cast(stats.row_groups_after_filter)); + state.counters["batches"] = ::benchmark::Counter(static_cast(stats.batches)); +} + +void BM_ParquetRead_FullScan(::benchmark::State& state) { + RunReadBenchmark(state, FlatFixture(), FlatSchema(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); +} + +// One column at a time, separating per-column transfer from materializing the whole row. +void BM_ParquetRead_Projection(::benchmark::State& state, const char* column) { + std::shared_ptr field = FlatSchema()->GetFieldByName(column); + if (!field) { + FailBenchmark(state, Status::Invalid("unknown projection column")); + return; + } + RunReadBenchmark(state, FlatFixture(), arrow::schema({field}), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); +} + +// args: percentage of rows the predicate keeps, and whether page-level pruning is on. The `id` +// column is ordered, so the surviving rows form a prefix; running both page-index settings makes +// the pruning gain attributable instead of merely visible. +void BM_ParquetRead_Filtered(::benchmark::State& state) { + const int64_t threshold = kRowsPerFile * state.range(0) / 100; + // `id` is ordered, so row-group statistics alone must already discard every group past the + // threshold - that holds with page-index filtering off too. Bounding at the row-group grain + // rather than at kRowsPerFile is what makes a pruning regression fail the case instead of + // quietly reading the whole file. + const int64_t row_group_bound = + ((threshold + kRowGroupLength - 1) / kRowGroupLength) * kRowGroupLength; + const bool enable_page_index = state.range(1) != 0; + std::shared_ptr predicate = PredicateBuilder::LessThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, Literal(threshold)); + std::map options; + options[paimon::parquet::PARQUET_READ_ENABLE_PAGE_INDEX_FILTER] = + enable_page_index ? "true" : "false"; + RunReadBenchmark(state, FlatFixture(), FlatSchema(), predicate, + /*selection_bitmap=*/std::nullopt, options, kReadBatchSize, + RowExpectation{threshold, row_group_bound}); +} + +// arg: distance between selected rows. The strides straddle the coalesce hole limit (32 rows by +// default), which splits this into two regimes: at or below it neighbouring single-row ranges +// merge into long spans, so rows_read runs far ahead of selected_rows and nothing is skipped; +// above it every row stays its own range and arrow's Skip decodes and discards each gap instead. +// Compare read_bytes and ns_per_input_row across the two. +void BM_ParquetRead_SkipHeavy(::benchmark::State& state) { + const int64_t stride = state.range(0); + RoaringBitmap32 bitmap; + for (int64_t row = 0; row < kRowsPerFile; row += stride) { + bitmap.Add(static_cast(row)); + } + const int64_t selected = bitmap.Cardinality(); + state.counters["selected_rows"] = ::benchmark::Counter(static_cast(selected)); + RunReadBenchmark(state, FlatFixture(), FlatSchema(), /*predicate=*/nullptr, bitmap, + /*options=*/{}, kReadBatchSize, RowExpectation{selected, kRowsPerFile}); +} + +// arg: percentage of nulls, the read side of BM_ParquetWrite_Nulls: definition levels have to be +// decoded back into a validity bitmap, and past some density arrow may do less work, not more. +void BM_ParquetRead_Nulls(::benchmark::State& state) { + RunReadBenchmark(state, NullableFlatFixture(state.range(0)), FlatSchema(), + /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, /*options=*/{}, + kReadBatchSize); +} + +// The VARCHAR column from a dictionary-encoded file and from a plain-encoded one. Both are +// decoded by arrow, so this is arrow's dictionary path against its plain path on equal values. +void BM_ParquetRead_Encoding(::benchmark::State& state, bool enable_dictionary) { + std::shared_ptr field = FlatSchema()->GetFieldByName("name"); + RunReadBenchmark(state, enable_dictionary ? FlatFixture() : PlainFlatFixture(), + arrow::schema({field}), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); +} + +// arg: decimal precision, the read side of BM_ParquetWrite_Decimal. Precision picks the physical +// type - INT32, INT64 or FIXED_LEN_BYTE_ARRAY, since ParquetWriterBuilder enables +// store_decimal_as_integer - and the three take different paths back to Decimal128Array. +void BM_ParquetRead_Decimal(::benchmark::State& state) { + const auto precision = static_cast(state.range(0)); + RunReadBenchmark(state, DecimalFixture(precision), DecimalSchema(precision), + /*predicate=*/nullptr, /*selection_bitmap=*/std::nullopt, /*options=*/{}, + kReadBatchSize); +} + +void BM_ParquetRead_Double(::benchmark::State& state) { + RunReadBenchmark(state, DoubleFixture(), DoubleSchema(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); +} + +// arg: rows per NextBatch call. At a fixed row count this turns the per-batch fixed cost of +// ParquetFileBatchReader::NextBatch - Validate, the two metrics counters, the ArrowSchema +// export - into a number, separated from the per-row decoding cost. +void BM_ParquetRead_BatchSize(::benchmark::State& state) { + RunReadBenchmark(state, FlatFixture(), FlatSchema(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, + static_cast(state.range(0))); +} + +// Each nested column costs definition and repetition levels a flat column does not pay. +void BM_ParquetRead_Nested(::benchmark::State& state, const char* column) { + std::shared_ptr full_schema = NestedReadSchema(); + std::shared_ptr read_schema = full_schema; + if (column != nullptr) { + std::shared_ptr field = full_schema->GetFieldByName(column); + if (!field) { + FailBenchmark(state, Status::Invalid("unknown nested column")); + return; + } + read_schema = arrow::schema({field}); + } + RunReadBenchmark(state, NestedFixture(), read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); +} + +} // namespace + +BENCHMARK(BM_ParquetWrite_Int64)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Double)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Boolean)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_String) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(1000) + ->Arg(kRowsPerFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_StringNoDictionary) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(1000) + ->Arg(kRowsPerFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_FlatInt32) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_DictionaryString) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_DictionaryInt32) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Decimal) + ->ArgName("precision") + ->Arg(9) + ->Arg(18) + ->Arg(38) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Struct)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_List)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Vector)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Map) + ->ArgName("entries") + ->Arg(3) + ->Arg(5) + ->Arg(10) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_Nulls) + ->ArgName("null_pct") + ->Arg(0) + ->Arg(20) + ->Arg(50) + ->Arg(70) + ->Arg(100) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_BatchSize) + ->ArgName("rows_per_batch") + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_ColumnCount) + ->ArgName("columns") + ->Arg(1) + ->Arg(5) + ->Arg(10) + ->Arg(20) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_RowGroupSize) + ->ArgName("row_group_rows") + ->Arg(5000) + ->Arg(25000) + ->Arg(kRowsPerFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetWrite_MemoryThreshold) + ->ArgNames({"max_memory_kib", "batches"}) + ->Args({512, 50}) + ->Args({512, 200}) + ->Args({4096, 200}) + ->Args({65536, 200}) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, none, "none") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, snappy, "snappy") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, gzip, "gzip") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, lz4_raw, "lz4_raw") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, lz4_hadoop, "lz4_hadoop") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, brotli, "brotli") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetWrite_Compression, zstd, "zstd") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); + +BENCHMARK(BM_ParquetRead_FullScan)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Projection, id, "id") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Projection, name, "name") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Projection, amount, "amount") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetRead_Filtered) + ->ArgNames({"keep_pct", "page_index"}) + ->Args({1, 1}) + ->Args({1, 0}) + ->Args({10, 1}) + ->Args({10, 0}) + ->Args({50, 1}) + ->Args({50, 0}) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetRead_SkipHeavy) + ->ArgName("stride") + ->Arg(8) + ->Arg(32) + ->Arg(64) + ->Arg(512) + ->Arg(4096) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetRead_Nulls) + ->ArgName("null_pct") + ->Arg(0) + ->Arg(20) + ->Arg(50) + ->Arg(70) + ->Arg(100) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Encoding, dictionary, true) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Encoding, plain, false) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetRead_Decimal) + ->ArgName("precision") + ->Arg(9) + ->Arg(18) + ->Arg(38) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK(BM_ParquetRead_Double)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_ParquetRead_BatchSize) + ->ArgName("batch_size") + ->Arg(512) + ->Arg(4096) + ->Arg(16384) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Nested, info_struct, "info") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Nested, tags_list, "tags") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Nested, embedding_vector, "embedding") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Nested, attrs_map, "attrs") + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +BENCHMARK_CAPTURE(BM_ParquetRead_Nested, all_columns, nullptr) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); + +int main(int argc, char** argv) { + ::benchmark::Initialize(&argc, argv); + if (::benchmark::ReportUnrecognizedArguments(argc, argv)) { + return 1; + } + ::benchmark::RunSpecifiedBenchmarks(); + ::benchmark::Shutdown(); + return g_failed.load() ? 1 : 0; +} diff --git a/benchmark/parquet_format_benchmark_test.cpp b/benchmark/parquet_format_benchmark_test.cpp new file mode 100644 index 000000000..e930449d0 --- /dev/null +++ b/benchmark/parquet_format_benchmark_test.cpp @@ -0,0 +1,533 @@ +/* + * 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. + */ + +// Smoke test for the assumptions parquet_format_benchmark.cpp is built on. +// +// The benchmark is only compiled under PAIMON_BUILD_BENCHMARKS, which CI does not set, so nothing +// there runs in CI. Every assumption it makes about the format layer - that a codec name is +// accepted, that a dictionary-encoded input array can be written, that a nested or high-precision +// column survives a round trip, that a predicate and a selection bitmap return the rows the +// benchmark asserts on, that the reader metrics it reports exist - is checked here instead, at a +// row count small enough to stay a test. A regression in any of them would otherwise surface as a +// benchmark that quietly measures the wrong thing. + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/concatenate.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/defs.h" +#include "paimon/format/format_writer.h" +#include "paimon/format/parquet/parquet_field_id_converter.h" +#include "paimon/format/parquet/parquet_file_batch_reader.h" +#include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/parquet/parquet_writer_builder.h" +#include "paimon/fs/file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/metrics.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/result.h" +#include "paimon/status.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" + +namespace paimon::parquet { +namespace { + +// Small enough to stay a test, large enough to span several batches and row groups. +constexpr int64_t kRows = 2'000; +constexpr int32_t kBatchSize = 256; +constexpr int64_t kRowGroupLength = 500; + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id) { + return arrow::field(name, type, + arrow::KeyValueMetadata::Make({ParquetFieldIdConverter::PARQUET_FIELD_ID}, + {std::to_string(field_id)})); +} + +class ParquetFormatBenchmarkTest : public ::testing::Test { + protected: + void SetUp() override { + dir_ = test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); + fs_ = dir_->GetFileSystem(); + pool_ = GetArrowPool(GetDefaultPool()); + } + + std::string PathOf(const std::string& name) const { + return PathUtil::JoinPath(dir_->Str(), name); + } + + // Writes the same struct array `batch_count` times through the builder the benchmark uses. + // More than one call matters for dictionary input: the benchmark always writes several + // batches, and every batch after the first hands the writer the same dictionary again. + Status Write(const std::string& path, const std::shared_ptr& schema, + const std::shared_ptr& batch, const std::string& compression, + const std::map& extra_options = {}, + int32_t batch_count = 1) { + std::map options = extra_options; + // emplace, not assignment: a caller that set its own row-group limit is testing that. + options.emplace(PARQUET_WRITE_MAX_ROW_GROUP_LENGTH, std::to_string(kRowGroupLength)); + ParquetWriterBuilder writer_builder(schema, kBatchSize, options); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr out, + fs_->Create(path, /*overwrite=*/true)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + writer_builder.Build(out, compression)); + for (int32_t i = 0; i < batch_count; ++i) { + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, &c_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array)); + } + PAIMON_RETURN_NOT_OK(writer->Finish()); + return out->Close(); + } + + // Concatenating `array` with itself `times` times, so a multi-batch write has an expected + // value to be compared against. + static Result> Repeat(const std::shared_ptr& array, + int32_t times) { + std::vector> chunks(times, array); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr repeated, + arrow::Concatenate(chunks)); + return repeated; + } + + static Result> MakeDictionary( + const std::shared_ptr& indices, const std::shared_ptr& values) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::DictionaryArray::FromArrays(indices, values)); + return array; + } + + struct ReadResult { + int64_t rows = 0; + uint64_t row_groups_total = 0; + uint64_t row_groups_after_filter = 0; + uint64_t batches = 0; + // Every batch, imported and concatenated. Row counts alone would let a decoding bug + // through, so the tests compare this against what was written. + std::shared_ptr data; + }; + + // Reads the file back the way the benchmark does, reporting what the benchmark reports on. + Result Read(const std::string& path, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate = nullptr, + const std::optional& selection = std::nullopt) { + PAIMON_ASSIGN_OR_RAISE(FileStatus file_status, fs_->GetFileStatus(path)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, fs_->Open(path)); + auto in_stream = + std::make_shared(input, file_status.GetLen(), pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + ParquetFileBatchReader::Create( + std::move(in_stream), /*options=*/{}, kBatchSize, + /*file_metadata=*/nullptr, /*storage_read_bytes=*/nullptr, pool_, + /*hints=*/std::nullopt)); + ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_schema)); + PAIMON_RETURN_NOT_OK(reader->SetReadSchema(&c_schema, predicate, selection)); + + ReadResult result; + std::vector> chunks; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + // ImportArray takes ownership of both C structs, so nothing is released by hand here. + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr chunk, + arrow::ImportArray(batch.first.get(), batch.second.get())); + result.rows += chunk->length(); + chunks.push_back(std::move(chunk)); + } + if (!chunks.empty()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(result.data, arrow::Concatenate(chunks)); + } + + std::shared_ptr metrics = reader->GetReaderMetrics(); + PAIMON_ASSIGN_OR_RAISE(result.row_groups_total, + metrics->GetCounter(ParquetMetrics::READ_ROW_GROUPS_TOTAL)); + PAIMON_ASSIGN_OR_RAISE(result.row_groups_after_filter, + metrics->GetCounter(ParquetMetrics::READ_ROW_GROUPS_AFTER_FILTER)); + PAIMON_ASSIGN_OR_RAISE(result.batches, + metrics->GetCounter(ParquetMetrics::READ_BATCH_COUNT)); + reader->Close(); + return result; + } + + static Result> Wrap( + const std::shared_ptr& schema, + const std::vector>& columns) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::StructArray::Make(columns, schema->fields())); + return checked_pointer_cast(array); + } + + std::unique_ptr dir_; + std::shared_ptr fs_; + std::shared_ptr pool_; +}; + +// Every codec name the benchmark registers has to be one Parquet accepts. "lz4" is the trap this +// guards: it resolves to arrow's LZ4_FRAME, which parquet::IsCodecSupported rejects, and the +// failure only shows up once a column chunk is actually written. +TEST_F(ParquetFormatBenchmarkTest, RegisteredCodecsWrite) { + std::shared_ptr schema = arrow::schema({MakeField("id", arrow::int64(), 0)}); + arrow::Int64Builder builder; + ASSERT_TRUE(builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + builder.UnsafeAppend(i); + } + std::shared_ptr ids; + ASSERT_TRUE(builder.Finish(&ids).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(schema, {ids})); + + for (const std::string& codec : + {"none", "snappy", "gzip", "brotli", "zstd", "lz4_raw", "lz4_hadoop"}) { + const std::string path = PathOf("codec_" + codec + ".parquet"); + ASSERT_OK(Write(path, schema, batch, codec)) << "codec " << codec; + ASSERT_OK_AND_ASSIGN(ReadResult result, Read(path, schema)); + EXPECT_EQ(kRows, result.rows) << "codec " << codec; + ASSERT_TRUE(result.data); + EXPECT_TRUE(result.data->Equals(*batch)) << "codec " << codec; + } + + // The name the benchmark deliberately does not register still has to be rejected; if arrow + // ever starts accepting it, the comment explaining its absence is stale. + EXPECT_FALSE(Write(PathOf("codec_lz4.parquet"), schema, batch, "lz4").ok()); +} + +// A dictionary-encoded input array must reach the writer intact. VARCHAR takes arrow's direct +// write path and INT32 gets densified first; both have to produce a readable file whose logical +// values match the flat equivalent. +TEST_F(ParquetFormatBenchmarkTest, DictionaryInputRoundTrip) { + constexpr int64_t kCardinality = 8; + arrow::Int32Builder index_builder; + ASSERT_TRUE(index_builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + index_builder.UnsafeAppend(static_cast(i % kCardinality)); + } + std::shared_ptr indices; + ASSERT_TRUE(index_builder.Finish(&indices).ok()); + + arrow::StringBuilder string_values; + arrow::Int32Builder int_values; + for (int64_t i = 0; i < kCardinality; ++i) { + ASSERT_TRUE(string_values.Append("value_" + std::to_string(i)).ok()); + ASSERT_TRUE(int_values.Append(static_cast(i * 7)).ok()); + } + std::shared_ptr string_dict; + std::shared_ptr int_dict; + ASSERT_TRUE(string_values.Finish(&string_dict).ok()); + ASSERT_TRUE(int_values.Finish(&int_dict).ok()); + + // The flat arrays the dictionary-encoded input has to decode back to. + arrow::StringBuilder flat_strings; + arrow::Int32Builder flat_ints; + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(flat_strings.Append("value_" + std::to_string(i % kCardinality)).ok()); + ASSERT_TRUE(flat_ints.Append(static_cast((i % kCardinality) * 7)).ok()); + } + std::shared_ptr flat_string_column; + std::shared_ptr flat_int_column; + ASSERT_TRUE(flat_strings.Finish(&flat_string_column).ok()); + ASSERT_TRUE(flat_ints.Finish(&flat_int_column).ok()); + + struct Case { + const char* name; + std::shared_ptr dictionary; + std::shared_ptr read_type; + std::shared_ptr flat; + }; + // The benchmark always writes several batches, so the writer sees the same dictionary more + // than once. + constexpr int32_t kBatches = 3; + for (const Case& c : {Case{"string", string_dict, arrow::utf8(), flat_string_column}, + Case{"int32", int_dict, arrow::int32(), flat_int_column}}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr dictionary_array, + MakeDictionary(indices, c.dictionary)); + std::shared_ptr write_schema = + arrow::schema({MakeField("v", dictionary_array->type(), 0)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, + Wrap(write_schema, {dictionary_array})); + const std::string path = PathOf(std::string("dict_") + c.name + ".parquet"); + ASSERT_OK(Write(path, write_schema, batch, "zstd", /*extra_options=*/{}, kBatches)) + << c.name; + + // Parquet has no dictionary type: the column comes back as its value type either way, and + // has to carry the values the flat equivalent would have. + std::shared_ptr read_schema = + arrow::schema({MakeField("v", c.read_type, 0)}); + ASSERT_OK_AND_ASSIGN(ReadResult result, Read(path, read_schema)); + EXPECT_EQ(kRows * kBatches, result.rows) << c.name; + ASSERT_TRUE(result.data); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, Repeat(c.flat, kBatches)); + std::shared_ptr actual = + checked_pointer_cast(result.data)->field(0); + EXPECT_TRUE(actual->Equals(*expected)) << c.name; + } +} + +// DECIMAL precision selects the Parquet physical type, and precision 38 is the only one that +// reaches FIXED_LEN_BYTE_ARRAY. The benchmark sweeps all three on both sides, so all three have +// to survive a round trip with their type intact. +TEST_F(ParquetFormatBenchmarkTest, DecimalPrecisionRoundTrip) { + for (int32_t precision : {9, 18, 38}) { + std::shared_ptr type = arrow::decimal128(precision, 4); + std::shared_ptr schema = arrow::schema({MakeField("amount", type, 0)}); + arrow::Decimal128Builder builder(type); + ASSERT_TRUE(builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(builder.Append(arrow::Decimal128(i)).ok()); + } + std::shared_ptr amounts; + ASSERT_TRUE(builder.Finish(&amounts).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(schema, {amounts})); + + const std::string path = PathOf("decimal_" + std::to_string(precision) + ".parquet"); + ASSERT_OK(Write(path, schema, batch, "zstd")) << "precision " << precision; + ASSERT_OK_AND_ASSIGN(ReadResult result, Read(path, schema)); + EXPECT_EQ(kRows, result.rows) << "precision " << precision; + ASSERT_TRUE(result.data); + EXPECT_TRUE(result.data->type()->field(0)->type()->Equals(*type)) + << "precision " << precision; + EXPECT_TRUE(result.data->Equals(*batch)) << "precision " << precision; + } +} + +// The nested fixture the benchmark reads is only meaningful if LIST and MAP survive the round +// trip with the shape the read schema asks for. +TEST_F(ParquetFormatBenchmarkTest, NestedRoundTrip) { + constexpr int32_t kEntries = 4; + auto list_values = std::make_shared(); + arrow::ListBuilder list_builder(arrow::default_memory_pool(), list_values, + arrow::list(arrow::int64())); + auto key_builder = std::make_shared(); + auto item_builder = std::make_shared(); + arrow::MapBuilder map_builder(arrow::default_memory_pool(), key_builder, item_builder, + arrow::map(arrow::utf8(), arrow::int64())); + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(list_builder.Append().ok()); + ASSERT_TRUE(map_builder.Append().ok()); + for (int32_t j = 0; j < kEntries; ++j) { + ASSERT_TRUE(list_values->Append(i + j).ok()); + ASSERT_TRUE(key_builder->Append("key_" + std::to_string(j)).ok()); + ASSERT_TRUE(item_builder->Append(i + j).ok()); + } + } + std::shared_ptr tags; + std::shared_ptr attrs; + ASSERT_TRUE(list_builder.Finish(&tags).ok()); + ASSERT_TRUE(map_builder.Finish(&attrs).ok()); + + std::shared_ptr schema = + arrow::schema({MakeField("tags", arrow::list(arrow::int64()), 0), + MakeField("attrs", arrow::map(arrow::utf8(), arrow::int64()), 1)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(schema, {tags, attrs})); + const std::string path = PathOf("nested.parquet"); + ASSERT_OK(Write(path, schema, batch, "zstd")); + + ASSERT_OK_AND_ASSIGN(ReadResult result, Read(path, schema)); + EXPECT_EQ(kRows, result.rows); + ASSERT_TRUE(result.data); + EXPECT_TRUE(result.data->Equals(*batch)); + + // Each nested column also has to be readable on its own, which is what the projected nested + // cases do. + for (const std::string& column : {"tags", "attrs"}) { + std::shared_ptr projected = arrow::schema({schema->GetFieldByName(column)}); + ASSERT_OK_AND_ASSIGN(ReadResult projected_result, Read(path, projected)); + EXPECT_EQ(kRows, projected_result.rows) << "column " << column; + } +} + +// The filtered and skip-heavy cases assert on row counts, so those counts have to mean what the +// benchmark assumes: a predicate keeps at least the matching rows and no more than the row groups +// that could hold them, and a selection bitmap keeps at least the rows it selected. +TEST_F(ParquetFormatBenchmarkTest, FilteredAndBitmapRowCounts) { + std::shared_ptr schema = arrow::schema({MakeField("id", arrow::int64(), 0)}); + arrow::Int64Builder builder; + ASSERT_TRUE(builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + builder.UnsafeAppend(i); + } + std::shared_ptr ids; + ASSERT_TRUE(builder.Finish(&ids).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(schema, {ids})); + const std::string path = PathOf("filtered.parquet"); + ASSERT_OK(Write(path, schema, batch, "zstd")); + + ASSERT_OK_AND_ASSIGN(ReadResult full, Read(path, schema)); + EXPECT_EQ(kRows, full.rows); + EXPECT_EQ(static_cast(kRows / kRowGroupLength), full.row_groups_total); + EXPECT_EQ(full.row_groups_total, full.row_groups_after_filter); + EXPECT_GT(full.batches, 0u); + + // `id` is ordered, so row-group statistics alone must discard everything past the threshold. + // This is the bound BM_ParquetRead_Filtered asserts on. + constexpr int64_t kThreshold = 300; + std::shared_ptr predicate = PredicateBuilder::LessThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, Literal(kThreshold)); + ASSERT_OK_AND_ASSIGN(ReadResult filtered, Read(path, schema, predicate)); + EXPECT_GE(filtered.rows, kThreshold); + EXPECT_LE(filtered.rows, kRowGroupLength); + EXPECT_LT(filtered.row_groups_after_filter, filtered.row_groups_total); + + // Selection is not precise, so the bitmap is a lower bound on what comes back. + RoaringBitmap32 bitmap; + int64_t selected = 0; + for (int64_t row = 0; row < kRows; row += 64) { + bitmap.Add(static_cast(row)); + ++selected; + } + ASSERT_OK_AND_ASSIGN(ReadResult skipped, Read(path, schema, /*predicate=*/nullptr, bitmap)); + EXPECT_GE(skipped.rows, selected); + EXPECT_LE(skipped.rows, kRows); +} + +// The encoding case compares a dictionary-encoded file against a plain one, which is only a +// comparison if both files read back identically. +TEST_F(ParquetFormatBenchmarkTest, PlainAndDictionaryFilesAgree) { + std::shared_ptr schema = arrow::schema({MakeField("name", arrow::utf8(), 0)}); + arrow::StringBuilder builder; + ASSERT_TRUE(builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(builder.Append("value_" + std::to_string(i % 16)).ok()); + } + std::shared_ptr names; + ASSERT_TRUE(builder.Finish(&names).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(schema, {names})); + + const std::string dict_path = PathOf("encoding_dict.parquet"); + const std::string plain_path = PathOf("encoding_plain.parquet"); + ASSERT_OK(Write(dict_path, schema, batch, "zstd")); + std::map plain_options; + plain_options[PARQUET_ENABLE_DICTIONARY] = "false"; + ASSERT_OK(Write(plain_path, schema, batch, "zstd", plain_options)); + + ASSERT_OK_AND_ASSIGN(ReadResult dict_result, Read(dict_path, schema)); + ASSERT_OK_AND_ASSIGN(ReadResult plain_result, Read(plain_path, schema)); + EXPECT_EQ(kRows, dict_result.rows); + EXPECT_EQ(dict_result.rows, plain_result.rows); + EXPECT_EQ(dict_result.row_groups_total, plain_result.row_groups_total); + ASSERT_TRUE(dict_result.data); + ASSERT_TRUE(plain_result.data); + EXPECT_TRUE(dict_result.data->Equals(*batch)); + EXPECT_TRUE(plain_result.data->Equals(*batch)); +} + +// BM_ParquetWrite_MemoryThreshold rests on one assumption: that a small +// parquet.writer.max.memory.use actually makes ParquetFormatWriter cut extra row groups. This +// covers the mechanism on both input shapes that case can present it with, plain and +// dictionary-encoded, with the row-count limit raised out of the way so the byte threshold is the +// only thing that can flush. +// +// It does not cover the benchmark's own 512 KiB setting: reaching that with cardinality-10 +// dictionary data takes the 500K to 2M rows the benchmark writes, which is not a unit test. What +// confirms that setting is the benchmark's own row_groups counter reading more than 1. +TEST_F(ParquetFormatBenchmarkTest, MemoryThresholdFlushesRowGroups) { + constexpr int32_t kBatches = 8; + constexpr int64_t kDictCardinality = 10; + + arrow::StringBuilder plain_builder; + ASSERT_TRUE(plain_builder.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(plain_builder.Append("value_" + std::to_string(i)).ok()); + } + std::shared_ptr plain_column; + ASSERT_TRUE(plain_builder.Finish(&plain_column).ok()); + + // The shape BM_ParquetWrite_MemoryThreshold writes: low-cardinality dictionary input. + arrow::StringBuilder dict_values; + for (int64_t i = 0; i < kDictCardinality; ++i) { + ASSERT_TRUE(dict_values.Append("value_" + std::to_string(i)).ok()); + } + std::shared_ptr dict_value_column; + ASSERT_TRUE(dict_values.Finish(&dict_value_column).ok()); + arrow::Int32Builder dict_indices; + ASSERT_TRUE(dict_indices.Reserve(kRows).ok()); + for (int64_t i = 0; i < kRows; ++i) { + dict_indices.UnsafeAppend(static_cast(i % kDictCardinality)); + } + std::shared_ptr index_column; + ASSERT_TRUE(dict_indices.Finish(&index_column).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr dict_column, + MakeDictionary(index_column, dict_value_column)); + + struct Case { + const char* name; + std::shared_ptr column; + }; + for (const Case& c : {Case{"plain", plain_column}, Case{"dictionary", dict_column}}) { + std::shared_ptr write_schema = + arrow::schema({MakeField("name", c.column->type(), 0)}); + // Parquet stores a dictionary column as its value type, so both read back as UTF8. + std::shared_ptr read_schema = + arrow::schema({MakeField("name", arrow::utf8(), 0)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, Wrap(write_schema, {c.column})); + + std::map options; + options[PARQUET_WRITE_MAX_ROW_GROUP_LENGTH] = std::to_string(kRows * kBatches * 10); + + options[PARQUET_WRITER_MAX_MEMORY_USE] = std::to_string(uint64_t{8} * 1024); + const std::string small_path = + PathOf(std::string("threshold_small_") + c.name + ".parquet"); + ASSERT_OK(Write(small_path, write_schema, batch, "zstd", options, kBatches)) << c.name; + ASSERT_OK_AND_ASSIGN(ReadResult small, Read(small_path, read_schema)); + EXPECT_EQ(kRows * kBatches, small.rows) << c.name; + EXPECT_GT(small.row_groups_total, 1u) + << c.name << ": byte threshold never triggered a flush"; + + options[PARQUET_WRITER_MAX_MEMORY_USE] = + std::to_string(DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE); + const std::string large_path = + PathOf(std::string("threshold_large_") + c.name + ".parquet"); + ASSERT_OK(Write(large_path, write_schema, batch, "zstd", options, kBatches)) << c.name; + ASSERT_OK_AND_ASSIGN(ReadResult large, Read(large_path, read_schema)); + EXPECT_EQ(kRows * kBatches, large.rows) << c.name; + EXPECT_EQ(1u, large.row_groups_total) << c.name; + + // Row-group boundaries must not change the data. + ASSERT_TRUE(small.data) << c.name; + ASSERT_TRUE(large.data) << c.name; + EXPECT_TRUE(small.data->Equals(*large.data)) << c.name; + } +} + +} // namespace +} // namespace paimon::parquet diff --git a/ci/scripts/build_paimon.sh b/ci/scripts/build_paimon.sh index 50983d4da..cea59da8b 100755 --- a/ci/scripts/build_paimon.sh +++ b/ci/scripts/build_paimon.sh @@ -147,7 +147,9 @@ CMAKE_ARGS=( "-G Ninja" "-DCMAKE_BUILD_TYPE=${build_type}" "-DPAIMON_BUILD_TESTS=ON" + "-DPAIMON_ENABLE_MOSAIC=ON" "-DPAIMON_ENABLE_JINDO=ON" + "-DPAIMON_ENABLE_OSS=ON" "-DPAIMON_ENABLE_S3=ON" "-DPAIMON_ENABLE_LUMINA=${ENABLE_LUMINA}" "-DPAIMON_ENABLE_LUCENE=ON" diff --git a/ci/scripts/setup_rust.sh b/ci/scripts/setup_rust.sh index bdb8f622e..8c1e58392 100755 --- a/ci/scripts/setup_rust.sh +++ b/ci/scripts/setup_rust.sh @@ -15,12 +15,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# Install the Rust toolchain + cbindgen required to build the -# tantivy-fts FFI crate (crates/tantivy_ffi) from CI. +# Install the Rust toolchain used by the Mosaic and tantivy-fts FFI builds, plus cbindgen required +# by tantivy-fts. # # The dev container (see .devcontainer/) already has these preinstalled; -# this script is for the GitHub Actions runners. Called by -# .github/workflows/build_and_test.yaml before ci/scripts/build_paimon.sh. +# this script is for the GitHub Actions runners and is called before ci/scripts/build_paimon.sh. # # Idempotent: a second invocation is a no-op when the tools already exist. diff --git a/cmake_modules/DefineOptions.cmake b/cmake_modules/DefineOptions.cmake index 21f653155..9f099208f 100644 --- a/cmake_modules/DefineOptions.cmake +++ b/cmake_modules/DefineOptions.cmake @@ -219,6 +219,11 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}") "" AUTO BUNDLED) + define_option_string(OSS_SDK_V2_SOURCE + "Dependency source for OSS SDK v2; SYSTEM is unsupported" + "" + AUTO + BUNDLED) define_option_string(fmt_SOURCE "Dependency source for fmt" "" diff --git a/cmake_modules/ThirdpartyToolchain.cmake b/cmake_modules/ThirdpartyToolchain.cmake index 776519104..0c64819c2 100644 --- a/cmake_modules/ThirdpartyToolchain.cmake +++ b/cmake_modules/ThirdpartyToolchain.cmake @@ -91,6 +91,18 @@ else() endif() endif() +if(DEFINED ENV{PAIMON_MOSAIC_URL}) + set(MOSAIC_SOURCE_URL "$ENV{PAIMON_MOSAIC_URL}") +else() + if(EXISTS "${THIRDPARTY_DIR}/${PAIMON_MOSAIC_PKG_NAME}") + set_urls(MOSAIC_SOURCE_URL "${THIRDPARTY_DIR}/${PAIMON_MOSAIC_PKG_NAME}") + else() + set_urls(MOSAIC_SOURCE_URL + "https://downloads.apache.org/paimon/paimon-mosaic-${PAIMON_MOSAIC_BUILD_VERSION}/${PAIMON_MOSAIC_PKG_NAME}" + ) + endif() +endif() + if(DEFINED ENV{PAIMON_RAPIDJSON_URL}) set(RAPIDJSON_SOURCE_URL "$ENV{PAIMON_RAPIDJSON_URL}") else() @@ -308,6 +320,16 @@ else() endif() endif() +if(DEFINED ENV{PAIMON_OSS_SDK_V2_URL}) + set(OSS_SDK_V2_SOURCE_URL "$ENV{PAIMON_OSS_SDK_V2_URL}") +elseif(EXISTS "${THIRDPARTY_DIR}/${PAIMON_OSS_SDK_V2_PKG_NAME}") + set_urls(OSS_SDK_V2_SOURCE_URL "${THIRDPARTY_DIR}/${PAIMON_OSS_SDK_V2_PKG_NAME}") +else() + set_urls(OSS_SDK_V2_SOURCE_URL + "${THIRDPARTY_MIRROR_URL}https://github.com/aliyun/alibabacloud-oss-cpp-sdk-v2/archive/refs/tags/${PAIMON_OSS_SDK_V2_BUILD_VERSION}.tar.gz" + ) +endif() + if(DEFINED ENV{PAIMON_LUMINA_URL}) set(LUMINA_SOURCE_URL "$ENV{PAIMON_LUMINA_URL}") elseif(EXISTS "${THIRDPARTY_DIR}/${PAIMON_LUMINA_PKG_NAME}") @@ -492,6 +514,25 @@ function(paimon_enforce_patched_dependency_policy) PARENT_SCOPE) endif() endif() + + if(PAIMON_ENABLE_OSS) + paimon_set_dependency_source_default( + OSS_SDK_V2 BUNDLED "OSS SDK v2 is only supported as a bundled dependency") + paimon_get_dependency_source(OSS_SDK_V2 _oss_sdk_v2_source) + if(_oss_sdk_v2_source STREQUAL "SYSTEM") + message(FATAL_ERROR "OSS_SDK_V2_SOURCE=SYSTEM is not supported. " + "Use OSS_SDK_V2_SOURCE=BUNDLED.") + elseif(_oss_sdk_v2_source STREQUAL "AUTO") + message(STATUS "Forcing OSS_SDK_V2_SOURCE to BUNDLED because paimon-cpp " + "only supports the bundled OSS SDK v2") + set(OSS_SDK_V2_SOURCE + "BUNDLED" + CACHE STRING "Dependency source for OSS SDK v2" FORCE) + set(OSS_SDK_V2_SOURCE + "BUNDLED" + PARENT_SCOPE) + endif() + endif() endfunction() function(paimon_apply_dependency_source_defaults) @@ -597,6 +638,8 @@ function(paimon_get_dependency_compat_target DEPENDENCY_NAME OUT_VAR) set(_target tbb) elseif("${DEPENDENCY_NAME}" STREQUAL "Avro") set(_target avro) + elseif("${DEPENDENCY_NAME}" STREQUAL "OSS_SDK_V2") + set(_target alibabacloud_oss_v2::oss) else() set(_target "${DEPENDENCY_NAME}") endif() @@ -671,6 +714,8 @@ macro(paimon_build_dependency DEPENDENCY_NAME) build_glog() elseif("${DEPENDENCY_NAME}" STREQUAL "Avro") build_avro() + elseif("${DEPENDENCY_NAME}" STREQUAL "OSS_SDK_V2") + build_oss_sdk_v2() elseif("${DEPENDENCY_NAME}" STREQUAL "GTest") build_gtest() elseif("${DEPENDENCY_NAME}" STREQUAL "benchmark") @@ -1314,6 +1359,50 @@ macro(build_lumina) install(FILES "${LUMINA_DYNAMIC_LIB}" DESTINATION ${CMAKE_INSTALL_LIBDIR}) endmacro() +macro(build_mosaic) + message(STATUS "Building Apache Paimon Mosaic Rust FFI from source") + find_program(PAIMON_CARGO_EXECUTABLE cargo REQUIRED) + + set(MOSAIC_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/mosaic_ep-install") + set(MOSAIC_INCLUDE_DIR "${MOSAIC_PREFIX}/include") + set(MOSAIC_LIB_DIR "${MOSAIC_PREFIX}/${CMAKE_INSTALL_LIBDIR}") + set(MOSAIC_DYNAMIC_LIB + "${MOSAIC_LIB_DIR}/${CMAKE_SHARED_LIBRARY_PREFIX}paimon_mosaic_ffi${CMAKE_SHARED_LIBRARY_SUFFIX}" + ) + set(MOSAIC_CARGO_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/mosaic_ep-cargo") + set(MOSAIC_CARGO_DYNAMIC_LIB + "${MOSAIC_CARGO_TARGET_DIR}/release/${CMAKE_SHARED_LIBRARY_PREFIX}paimon_mosaic_ffi${CMAKE_SHARED_LIBRARY_SUFFIX}" + ) + + file(MAKE_DIRECTORY "${MOSAIC_INCLUDE_DIR}") + file(MAKE_DIRECTORY "${MOSAIC_LIB_DIR}") + + externalproject_add(mosaic_ep + URL ${MOSAIC_SOURCE_URL} + URL_HASH "SHA256=${PAIMON_MOSAIC_BUILD_SHA256_CHECKSUM}" + ${THIRDPARTY_LOG_OPTIONS} + CONFIGURE_COMMAND "" + BUILD_COMMAND ${CMAKE_COMMAND} -E env + CARGO_TARGET_DIR=${MOSAIC_CARGO_TARGET_DIR} + ${PAIMON_CARGO_EXECUTABLE} build --release + --manifest-path /ffi/Cargo.toml + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${MOSAIC_CARGO_DYNAMIC_LIB} ${MOSAIC_DYNAMIC_LIB} + INSTALL_COMMAND ${CMAKE_COMMAND} -E copy_directory + /include ${MOSAIC_INCLUDE_DIR} + BUILD_BYPRODUCTS "${MOSAIC_DYNAMIC_LIB}") + + add_library(paimon_mosaic_ffi SHARED IMPORTED GLOBAL) + set_target_properties(paimon_mosaic_ffi + PROPERTIES IMPORTED_LOCATION "${MOSAIC_DYNAMIC_LIB}" + IMPORTED_NO_SONAME TRUE + INTERFACE_INCLUDE_DIRECTORIES + "${MOSAIC_INCLUDE_DIR}") + add_dependencies(paimon_mosaic_ffi mosaic_ep) + + install(FILES "${MOSAIC_DYNAMIC_LIB}" DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endmacro() + macro(build_jindosdk_nextarch) message(STATUS "Building jindosdk-nextarch from local source") @@ -1359,6 +1448,89 @@ macro(build_jindosdk_nextarch) add_dependencies(jindosdk::nextarch jindosdk-nextarch_ep) endmacro() +macro(build_oss_sdk_v2) + message(STATUS "Building Alibaba Cloud OSS C++ SDK v2 from source") + find_package(CURL REQUIRED) + find_package(Threads REQUIRED) + + set(OSS_SDK_V2_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/oss_sdk_v2_ep-install") + set(OSS_SDK_V2_INCLUDE_DIR "${OSS_SDK_V2_PREFIX}/include") + set(OSS_SDK_V2_INSTALL_LIBDIR "${CMAKE_INSTALL_LIBDIR}") + set(OSS_SDK_V2_LIB_DIR "${OSS_SDK_V2_PREFIX}/${OSS_SDK_V2_INSTALL_LIBDIR}") + set(OSS_SDK_V2_STATIC_LIB + "${OSS_SDK_V2_LIB_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}alibabacloud-oss-cpp-sdk-v2${CMAKE_STATIC_LIBRARY_SUFFIX}" + ) + + set(OSS_SDK_V2_CMAKE_ARGS + ${EP_COMMON_CMAKE_ARGS} + "-DCMAKE_INSTALL_PREFIX=${OSS_SDK_V2_PREFIX}" + "-DCMAKE_INSTALL_LIBDIR=${OSS_SDK_V2_INSTALL_LIBDIR}" + -DCMAKE_PARENT_CXX_STANDARD=17 + -DBUILD_SHARED_LIBS=OFF + -DBUILD_TESTS=OFF + -DBUILD_SAMPLES=OFF + -DENABLE_RTTI=OFF + -DUSE_CURL_TRANSPORT=ON + -DUSE_SYSTEM_CURL=ON + -DUSE_SYSTEM_OPENSSL=OFF + -DUSE_SYSTEM_MBEDTLS=OFF + -DUSE_SYSTEM_TINYXML2=OFF + -DUSE_STD_EXPECTED=OFF + -DENABLE_ENCRYPTION=OFF) + set(OSS_SDK_V2_CURL_INCLUDE_DIR "${CURL_INCLUDE_DIR}") + if(NOT OSS_SDK_V2_CURL_INCLUDE_DIR AND CURL_INCLUDE_DIRS) + list(GET CURL_INCLUDE_DIRS 0 OSS_SDK_V2_CURL_INCLUDE_DIR) + endif() + set(OSS_SDK_V2_CURL_LIBRARY "${CURL_LIBRARY_RELEASE}") + if(NOT OSS_SDK_V2_CURL_LIBRARY) + set(OSS_SDK_V2_CURL_LIBRARY "${CURL_LIBRARY}") + endif() + if(TARGET CURL::libcurl) + if(NOT OSS_SDK_V2_CURL_INCLUDE_DIR) + get_target_property(OSS_SDK_V2_CURL_INCLUDE_DIR CURL::libcurl + INTERFACE_INCLUDE_DIRECTORIES) + endif() + if(NOT OSS_SDK_V2_CURL_LIBRARY) + foreach(CURL_CONFIG RELEASE RELWITHDEBINFO DEBUG NOCONFIG) + get_target_property(OSS_SDK_V2_CURL_LIBRARY CURL::libcurl + "IMPORTED_LOCATION_${CURL_CONFIG}") + if(OSS_SDK_V2_CURL_LIBRARY) + break() + endif() + endforeach() + endif() + if(NOT OSS_SDK_V2_CURL_LIBRARY) + get_target_property(OSS_SDK_V2_CURL_LIBRARY CURL::libcurl IMPORTED_LOCATION) + endif() + endif() + if(OSS_SDK_V2_CURL_INCLUDE_DIR AND OSS_SDK_V2_CURL_LIBRARY) + list(APPEND + OSS_SDK_V2_CMAKE_ARGS + "-DCURL_INCLUDE_DIR=${OSS_SDK_V2_CURL_INCLUDE_DIR}" + "-DCURL_LIBRARY=${OSS_SDK_V2_CURL_LIBRARY}" + "-DCURL_LIBRARY_RELEASE=${OSS_SDK_V2_CURL_LIBRARY}") + endif() + + externalproject_add(oss_sdk_v2_ep + ${EP_COMMON_OPTIONS} + URL ${OSS_SDK_V2_SOURCE_URL} + URL_HASH "SHA256=${PAIMON_OSS_SDK_V2_BUILD_SHA256_CHECKSUM}" + CMAKE_ARGS ${OSS_SDK_V2_CMAKE_ARGS} ${THIRDPARTY_LOG_OPTIONS} + BUILD_BYPRODUCTS "${OSS_SDK_V2_STATIC_LIB}") + + file(MAKE_DIRECTORY "${OSS_SDK_V2_INCLUDE_DIR}") + file(MAKE_DIRECTORY "${OSS_SDK_V2_LIB_DIR}") + + add_library(alibabacloud_oss_v2::oss STATIC IMPORTED) + set_target_properties(alibabacloud_oss_v2::oss + PROPERTIES IMPORTED_LOCATION "${OSS_SDK_V2_STATIC_LIB}" + INTERFACE_INCLUDE_DIRECTORIES + "${OSS_SDK_V2_INCLUDE_DIR}") + target_link_libraries(alibabacloud_oss_v2::oss INTERFACE CURL::libcurl + Threads::Threads) + add_dependencies(alibabacloud_oss_v2::oss oss_sdk_v2_ep) +endmacro() + macro(build_protobuf) message(STATUS "Building protobuf from source") set(PROTOBUF_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/protobuf_ep-install") @@ -1953,6 +2125,9 @@ paimon_warn_if_mixed_arrow_dependencies() resolve_dependency(TBB) resolve_dependency(glog) +if(PAIMON_ENABLE_MOSAIC) + build_mosaic() +endif() if(PAIMON_ENABLE_AVRO) resolve_dependency(Avro) endif() @@ -1967,6 +2142,9 @@ if(PAIMON_ENABLE_JINDO) build_jindosdk_c() build_jindosdk_nextarch() endif() +if(PAIMON_ENABLE_OSS) + resolve_dependency(OSS_SDK_V2) +endif() if(PAIMON_ENABLE_S3) include(BuildAwsAuth) build_aws_auth() diff --git a/cmake_modules/arrow.diff b/cmake_modules/arrow.diff index f86f36e8a..bb71e9c39 100644 --- a/cmake_modules/arrow.diff +++ b/cmake_modules/arrow.diff @@ -1,22 +1,345 @@ -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 +@@ -814,5 +814,6 @@ if(DEFINED ENV{ARROW_THRIFT_URL}) + set(THRIFT_SOURCE_URL "$ENV{ARROW_THRIFT_URL}") + else() + set_urls(THRIFT_SOURCE_URL ++ "https://archive.apache.org/dist/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" + "https://www.apache.org/dyn/closer.cgi?action=download&filename=/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" + "https://downloads.apache.org/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" +@@ -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 +378,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 +634,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 +968,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 +980,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 +1007,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 +1030,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/cmake_modules/orc.diff b/cmake_modules/orc.diff index e4ca4e299..88742dbdc 100644 --- a/cmake_modules/orc.diff +++ b/cmake_modules/orc.diff @@ -435,3 +435,12 @@ index 9b2c829c7..434841224 100644 set(ORC_FORMAT_VERSION "1.0.0") set(LZ4_VERSION "1.10.0") set(SNAPPY_VERSION "1.2.1") +@@ -140,7 +142,7 @@ if(DEFINED ENV{ORC_FORMAT_URL}) + set(ORC_FORMAT_SOURCE_URL "$ENV{ORC_FORMAT_URL}") + message(STATUS "Using ORC_FORMAT_URL: ${ORC_FORMAT_SOURCE_URL}") + else() +- set(ORC_FORMAT_SOURCE_URL "https://www.apache.org/dyn/closer.lua/orc/orc-format-${ORC_FORMAT_VERSION}/orc-format-${ORC_FORMAT_VERSION}.tar.gz?action=download" ) ++ set(ORC_FORMAT_SOURCE_URL "https://archive.apache.org/dist/orc/orc-format-${ORC_FORMAT_VERSION}/orc-format-${ORC_FORMAT_VERSION}.tar.gz" ) + message(STATUS "Using DEFAULT URL: ${ORC_FORMAT_SOURCE_URL}") + endif() + ExternalProject_Add (orc-format_ep diff --git a/docs/source/build_system.rst b/docs/source/build_system.rst index a2a1a9e7b..40964f6f8 100644 --- a/docs/source/build_system.rst +++ b/docs/source/build_system.rst @@ -106,6 +106,8 @@ Paimon provides a set of built-in optional plugins that you can link to as neede - ``Paimon::paimon_local_file_system_shared`` - ``Paimon::paimon_jindo_file_system_shared`` + - ``Paimon::paimon_oss_file_system_shared`` + - ``Paimon::paimon_s3_file_system_shared`` - Index plugins: diff --git a/docs/source/building.rst b/docs/source/building.rst index 466d461b4..a057ff539 100644 --- a/docs/source/building.rst +++ b/docs/source/building.rst @@ -178,11 +178,14 @@ boolean flags to ``cmake``. * ``-DPAIMON_ENABLE_ORC=ON``: Paimon integration with Apache ORC * ``-DPAIMON_ENABLE_AVRO=ON``: Apache Avro libraries and Paimon integration * ``-DPAIMON_ENABLE_JINDO=ON``: Support for Alibaba Jindo filesystems +* ``-DPAIMON_ENABLE_OSS=ON``: Support for Alibaba Cloud OSS through OSS SDK V2 +* ``-DPAIMON_ENABLE_S3=ON``: Support for Amazon S3-compatible filesystems * ``-DPAIMON_ENABLE_LUMINA=ON``: Support for the Lumina vector index. Requires Linux ``x86_64``; see :ref:`cpp-building-platforms`. * ``-DPAIMON_ENABLE_LUCENE=ON``: Support for Lucene full-text search indexes * ``-DPAIMON_ENABLE_TANTIVY=ON``: Enable the experimental Tantivy full-text index Rust FFI. -* ``-DPAIMON_ENABLE_REST=ON``: Support for the REST catalog (``metastore=rest``), requires the libcurl development package. +* ``-DPAIMON_ENABLE_REST=ON``: Support for the REST catalog + (``metastore=rest``), requires the libcurl and OpenSSL development packages. Third-party dependency source ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/examples/benchmark.rst b/docs/source/examples/benchmark.rst index d2ff1cdef..b82114eea 100644 --- a/docs/source/examples/benchmark.rst +++ b/docs/source/examples/benchmark.rst @@ -19,8 +19,17 @@ Benchmark Usage ================ -Paimon C++ provides Google Benchmark based cases for append-table write/read and -primary-key table write/MOR read paths. Benchmarks are disabled by default. +Paimon C++ provides Google Benchmark based cases at two levels: + +``paimon-read-write-benchmark`` + Table-level cases for append-table write/read and primary-key table write/MOR + read paths. + +``paimon-parquet-format-benchmark`` + Format-level cases that drive the Parquet writer and reader directly, without + catalog lookup, split planning, merge/sort or commit. + +Benchmarks are disabled by default. Build ===== @@ -29,13 +38,14 @@ Enable benchmarks when configuring CMake:: cmake -S . -B build -DPAIMON_BUILD_BENCHMARKS=ON cmake --build build --target paimon-read-write-benchmark + cmake --build build --target paimon-parquet-format-benchmark Run all benchmark cases through CTest:: cmake --build build --target benchmark -Custom Options -============== +Table-level Custom Options +========================== ``paimon-read-write-benchmark`` accepts Google Benchmark options plus the Paimon specific options below: @@ -88,3 +98,97 @@ MOR read from an existing table:: --paimon_source_table_path /path/table \ --paimon_pk_columns=id \ --benchmark_filter=BM_MOR_Read/4 + +Parquet Format Benchmark +======================== + +``paimon-parquet-format-benchmark`` takes only Google Benchmark options. It +generates its own data and writes it to a temporary directory, so it needs no +source file or table. + +Two things shape how the results should be read: + +- Every axis is swept on its own rather than as a combined matrix, so each + case answers one question and a change can be attributed to it. +- Writes go through the local FileSystem into a temporary directory, so + absolute numbers carry the cost of that path. Comparisons are meaningful + only under the same environment and methodology - the same machine, build + configuration and options - which is what makes a before/after comparison + useful. + +Writer cases (``BM_ParquetWrite_*``) cover flat primitives, ``VARCHAR`` at low / +medium / high cardinality with and without file-level dictionary encoding, +already dictionary-encoded ``VARCHAR`` / ``INTEGER`` input arrays against their +flat equivalents, ``DECIMAL`` at precision 9 / 18 / 38, nested ``STRUCT`` / +``LIST`` / ``VECTOR`` / ``MAP``, null density from 0 to 100 percent, rows per +``AddBatch`` call, column count at a fixed row count, row group size, the +writer memory threshold that triggers a byte-based row-group flush, and the +codecs Parquet accepts - ``none``, ``snappy``, ``gzip``, ``brotli``, ``zstd``, +``lz4_raw`` and ``lz4_hadoop``. Note that ``lz4`` is deliberately not among +them: it resolves to Arrow's ``LZ4_FRAME``, which +``parquet::IsCodecSupported`` rejects. + +The two dictionary axes are different questions. ``BM_ParquetWrite_String`` and +``BM_ParquetWrite_StringNoDictionary`` vary whether the *file* is dictionary +encoded; ``BM_ParquetWrite_Dictionary*`` vary whether the *input array* already +is, which is what decides whether Arrow can pass indices through to Parquet or +has to materialize them first. + +Reader cases (``BM_ParquetRead_*``) cover full scan, single-column projection, +predicate-filtered reads at varying selectivity with page-index filtering on and +off, skip-heavy reads driven by a strided selection bitmap, null density, +``DECIMAL`` at precision 9 / 18 / 38, ``DOUBLE``, dictionary-encoded against +plain-encoded files, rows per ``NextBatch`` call, and nested column reads. + +Every case reports ``ns_per_row`` next to ``bytes_per_row`` - ``file_bytes`` for +writes, ``read_bytes`` for reads - so a change that trades CPU for size is +visible in both directions. Read cases additionally report ``rows_read``, +``batches``, and ``row_groups`` / ``row_groups_after_filter`` from the reader's +own metrics. + +Compare filtered cases on ``ns_per_input_row`` and ``bytes_per_input_row``, not +``ns_per_row`` and ``bytes_per_row``. The latter pair divides by the rows a case +actually materialized, so pruning shrinks numerator and denominator together and +they can rise even as the run gets faster; the ``_input_row`` pair divides by the +rows the file holds, which every setting shares. + +``row_groups_after_filter`` counts row groups only. It does not show page-level +pruning: on the ordered ``id`` column both page-index settings usually keep the +same row groups, and the page-index gain shows up in ``rows_read``, +``read_bytes`` and ``ns_per_input_row`` instead. + +A case that cannot run - an unsupported codec, a schema the reader rejects - +calls ``SkipWithError`` and makes the process exit non-zero, so ``ctest -L +benchmark`` fails instead of reporting a silent skip. Read cases also assert on +the number of rows they materialized, so a fixture that stopped producing rows +fails rather than looking fast. + +Because the benchmark is only compiled under ``PAIMON_BUILD_BENCHMARKS``, the +format-layer assumptions it relies on are covered separately by +``paimon-parquet-format-benchmark-test``, which builds with the normal test +suite. + +Each read case scans a file that is generated once on first use and reused for +the rest of the run, so a filtered run only pays to build the fixtures its own +cases need. + +All Parquet writer cases:: + + paimon-parquet-format-benchmark --benchmark_filter=BM_ParquetWrite + +Page-index filtering at 1% selectivity, on and off - compare ``rows_read``, +``read_bytes`` and ``ns_per_input_row`` between the two:: + + paimon-parquet-format-benchmark \ + --benchmark_filter='BM_ParquetRead_Filtered/keep_pct:1/' + +Read batch size sweep, repeated for a stable comparison:: + + paimon-parquet-format-benchmark \ + --benchmark_filter=BM_ParquetRead_BatchSize \ + --benchmark_repetitions=5 \ + --benchmark_report_aggregates_only=true + +Null density on both sides, to see what definition levels cost:: + + paimon-parquet-format-benchmark --benchmark_filter='Parquet(Write|Read)_Nulls' diff --git a/docs/source/user_guide.rst b/docs/source/user_guide.rst index 38ef0fc64..fdee89d9a 100644 --- a/docs/source/user_guide.rst +++ b/docs/source/user_guide.rst @@ -38,6 +38,7 @@ User Guide user_guide/commit user_guide/compaction user_guide/read + user_guide/metrics user_guide/clean user_guide/prefetch user_guide/arrow diff --git a/docs/source/user_guide/catalog.rst b/docs/source/user_guide/catalog.rst index c57c50c12..b695044dd 100644 --- a/docs/source/user_guide/catalog.rst +++ b/docs/source/user_guide/catalog.rst @@ -50,9 +50,25 @@ registered on the REST server. The catalog is configured through the * ``metastore``: must be ``rest`` to select the REST catalog. * ``uri``: server url of the REST catalog server. -* ``token.provider``: authentication provider of the REST catalog; currently only - ``bear`` is supported (the protocol's historical spelling of "bearer"). +* ``token.provider``: authentication provider of the REST catalog. ``bear`` + (the protocol's historical spelling of "bearer") and ``dlf`` are supported. * ``token``: token of the ``bear`` token provider. +* ``dlf.region``: region used by DLF request signing. It is inferred from the + endpoint URI when omitted. +* ``dlf.access-key-id`` and ``dlf.access-key-secret``: static DLF access key. +* ``dlf.security-token``: optional STS security token used with a static access + key. +* ``dlf.token-path``: path to a JSON file containing refreshable DLF credentials. +* ``dlf.token-loader``: refreshable credential loader. ``local_file`` reads + ``dlf.token-path`` and ``ecs`` obtains an STS token from an ECS RAM role. +* ``dlf.token-ecs-metadata-url``: ECS RAM role metadata endpoint. It defaults to + ``http://100.100.100.200/latest/meta-data/Ram/security-credentials/``. +* ``dlf.token-ecs-role-name``: optional ECS RAM role name. The loader discovers + the role from the metadata endpoint when it is omitted. +* ``dlf.signing-algorithm``: ``default`` selects DLF4-HMAC-SHA256 for DLF VPC + endpoints and ``openapi`` selects ROA HMAC-SHA1 for DlfNext OpenAPI endpoints. + When omitted, an endpoint containing ``dlfnext`` selects ``openapi`` and other + endpoints select ``default``. * ``table-default.``: table option defaults applied when a created table left ```` unset. * ``header.``: sent as the ```` http header on every request to the @@ -70,6 +86,25 @@ registered on the REST server. The catalog is configured through the PAIMON_ASSIGN_OR_RAISE(std::unique_ptr catalog, paimon::Catalog::Create(/*root_path=*/"my_instance", options)); +For DLF, configure one credential source. Static AK/SK credentials, an optional +STS token, a refreshable local token file, and ECS RAM role credentials are +supported. A local or ECS token has the Java-compatible JSON fields +``AccessKeyId``, ``AccessKeySecret``, ``SecurityToken`` and ``Expiration``. The +last field uses UTC ``yyyy-MM-dd'T'HH:mm:ss'Z'`` format. Refreshable credentials +are reloaded when less than one hour of validity remains. + +.. code-block:: cpp + + std::map options = { + {"metastore", "rest"}, + {"uri", "https://dlfnext.cn-hangzhou.aliyuncs.com"}, + {"token.provider", "dlf"}, + {"dlf.access-key-id", ""}, + {"dlf.access-key-secret", ""}, + // Optional for temporary credentials: + {"dlf.security-token", ""}, + }; + On creation the catalog queries the server's ``/v1/config`` endpoint and merges its response with the options above: the server's overrides win over the client options, which in turn win over the server's defaults. @@ -81,5 +116,4 @@ through the regular ``Catalog`` API, and table snapshots can be listed through The C++ REST catalog covers the database, table and snapshot operations of the ``Catalog`` API. The parts of the Java REST catalog that have no C++ counterpart yet — altering a database or a table, views, functions, partitions, tags, branch -management and consumers — are not supported, and neither is the ``dlf`` token -provider. +management and consumers — are not supported. diff --git a/docs/source/user_guide/data_types.rst b/docs/source/user_guide/data_types.rst index 3d529332b..add537adb 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/docs/source/user_guide/metrics.rst b/docs/source/user_guide/metrics.rst new file mode 100644 index 000000000..72ee2a176 --- /dev/null +++ b/docs/source/user_guide/metrics.rst @@ -0,0 +1,116 @@ +.. 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. + +Metrics reference +================= + +``Metrics`` contains counters, gauges, and histogram snapshots. A counter is a non-negative +integer, a gauge represents current state, and a histogram records a distribution of observed +values. The scan and prefetch-reader metrics described below are returned as point-in-time +snapshots. Modifying a returned snapshot does not modify the component that produced it. + +Scan planning +------------- + +The names below are declared by ``ScanMetrics``. ``last*`` counters are replaced after each +successful ``CreatePlan()`` call. Histograms and the cache hit/miss counters accumulate for the +lifetime of the scan. They are unrelated to ``SetReadSchema()`` and ``ReadAheadCache::Reset()``. +The first six names match Java ``ScanMetrics``; the remaining names are C++-only. + +.. csv-table:: + :header: "Name", "Type", "Unit", "Meaning" + :widths: 34, 12, 12, 52 + + "lastScanDuration", "counter", "milliseconds", "Duration of the last successful plan" + "scanDuration", "histogram", "milliseconds", "Distribution of successful plan durations" + "lastScannedSnapshotId", "counter", "snapshot ID", "Snapshot used by the last plan, or 0" + "lastScannedManifests", "counter", "files", "Manifest files selected by the last plan" + "lastScanSkippedTableFiles", "counter", "files", "Table files skipped by the last plan" + "lastScanResultedTableFiles", "counter", "files", "Table files returned by the last plan" + "lastManifestReadDuration", "counter", "milliseconds", "Manifest-list and entry read time for the last plan" + "manifestReadDuration", "histogram", "milliseconds", "Distribution of manifest read times" + "lastSnapshotCacheEnabled", "counter", "boolean", "Whether snapshot manifest-entry cache was eligible" + "lastSnapshotCacheHit", "counter", "boolean", "Whether the last eligible lookup hit" + "snapshotCacheHits", "counter", "lookups", "Cumulative exact-snapshot cache hits" + "snapshotCacheMisses", "counter", "lookups", "Cumulative eligible cache misses" + "lastSnapshotCacheLoadDuration", "counter", "milliseconds", "Cache load time for the last plan" + "snapshotCacheLoadDuration", "histogram", "milliseconds", "Distribution of cache load times" + "lastSnapshotCacheStoreDuration", "counter", "milliseconds", "Cache store time for the last plan; 0 when not stored" + "snapshotCacheStoreDuration", "histogram", "milliseconds", "Distribution of cache store times" + "lastLazyDecodeScannedRows", "counter", "manifest rows", "Candidate manifest rows inspected by the last plan" + "lastLazyDecodeMaterializedRows", "counter", "manifest rows", "Manifest rows retained after lazy filtering" + +Prefetch reader +--------------- + +The names below are declared by ``PrefetchMetrics``. Counters and histograms accumulate for the +lifetime of the prefetch reader, including across ``SetReadSchema()``. ``enabled`` and +``parallelism`` describe the most recently initialized schema. ``queue-depth`` is reset by +``SetReadSchema()`` and ``Close()``; ``queue-depth.max`` remains the lifetime maximum. +These metrics are C++-only and have no counterparts in Java Paimon. + +.. csv-table:: + :header: "Name", "Type", "Unit", "Meaning" + :widths: 38, 12, 12, 48 + + "prefetch.enabled", "gauge", "boolean", "Whether the most recently initialized schema selected prefetch" + "prefetch.parallelism", "gauge", "readers", "Effective reader parallelism" + "prefetch.read-ranges.total", "counter", "ranges", "Generated ranges before bitmap filtering" + "prefetch.read-ranges.after-bitmap", "counter", "ranges", "Ranges retained after bitmap filtering" + "prefetch.seek.count", "counter", "operations", "Underlying reader seek operations" + "prefetch.produced-batches", "counter", "batches", "Data batches placed into prefetch queues" + "prefetch.consumed-batches", "counter", "batches", "Data batches returned to the consumer" + "prefetch.discarded-batches", "counter", "batches", "Data batches released without consumption, plus EOF entries released during cleanup" + "prefetch.errors", "counter", "errors", "Errors recorded by the background prefetch loop" + "prefetch.adaptive-disabled-count", "counter", "decisions", "Times adaptive strategy disabled prefetch" + "prefetch.queue-full-count", "counter", "events", "Times production found a full queue" + "prefetch.queue-depth", "gauge", "queue entries", "Current queued entries, including retained EOF markers" + "prefetch.queue-depth.max", "gauge", "queue entries", "Maximum queued entries in the reader lifetime" + "prefetch.reader-read-latency-us", "histogram", "microseconds", "Underlying reader batch latency" + "prefetch.consumer-wait-latency-us", "histogram", "microseconds", "Consumer wait latency per returned batch or EOF" + +Prefetch I/O +------------ + +``PrefetchIoMetrics`` describes only I/O that passes through the prefetch reader's instrumented +input streams. It is not a whole-query or whole-table I/O total. All counters accumulate for the +reader lifetime and are retained across ``SetReadSchema()`` and cache reset. Latency uses relaxed +atomic count and sum counters instead of per-I/O histograms to reduce hot-path cost. Collection is +disabled by default; set ``prefetch.io-metrics.enabled`` to ``true`` in the read options to enable +it. When disabled, these per-I/O metrics are absent and the input streams have no metrics +instrumentation. +``io.async.pending`` is current state and returns to zero when all callbacks complete. +These metrics are C++-only and have no counterparts in Java Paimon. + +.. csv-table:: + :header: "Name", "Type", "Unit", "Meaning" + :widths: 34, 12, 12, 52 + + "io.read.requests", "counter", "requests", "Synchronous read requests" + "io.read.requested-bytes", "counter", "bytes", "Bytes requested by synchronous reads" + "io.read.physical-bytes", "counter", "bytes", "Bytes returned by successful synchronous reads" + "io.read.failed", "counter", "requests", "Failed synchronous reads" + "io.read.latency.count", "counter", "requests", "Completed synchronous read latency samples" + "io.read.latency.sum-us", "counter", "microseconds", "Sum of synchronous read latency" + "io.async.requests", "counter", "requests", "Asynchronous read requests" + "io.async.requested-bytes", "counter", "bytes", "Bytes requested by asynchronous reads" + "io.async.physical-bytes", "counter", "bytes", "Bytes attributed to successful asynchronous reads" + "io.async.completed", "counter", "requests", "Successful asynchronous reads" + "io.async.failed", "counter", "requests", "Failed asynchronous reads" + "io.async.pending", "gauge", "requests", "Asynchronous callbacks not yet completed" + "io.async.latency.count", "counter", "requests", "Completed asynchronous callback latency samples" + "io.async.latency.sum-us", "counter", "microseconds", "Sum of asynchronous callback latency" diff --git a/include/paimon/api.h b/include/paimon/api.h index d75666850..81f236bc4 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/catalog_options.h b/include/paimon/catalog_options.h index f58a876c3..05b1af278 100644 --- a/include/paimon/catalog_options.h +++ b/include/paimon/catalog_options.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -32,10 +34,37 @@ struct PAIMON_EXPORT CatalogOptions { /// "token" - Token of the "bear" token provider of the REST catalog. static const char TOKEN[]; - /// "token.provider" - Authentication provider of the REST catalog. Only "bear" is - /// supported ("bear" is the protocol's historical spelling of "bearer", do not "fix" it). + /// "token.provider" - Authentication provider of the REST catalog. Supported values are + /// "bear" (the protocol's historical spelling of "bearer") and "dlf". static const char TOKEN_PROVIDER[]; + /// "dlf.region" - Region used by DLF request signing. Inferred from URI when absent. + static const char DLF_REGION[]; + + /// "dlf.token-path" - Path of a JSON file containing refreshable DLF credentials. + static const char DLF_TOKEN_PATH[]; + + /// "dlf.access-key-id" - DLF access key id. + static const char DLF_ACCESS_KEY_ID[]; + + /// "dlf.access-key-secret" - DLF access key secret. + static const char DLF_ACCESS_KEY_SECRET[]; + + /// "dlf.security-token" - Optional STS security token used with a DLF access key. + static const char DLF_SECURITY_TOKEN[]; + + /// "dlf.token-loader" - Refreshable DLF token loader ("ecs" or "local_file"). + static const char DLF_TOKEN_LOADER[]; + + /// "dlf.token-ecs-metadata-url" - ECS RAM role metadata endpoint. + static const char DLF_TOKEN_ECS_METADATA_URL[]; + + /// "dlf.token-ecs-role-name" - Optional ECS RAM role name. + static const char DLF_TOKEN_ECS_ROLE_NAME[]; + + /// "dlf.signing-algorithm" - DLF signer ("default" or "openapi"). + static const char DLF_SIGNING_ALGORITHM[]; + /// "table-default." - Prefix of the catalog options that provide table option /// defaults: "table-default.=" applies "=" to a created /// table when the caller left "" unset. diff --git a/include/paimon/defs.h b/include/paimon/defs.h index 9fcf8e34d..e05faefd8 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, }; @@ -154,6 +156,10 @@ struct PAIMON_EXPORT Options { /// 9, but the read and write speed will significantly decrease. Default value is 1. static const char FILE_COMPRESSION_ZSTD_LEVEL[]; + /// "file.block-size" - File block size of format, default value of orc stripe is 64 MB, + /// parquet row group is 128 MB, and Mosaic row group is 256 MB. + static const char FILE_BLOCK_SIZE[]; + /// "manifest.target-file-size" - Suggested file size of a manifest file. /// Default value is 8MB. static const char MANIFEST_TARGET_FILE_SIZE[]; @@ -198,6 +204,14 @@ 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[]; + + /// "prefetch.io-metrics.enabled" - Whether to collect per-I/O metrics for prefetch reads. + /// Default value is false. + static const char PREFETCH_IO_METRICS_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[]; @@ -372,14 +386,34 @@ struct PAIMON_EXPORT Options { /// @note: bitmap64 dv is not supported. static const char DELETION_VECTOR_BITMAP64[]; - /// @note `CHANGELOG_PRODUCER` currently only support `none` - /// /// "changelog-producer" - Whether to double write to a changelog file. This changelog file /// keeps the details of data changes, it can be read directly during stream reads. This can be /// applied to tables with primary keys. Values can be "none", "input", "lookup", /// "full-compaction". Default value is "none". + /// @note C++ Paimon currently supports "none", "input", and "lookup". static const char CHANGELOG_PRODUCER[]; + /// "changelog-producer.row-deduplicate" - Whether to generate update-before and update-after + /// changelog records when the row has not changed. This option is only valid for "lookup" or + /// "full-compaction" changelog producers. Default value is "false". + static const char CHANGELOG_PRODUCER_ROW_DEDUPLICATE[]; + + /// "changelog-producer.row-deduplicate-ignore-fields" - Comma-separated fields to ignore when + /// comparing rows for changelog deduplication. This option is only valid when + /// "changelog-producer.row-deduplicate" is "true". + static const char CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS[]; + + /// "changelog-file.prefix" - Specify the file name prefix of changelog files. Default value is + /// "changelog-". + static const char CHANGELOG_FILE_PREFIX[]; + + /// "changelog-file.format" - Specify the file format of changelog files. No default value. + static const char CHANGELOG_FILE_FORMAT[]; + + /// "changelog-file.compression" - Specify the compression of changelog files. No default + /// value. + static const char CHANGELOG_FILE_COMPRESSION[]; + /// "force-lookup" - Whether to force the use of lookup for compaction. Default value is /// "false". static const char FORCE_LOOKUP[]; @@ -403,6 +437,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 +590,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 6e12d4561..a3030e5ea 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 b46dee8c4..3993b6247 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 63efb5621..8af776959 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 fdc172c3a..1d7ca0888 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 f16cb2547..a4e3de484 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/src/paimon/core/operation/metrics/scan_metrics.h b/include/paimon/format/read_hints.h similarity index 59% rename from src/paimon/core/operation/metrics/scan_metrics.h rename to include/paimon/format/read_hints.h index 483ab7afc..d60b41320 100644 --- a/src/paimon/core/operation/metrics/scan_metrics.h +++ b/include/paimon/format/read_hints.h @@ -18,18 +18,18 @@ #pragma once +#include "paimon/visibility.h" + namespace paimon { -/// Metrics to measure scan operation. -class ScanMetrics { - public: - static constexpr char LAST_SCAN_DURATION[] = "lastScanDuration"; - // Histogram metric for scan plan duration (milliseconds). - static constexpr char SCAN_DURATION[] = "scanDuration"; - static constexpr char LAST_SCANNED_SNAPSHOT_ID[] = "lastScannedSnapshotId"; - static constexpr char LAST_SCANNED_MANIFESTS[] = "lastScannedManifests"; - static constexpr char LAST_SCAN_SKIPPED_TABLE_FILES[] = "lastScanSkippedTableFiles"; - static constexpr char LAST_SCAN_RESULTED_TABLE_FILES[] = "lastScanResultedTableFiles"; +/// 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 b0837a262..5a28077a8 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/predicate/literal.h b/include/paimon/predicate/literal.h index ef5864573..168d6a451 100644 --- a/include/paimon/predicate/literal.h +++ b/include/paimon/predicate/literal.h @@ -91,13 +91,12 @@ class PAIMON_EXPORT Literal { std::string ToString() const; /// Gets the hash code for this literal. - /// @note HashCode() hashes the exact bit representation (including Decimal scale), while - /// operator== delegates to CompareTo() which uses numeric equality (e.g. decimals with - /// different scales can compare equal). This means the hash-equality contract (equal objects - /// must have equal hashes) may be violated for Decimal literals with different scales. In - /// practice this is safe because all current std::unordered_map usages (bitmap - /// file index) only store values from the same column, which guarantees a fixed precision and - /// scale. + /// @note HashCode() canonicalizes all floating-point NaNs so that values considered equal by + /// CompareTo() have the same hash. Decimal values include their scale in the hash, while + /// CompareTo() uses numeric equality, so Decimal literals with different scales can still + /// violate the hash-equality contract. In practice this is safe because all current + /// std::unordered_map usages only store values from the same column, which has a + /// fixed precision and scale. size_t HashCode() const; /// Compares this literal with another literal. The comparison follows SQL semantics for the diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 9bed54024..e645645ba 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 { @@ -51,7 +51,7 @@ class PAIMON_EXPORT ReadContext { const std::vector& read_field_names, const std::vector& read_field_ids, const std::shared_ptr& predicate, bool enable_predicate_filter, - bool enable_prefetch, uint32_t prefetch_batch_count, + bool enable_prefetch, bool enable_late_materializing, uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, const std::optional& table_schema, const std::shared_ptr& memory_pool, @@ -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 { @@ -98,6 +97,9 @@ class PAIMON_EXPORT ReadContext { bool EnablePrefetch() const { return enable_prefetch_; } + bool EnableLateMaterializing() const { + return enable_late_materializing_; + } uint32_t GetPrefetchBatchCount() const { return prefetch_batch_count_; } @@ -128,8 +130,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 { @@ -164,6 +166,7 @@ class PAIMON_EXPORT ReadContext { std::shared_ptr predicate_; bool enable_predicate_filter_; bool enable_prefetch_; + bool enable_late_materializing_; uint32_t prefetch_batch_count_; uint32_t prefetch_max_parallel_num_; bool enable_multi_thread_row_to_batch_; @@ -175,7 +178,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 +310,22 @@ 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 late materialization (probe/payload two-phase reads). When enabled, + /// each parallel reader under the prefetch layer performs a probe read of predicate + /// columns first and only materializes payload columns for matched rows. + /// @param enabled Whether to enable late materialization (default: false) + /// @return Reference to this builder for method chaining. + /// @note Without a pushed-down predicate the late-materializing reader degrades to a + /// plain passthrough. + ReadContextBuilder& EnableLateMaterializing(bool enabled); + + /// 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/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index acc7d0bbd..842279956 100644 --- a/include/paimon/reader/prefetch_file_batch_reader.h +++ b/include/paimon/reader/prefetch_file_batch_reader.h @@ -27,6 +27,47 @@ namespace paimon { +/// C++-only prefetch reader metrics. Java Paimon has no corresponding metrics. +class PAIMON_EXPORT PrefetchMetrics { + public: + static constexpr char ENABLED[] = "prefetch.enabled"; + static constexpr char PARALLELISM[] = "prefetch.parallelism"; + static constexpr char READ_RANGES_TOTAL[] = "prefetch.read-ranges.total"; + static constexpr char READ_RANGES_AFTER_BITMAP[] = "prefetch.read-ranges.after-bitmap"; + static constexpr char SEEK_COUNT[] = "prefetch.seek.count"; + static constexpr char PRODUCED_BATCHES[] = "prefetch.produced-batches"; + static constexpr char CONSUMED_BATCHES[] = "prefetch.consumed-batches"; + static constexpr char DISCARDED_BATCHES[] = "prefetch.discarded-batches"; + static constexpr char ERRORS[] = "prefetch.errors"; + static constexpr char ADAPTIVE_DISABLED_COUNT[] = "prefetch.adaptive-disabled-count"; + static constexpr char QUEUE_FULL_COUNT[] = "prefetch.queue-full-count"; + static constexpr char QUEUE_DEPTH[] = "prefetch.queue-depth"; + static constexpr char QUEUE_DEPTH_MAX[] = "prefetch.queue-depth.max"; + static constexpr char READER_READ_LATENCY_US[] = "prefetch.reader-read-latency-us"; + static constexpr char CONSUMER_WAIT_LATENCY_US[] = "prefetch.consumer-wait-latency-us"; +}; + +/// C++-only metric names for I/O observed by the prefetch reader's instrumented input streams. +/// Java Paimon has no corresponding metrics. +/// These metrics do not represent whole-query or whole-table I/O. +class PAIMON_EXPORT PrefetchIoMetrics { + public: + static constexpr char READ_REQUESTS[] = "io.read.requests"; + static constexpr char READ_REQUESTED_BYTES[] = "io.read.requested-bytes"; + static constexpr char READ_PHYSICAL_BYTES[] = "io.read.physical-bytes"; + static constexpr char READ_FAILED[] = "io.read.failed"; + static constexpr char READ_LATENCY_COUNT[] = "io.read.latency.count"; + static constexpr char READ_LATENCY_SUM_US[] = "io.read.latency.sum-us"; + static constexpr char ASYNC_REQUESTS[] = "io.async.requests"; + static constexpr char ASYNC_REQUESTED_BYTES[] = "io.async.requested-bytes"; + static constexpr char ASYNC_PHYSICAL_BYTES[] = "io.async.physical-bytes"; + static constexpr char ASYNC_COMPLETED[] = "io.async.completed"; + static constexpr char ASYNC_FAILED[] = "io.async.failed"; + static constexpr char ASYNC_PENDING[] = "io.async.pending"; + static constexpr char ASYNC_LATENCY_COUNT[] = "io.async.latency.count"; + static constexpr char ASYNC_LATENCY_SUM_US[] = "io.async.latency.sum-us"; +}; + /// The prefetch file batch reader extends the basic FileBatchReader interface for prefetch read, /// if a format implementation inherits from this class, it will automatically support the C++ /// Paimon prefetch capability and integrate with the Paimon prefetch framework. @@ -40,7 +81,7 @@ class PAIMON_EXPORT PrefetchFileBatchReader : public FileBatchReader { /// Retrieves the row number of the next row to be read. /// This method indicates the current read position within the file. /// @return The row number of the next row to read. - virtual uint64_t GetNextRowToRead() const = 0; + virtual Result GetNextRowToRead() const = 0; /// Generates a list of row ranges to be read in batches. /// Each range specifies the start and end row numbers for a batch, diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 3dc257fc8..b3fa630de 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -27,10 +27,7 @@ namespace paimon { 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 9bc870fe8..200e4ba4c 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 9177fc06b..6ae81f1f4 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,33 @@ 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 realtime primary-key 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; + /// 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 realtime primary-key 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 +103,11 @@ 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 requested realtime primary-key + /// transport 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 +144,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 realtime primary-key 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 +157,16 @@ 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 requested realtime primary-key transport schema, including + /// nested field-ID alignment, 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 +189,10 @@ 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 configured with the supplied schema, statistics, options, and memory pool. + /// Creates a store for the requested table mode. + /// 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 dc780ad7e..9c0b9d470 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 000000000..0be0e9104 --- /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/table/source/scan_metrics.h b/include/paimon/table/source/scan_metrics.h new file mode 100644 index 000000000..15bc41ff1 --- /dev/null +++ b/include/paimon/table/source/scan_metrics.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 "paimon/visibility.h" + +namespace paimon { + +/// Metric names for scan planning operations. +class PAIMON_EXPORT ScanMetrics { + public: + static constexpr char LAST_SCAN_DURATION[] = "lastScanDuration"; + // Histogram metric for scan plan duration (milliseconds). + static constexpr char SCAN_DURATION[] = "scanDuration"; + static constexpr char LAST_SCANNED_SNAPSHOT_ID[] = "lastScannedSnapshotId"; + static constexpr char LAST_SCANNED_MANIFESTS[] = "lastScannedManifests"; + static constexpr char LAST_SCAN_SKIPPED_TABLE_FILES[] = "lastScanSkippedTableFiles"; + static constexpr char LAST_SCAN_RESULTED_TABLE_FILES[] = "lastScanResultedTableFiles"; + + // The metrics below are C++-only and do not have counterparts in Java ScanMetrics. + static constexpr char LAST_MANIFEST_READ_DURATION[] = "lastManifestReadDuration"; + // Histogram metric for manifest-list and manifest-entry read duration (milliseconds). + static constexpr char MANIFEST_READ_DURATION[] = "manifestReadDuration"; + static constexpr char LAST_SNAPSHOT_CACHE_ENABLED[] = "lastSnapshotCacheEnabled"; + static constexpr char LAST_SNAPSHOT_CACHE_HIT[] = "lastSnapshotCacheHit"; + static constexpr char SNAPSHOT_CACHE_HITS[] = "snapshotCacheHits"; + static constexpr char SNAPSHOT_CACHE_MISSES[] = "snapshotCacheMisses"; + static constexpr char LAST_SNAPSHOT_CACHE_LOAD_DURATION[] = "lastSnapshotCacheLoadDuration"; + static constexpr char SNAPSHOT_CACHE_LOAD_DURATION[] = "snapshotCacheLoadDuration"; + static constexpr char LAST_SNAPSHOT_CACHE_STORE_DURATION[] = "lastSnapshotCacheStoreDuration"; + static constexpr char SNAPSHOT_CACHE_STORE_DURATION[] = "snapshotCacheStoreDuration"; + // Candidate manifest-entry rows inspected by lazy scan filtering. + static constexpr char LAST_LAZY_DECODE_SCANNED_ROWS[] = "lastLazyDecodeScannedRows"; + // Full manifest entries retained after lazy scan filtering. + static constexpr char LAST_LAZY_DECODE_MATERIALIZED_ROWS[] = "lastLazyDecodeMaterializedRows"; +}; + +} // namespace paimon diff --git a/include/paimon/table/source/table_scan.h b/include/paimon/table/source/table_scan.h index c9b42915f..14c447568 100644 --- a/include/paimon/table/source/table_scan.h +++ b/include/paimon/table/source/table_scan.h @@ -23,6 +23,7 @@ #include "paimon/result.h" #include "paimon/table/source/plan.h" +#include "paimon/table/source/scan_metrics.h" #include "paimon/type_fwd.h" #include "paimon/visibility.h" @@ -44,5 +45,11 @@ class PAIMON_EXPORT TableScan { /// /// @return A Result containing a shared pointer to the created `Plan` or an error status. virtual Result> CreatePlan() = 0; + + /// Retrieve metrics related to scan planning operations. + /// + /// @return A point-in-time snapshot of scan metrics. Mutating the returned object does not + /// affect metrics collected by this scan. + virtual std::shared_ptr GetMetrics() const; }; } // 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 000000000..4bbf1ecd3 --- /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 829f29889..5219d72db 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 fa549eef1..a9dd367ec 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 b0fe91b08..051eba324 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 @@ -133,6 +134,7 @@ set(PAIMON_COMMON_SRCS common/predicate/starts_with.cpp common/reader/batch_reader.cpp common/reader/concat_batch_reader.cpp + common/reader/late_materializing_file_batch_reader.cpp common/reader/predicate_batch_reader.cpp common/reader/prefetch_file_batch_reader_impl.cpp common/reader/reader_utils.cpp @@ -157,6 +159,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 @@ -171,7 +174,7 @@ set(PAIMON_COMMON_SRCS common/data/shredding/map_shared_shredding_batch_converter.cpp common/data/shredding/map_shared_shredding_column_allocator.cpp common/data/shredding/lru_map_shared_shredding_column_allocator.cpp - common/data/shredding/map_shared_shredding_file_reader.cpp + common/data/shredding/map_shared_shredding_read_plan_factory.cpp common/data/shredding/shredding_file_reader.cpp common/utils/delta_varint_compressor.cpp common/utils/fields_comparator.cpp @@ -194,7 +197,11 @@ if(PAIMON_ENABLE_S3 OR PAIMON_ENABLE_REST) list(APPEND PAIMON_COMMON_SRCS common/utils/http_client.cpp) set(PAIMON_CURL_LINK_LIBS CURL::libcurl) endif() -if(PAIMON_ENABLE_S3) +set(PAIMON_REST_LINK_LIBS) +if(PAIMON_ENABLE_REST) + set(PAIMON_REST_LINK_LIBS OpenSSL::Crypto) +endif() +if(PAIMON_ENABLE_OSS OR PAIMON_ENABLE_S3) list(APPEND PAIMON_COMMON_SRCS common/fs/object_store_file_system.cpp) endif() @@ -268,15 +275,20 @@ set(PAIMON_CORE_SRCS core/io/data_file_meta_first_row_id_legacy_serializer.cpp core/io/data_file_meta.cpp core/io/data_file_meta_serializer.cpp + core/io/data_file_meta_write_cols_legacy_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_factories.cpp core/io/key_value_data_file_writer_factory.cpp core/io/key_value_data_file_writer.cpp core/io/key_value_in_memory_record_reader.cpp @@ -373,9 +385,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/realtime_primary_key_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 @@ -444,6 +459,7 @@ if(PAIMON_ENABLE_REST) rest/resource_paths.cpp rest/rest_api.cpp rest/rest_auth.cpp + rest/dlf_auth.cpp rest/rest_catalog.cpp rest/rest_messages.cpp rest/rest_util.cpp) @@ -463,6 +479,7 @@ add_paimon_lib(paimon Threads::Threads RapidJSON ${PAIMON_CURL_LINK_LIBS} + ${PAIMON_REST_LINK_LIBS} DataSketches STATIC_LINK_LIBS arrow @@ -475,6 +492,7 @@ add_paimon_lib(paimon RapidJSON DataSketches ${PAIMON_CURL_LINK_LIBS} + ${PAIMON_REST_LINK_LIBS} SHARED_LINK_FLAGS ${PAIMON_VERSION_SCRIPT_FLAGS}) @@ -575,6 +593,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 @@ -593,6 +612,7 @@ if(PAIMON_BUILD_TESTS) common/predicate/predicate_utils_test.cpp common/predicate/predicate_validator_test.cpp common/reader/concat_batch_reader_test.cpp + common/reader/late_materializing_file_batch_reader_test.cpp common/reader/predicate_batch_reader_test.cpp common/reader/prefetch_file_batch_reader_impl_test.cpp common/reader/reader_utils_test.cpp @@ -611,6 +631,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 +665,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 @@ -666,7 +689,7 @@ if(PAIMON_BUILD_TESTS) common/data/shredding/sequential_map_shared_shredding_column_allocator_test.cpp common/data/shredding/map_shared_shredding_field_dict_test.cpp common/data/shredding/map_shared_shredding_context_test.cpp - common/data/shredding/map_shared_shredding_file_reader_test.cpp + common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp STATIC_LINK_LIBS paimon_shared test_utils_static @@ -675,6 +698,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 +771,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 +794,8 @@ 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_primary_key_reader_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp @@ -791,6 +820,7 @@ if(PAIMON_BUILD_TESTS) core/mergetree/compact/deduplicate_merge_function_test.cpp core/mergetree/compact/first_row_merge_function_test.cpp core/mergetree/compact/first_row_merge_function_wrapper_test.cpp + core/mergetree/compact/internal_row_equalizer_test.cpp core/mergetree/compact/interval_partition_test.cpp core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp @@ -895,6 +925,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 @@ -924,7 +955,7 @@ if(PAIMON_BUILD_TESTS) endif() set(PAIMON_OBJECT_STORE_FS_TEST_SOURCES) - if(PAIMON_ENABLE_S3) + if(PAIMON_ENABLE_OSS OR PAIMON_ENABLE_S3) list(APPEND PAIMON_OBJECT_STORE_FS_TEST_SOURCES common/fs/object_store_file_system_test.cpp) endif() @@ -951,6 +982,7 @@ if(PAIMON_BUILD_TESTS) rest/rest_http_client_test.cpp rest/mock_rest_server.cpp rest/resource_paths_test.cpp + rest/dlf_auth_test.cpp rest/rest_catalog_test.cpp rest/rest_messages_test.cpp rest/rest_util_test.cpp diff --git a/src/paimon/common/catalog_options.cpp b/src/paimon/common/catalog_options.cpp index 6e89e80a8..1a856dad7 100644 --- a/src/paimon/common/catalog_options.cpp +++ b/src/paimon/common/catalog_options.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -22,6 +24,15 @@ const char CatalogOptions::METASTORE[] = "metastore"; const char CatalogOptions::URI[] = "uri"; const char CatalogOptions::TOKEN[] = "token"; const char CatalogOptions::TOKEN_PROVIDER[] = "token.provider"; +const char CatalogOptions::DLF_REGION[] = "dlf.region"; +const char CatalogOptions::DLF_TOKEN_PATH[] = "dlf.token-path"; +const char CatalogOptions::DLF_ACCESS_KEY_ID[] = "dlf.access-key-id"; +const char CatalogOptions::DLF_ACCESS_KEY_SECRET[] = "dlf.access-key-secret"; +const char CatalogOptions::DLF_SECURITY_TOKEN[] = "dlf.security-token"; +const char CatalogOptions::DLF_TOKEN_LOADER[] = "dlf.token-loader"; +const char CatalogOptions::DLF_TOKEN_ECS_METADATA_URL[] = "dlf.token-ecs-metadata-url"; +const char CatalogOptions::DLF_TOKEN_ECS_ROLE_NAME[] = "dlf.token-ecs-role-name"; +const char CatalogOptions::DLF_SIGNING_ALGORITHM[] = "dlf.signing-algorithm"; const char CatalogOptions::TABLE_DEFAULT_OPTION_PREFIX[] = "table-default."; } // namespace paimon diff --git a/src/paimon/common/data/binary_array.h b/src/paimon/common/data/binary_array.h index c505ce7fa..8ab07e5f5 100644 --- a/src/paimon/common/data/binary_array.h +++ b/src/paimon/common/data/binary_array.h @@ -52,7 +52,7 @@ class MemorySegment; /// [size(int)] + [null bits(4-byte word boundaries)] + [values or offset&length] + [variable length /// part]. -class BinaryArray final : public BinarySection, public InternalArray { +class PAIMON_EXPORT BinaryArray final : public BinarySection, public InternalArray { public: BinaryArray() = default; diff --git a/src/paimon/common/data/binary_row.h b/src/paimon/common/data/binary_row.h index e3dfe8b07..7dda5492a 100644 --- a/src/paimon/common/data/binary_row.h +++ b/src/paimon/common/data/binary_row.h @@ -59,7 +59,7 @@ class MemoryPool; /// @note: Unlike the Java implementation where variable-length data may span multiple /// MemorySegments, in this C++ implementation both the fixed-length part and the /// variable-length part reside within a single MemorySegment. -class BinaryRow final : public BinarySection, public InternalRow, public DataSetters { +class PAIMON_EXPORT BinaryRow final : public BinarySection, public InternalRow, public DataSetters { public: BinaryRow() : BinaryRow(0) {} explicit BinaryRow(int32_t arity); diff --git a/src/paimon/common/data/columnar/columnar_row.h b/src/paimon/common/data/columnar/columnar_row.h index 2156c814d..f3d20dba0 100644 --- a/src/paimon/common/data/columnar/columnar_row.h +++ b/src/paimon/common/data/columnar/columnar_row.h @@ -78,6 +78,12 @@ class ColumnarRow : public InternalRow { row_kind_ = kind; } + /// Update the row represented by this view without rebuilding its column pointers. + /// @param row_id Zero-based row index in the underlying arrays. + void SetRowId(int64_t row_id) { + row_id_ = row_id; + } + int32_t GetFieldCount() const override { return array_vec_.size(); } diff --git a/src/paimon/common/data/columnar/columnar_row_test.cpp b/src/paimon/common/data/columnar/columnar_row_test.cpp index 6301cd293..d1a5c6003 100644 --- a/src/paimon/common/data/columnar/columnar_row_test.cpp +++ b/src/paimon/common/data/columnar/columnar_row_test.cpp @@ -70,6 +70,16 @@ TEST(ColumnarRowTest, TestSimple) { ASSERT_EQ(row.GetDouble(6), 5.5); ASSERT_EQ(row.GetString(7).ToString(), "Hello"); ASSERT_EQ(std::string(row.GetStringView(7)), "Hello"); + + row.SetRowId(3); + ASSERT_TRUE(row.GetBoolean(0)); + ASSERT_EQ(row.GetByte(1), 3); + ASSERT_EQ(row.GetShort(2), 7); + ASSERT_EQ(row.GetInt(3), 13); + ASSERT_EQ(row.GetLong(4), 18); + ASSERT_EQ(row.GetFloat(5), 3.3f); + ASSERT_EQ(row.GetDouble(6), 8.8); + ASSERT_EQ(std::string(row.GetStringView(7)), "WORLD"); } TEST(ColumnarRowRefTest, TestSimple) { diff --git a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.h b/src/paimon/common/data/shredding/map_shared_shredding_file_reader.h deleted file mode 100644 index 744d09fd1..000000000 --- a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * 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 - -#include "arrow/api.h" -#include "paimon/common/data/shredding/map_shared_shredding_utils.h" -#include "paimon/memory/memory_pool.h" -#include "paimon/reader/file_batch_reader.h" - -namespace paimon { - -class MapFieldReadPlan { - public: - virtual ~MapFieldReadPlan() = default; - - MapFieldReadPlan(const std::shared_ptr& logical_field, - const std::shared_ptr& physical_read_field) - : logical_field_(logical_field), physical_read_field_(physical_read_field) {} - - const std::shared_ptr& LogicalField() const { - return logical_field_; - } - - const std::shared_ptr& PhysicalReadField() const { - return physical_read_field_; - } - - virtual Result> Materialize( - const std::shared_ptr& physical_array, - arrow::MemoryPool* arrow_pool) const = 0; - - private: - std::shared_ptr logical_field_; - std::shared_ptr physical_read_field_; -}; - -class MapFieldReadPlanFactory { - public: - static Result> CreateMapReadPlan( - const std::shared_ptr& logical_map_field, - const MapSharedShreddingFieldMeta& meta); - - static Result> CreateSharedSelectedKeysReadPlan( - const std::shared_ptr& selected_keys_field, - const MapSharedShreddingFieldMeta& meta); - - static Result> CreateDefaultSelectedKeysReadPlan( - const std::shared_ptr& file_map_field, - const std::shared_ptr& selected_keys_field); -}; - -class MapSharedShreddingFileReader : public FileBatchReader { - public: - MapSharedShreddingFileReader( - std::unique_ptr&& reader, - std::map>&& field_read_plans, - const std::shared_ptr& pool); - - Result> GetFileSchema() const override; - - 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; - - void Close() override; - - Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; - - Result GetNumberOfRows() const override; - - bool SupportPreciseBitmapSelection() const override; - - private: - static Result> ToLogicalMapField( - const std::shared_ptr& physical_field); - - private: - std::shared_ptr arrow_pool_; - std::unique_ptr reader_; - std::map> field_read_plans_; -}; - -} // namespace paimon diff --git a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.cpp similarity index 63% rename from src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp rename to src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.cpp index b7c84095c..4cb28f694 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.cpp @@ -17,22 +17,16 @@ * under the License. */ -#include "paimon/common/data/shredding/map_shared_shredding_file_reader.h" +#include "paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h" -#include #include #include #include #include -#include "arrow/c/bridge.h" -#include "arrow/util/key_value_metadata.h" #include "fmt/format.h" -#include "paimon/common/reader/reader_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/core/casting/casting_utils.h" #include "paimon/core/utils/nested_projection_utils.h" namespace paimon { @@ -70,16 +64,35 @@ void CollectPhysicalColumns( } } -class FullMapReadPlan : public MapFieldReadPlan { +class MapShreddingColumnReadPlan : public ShreddingColumnReadPlan { + public: + MapShreddingColumnReadPlan(std::shared_ptr logical_field, + std::shared_ptr physical_field) + : logical_field_(std::move(logical_field)), physical_field_(std::move(physical_field)) {} + + const std::shared_ptr& LogicalField() const override { + return logical_field_; + } + + const std::shared_ptr& PhysicalField() const override { + return physical_field_; + } + + private: + std::shared_ptr logical_field_; + std::shared_ptr physical_field_; +}; + +class FullMapReadPlan : public MapShreddingColumnReadPlan { public: FullMapReadPlan(const std::shared_ptr& logical_field, const std::shared_ptr& physical_read_field, std::vector>&& selected_key_ids) - : MapFieldReadPlan(logical_field, physical_read_field), + : MapShreddingColumnReadPlan(logical_field, physical_read_field), selected_key_ids_(std::move(selected_key_ids)), logical_map_type_(checked_pointer_cast(logical_field->type())) {} - Result> Materialize( + Result> Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const override; @@ -88,7 +101,7 @@ class FullMapReadPlan : public MapFieldReadPlan { std::shared_ptr logical_map_type_; }; -class SharedSelectedKeysReadPlan : public MapFieldReadPlan { +class SharedSelectedKeysReadPlan : public MapShreddingColumnReadPlan { public: struct SelectedKey { int32_t field_id = -1; @@ -99,10 +112,10 @@ class SharedSelectedKeysReadPlan : public MapFieldReadPlan { SharedSelectedKeysReadPlan(const std::shared_ptr& logical_field, const std::shared_ptr& physical_read_field, std::vector&& selected_keys) - : MapFieldReadPlan(logical_field, physical_read_field), + : MapShreddingColumnReadPlan(logical_field, physical_read_field), selected_keys_(std::move(selected_keys)) {} - Result> Materialize( + Result> Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const override; @@ -110,14 +123,65 @@ class SharedSelectedKeysReadPlan : public MapFieldReadPlan { std::vector selected_keys_; }; -class DefaultSelectedKeysReadPlan : public MapFieldReadPlan { +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 MapShreddingColumnReadPlan { public: DefaultSelectedKeysReadPlan(const std::shared_ptr& logical_field, const std::shared_ptr& physical_read_field, const std::vector& selected_keys) - : MapFieldReadPlan(logical_field, physical_read_field), selected_keys_(selected_keys) {} + : MapShreddingColumnReadPlan(logical_field, physical_read_field), + selected_keys_(selected_keys) {} - Result> Materialize( + Result> Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const override; @@ -127,7 +191,8 @@ class DefaultSelectedKeysReadPlan : public MapFieldReadPlan { } // namespace -Result> MapFieldReadPlanFactory::CreateMapReadPlan( +Result> +MapSharedShreddingReadPlanFactory::CreateMapReadPlan( const std::shared_ptr& logical_map_field, const MapSharedShreddingFieldMeta& meta) { if (logical_map_field->type()->id() != arrow::Type::MAP) { @@ -164,12 +229,13 @@ Result> MapFieldReadPlanFactory::CreateMapRead logical_map_type->item_type(), selected_physical_column_ids, logical_map_type->item_field()->nullable(), include_overflow); auto physical_read_field = logical_map_field->WithType(physical_type); - std::unique_ptr read_plan = std::make_unique( + std::shared_ptr read_plan = std::make_shared( logical_map_field, physical_read_field, ResolveSelectedKeyIds(meta, selected_keys)); return read_plan; } -Result> MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan( +Result> +MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( const std::shared_ptr& selected_keys_field, const MapSharedShreddingFieldMeta& meta) { PAIMON_ASSIGN_OR_RAISE( @@ -204,13 +270,14 @@ Result> MapFieldReadPlanFactory::CreateSharedS value_field->type(), selected_physical_column_ids, value_field->nullable(), include_overflow); auto physical_read_field = selected_keys_field->WithType(physical_type); - std::unique_ptr read_plan = std::make_unique( - selected_keys_field, physical_read_field, std::move(selected_key_plans)); + std::shared_ptr read_plan = + std::make_shared(selected_keys_field, physical_read_field, + std::move(selected_key_plans)); return read_plan; } -Result> -MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( +Result> +MapSharedShreddingReadPlanFactory::CreateDefaultSelectedKeysReadPlan( const std::shared_ptr& file_map_field, const std::shared_ptr& selected_keys_field) { if (file_map_field->type()->id() != arrow::Type::MAP) { @@ -222,145 +289,13 @@ MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( std::vector selected_keys, NestedProjectionUtils::ValidateMapSharedShreddingAccessField(selected_keys_field)); auto physical_read_field = selected_keys_field->WithType(file_map_field->type()); - std::unique_ptr read_plan = std::make_unique( - selected_keys_field, physical_read_field, selected_keys); + std::shared_ptr read_plan = + std::make_shared(selected_keys_field, physical_read_field, + selected_keys); return read_plan; } -MapSharedShreddingFileReader::MapSharedShreddingFileReader( - std::unique_ptr&& reader, - std::map>&& field_read_plans, - const std::shared_ptr& pool) - : arrow_pool_(GetArrowPool(pool)), - reader_(std::move(reader)), - field_read_plans_(std::move(field_read_plans)) {} - -Result> MapSharedShreddingFileReader::GetFileSchema() const { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> physical_schema, - reader_->GetFileSchema()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr physical_arrow_schema, - arrow::ImportSchema(physical_schema.get())); - - arrow::FieldVector logical_fields = physical_arrow_schema->fields(); - for (int32_t i = 0; i < physical_arrow_schema->num_fields(); ++i) { - const auto& field = physical_arrow_schema->field(i); - std::shared_ptr metadata = - std::const_pointer_cast(field->metadata()); - if (!MapSharedShreddingUtils::HasShreddingMetadata(metadata)) { - continue; - } - PAIMON_ASSIGN_OR_RAISE(logical_fields[i], ToLogicalMapField(field)); - } - - auto logical_schema = arrow::schema(std::move(logical_fields)); - auto c_logical_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*logical_schema, c_logical_schema.get())); - return c_logical_schema; -} - -Result> MapSharedShreddingFileReader::ToLogicalMapField( - const std::shared_ptr& physical_field) { - if (!physical_field || !physical_field->type() || - physical_field->type()->id() != arrow::Type::STRUCT) { - return Status::Invalid(fmt::format("shared-shredding field {} is not a physical struct", - physical_field ? physical_field->name() : "")); - } - auto physical_type = checked_pointer_cast(physical_field->type()); - std::shared_ptr value_type; - bool value_nullable = true; - for (const auto& child : physical_type->fields()) { - if (child->name() == MapSharedShreddingDefine::kFieldMapping || - child->name() == MapSharedShreddingDefine::kOverflow) { - continue; - } - value_type = child->type(); - value_nullable = child->nullable(); - break; - } - if (!value_type) { - return Status::Invalid(fmt::format("cannot infer shared-shredding value type for field {}", - physical_field->name())); - } - return arrow::field( - physical_field->name(), - arrow::map(arrow::utf8(), arrow::field("value", value_type, value_nullable)), - physical_field->nullable()); -} - -Status MapSharedShreddingFileReader::SetReadSchema( - ::ArrowSchema* read_schema, const std::shared_ptr& predicate, - const std::optional& selection_bitmap) { - if (!read_schema) { - return Status::Invalid( - "invalid read schema in MapSharedShreddingFileReader, cannot be null"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_read_schema, - arrow::ImportSchema(read_schema)); - bool converted = false; - arrow::FieldVector physical_read_fields = logical_read_schema->fields(); - for (size_t i = 0; i < logical_read_schema->fields().size(); ++i) { - const auto& field = logical_read_schema->field(i); - auto plan_iter = field_read_plans_.find(field->name()); - if (plan_iter != field_read_plans_.end()) { - physical_read_fields[i] = plan_iter->second->PhysicalReadField(); - converted = true; - } - } - if (!converted) { - return Status::Invalid("suppose not fall into MapSharedShreddingFileReader"); - } - auto physical_read_schema = arrow::schema(std::move(physical_read_fields)); - std::unique_ptr c_physical_read_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportSchema(*physical_read_schema, c_physical_read_schema.get())); - return reader_->SetReadSchema(c_physical_read_schema.get(), predicate, selection_bitmap); -} - -Result MapSharedShreddingFileReader::NextBatch() { - return Status::Invalid( - "paimon inner reader MapSharedShreddingFileReader should use NextBatchWithBitmap"); -} - -Result MapSharedShreddingFileReader::NextBatchWithBitmap() { - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, - reader_->NextBatchWithBitmap()); - if (BatchReader::IsEofBatch(batch_with_bitmap)) { - return batch_with_bitmap; - } - - auto& [batch, bitmap] = batch_with_bitmap; - 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 batch to StructArray in MapSharedShreddingFileReader"); - } - auto struct_array = checked_pointer_cast(arrow_array); - - arrow::ArrayVector resolved_arrays = struct_array->fields(); - arrow::FieldVector resolved_fields = struct_array->struct_type()->fields(); - for (int32_t field_idx = 0; field_idx < struct_array->num_fields(); ++field_idx) { - const auto& physical_field = struct_array->struct_type()->field(field_idx); - auto plan_iter = field_read_plans_.find(physical_field->name()); - if (plan_iter == field_read_plans_.end()) { - continue; - } - PAIMON_ASSIGN_OR_RAISE( - resolved_arrays[field_idx], - plan_iter->second->Materialize(struct_array->field(field_idx), arrow_pool_.get())); - resolved_fields[field_idx] = plan_iter->second->LogicalField(); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr new_struct_array, - arrow::StructArray::Make(resolved_arrays, resolved_fields)); - auto new_c_array = std::make_unique(); - auto new_c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportArray(*new_struct_array, new_c_array.get(), new_c_schema.get())); - batch = std::make_pair(std::move(new_c_array), std::move(new_c_schema)); - return batch_with_bitmap; -} - -Result> FullMapReadPlan::Materialize( +Result> FullMapReadPlan::Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const { if (!physical_array || physical_array->type_id() != arrow::Type::STRUCT) { return Status::Invalid(fmt::format("cannot cast physical shredding field {} to StructArray", @@ -386,12 +321,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 +339,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, @@ -503,7 +433,7 @@ Result> FullMapReadPlan::Materialize( return map_array; } -Result> SharedSelectedKeysReadPlan::Materialize( +Result> SharedSelectedKeysReadPlan::Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const { if (!physical_array || physical_array->type_id() != arrow::Type::STRUCT) { return Status::Invalid(fmt::format("cannot cast physical shredding field {} to StructArray", @@ -530,12 +460,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 +477,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)); - } - } - - 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())); + PAIMON_ASSIGN_OR_RAISE(overflow_items, NestedProjectionUtils::AlignArrayToReadType( + overflow_items, value_type, arrow_pool)); } - 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,13 +581,28 @@ 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; } -Result> DefaultSelectedKeysReadPlan::Materialize( +Result> DefaultSelectedKeysReadPlan::Assemble( const std::shared_ptr& physical_array, arrow::MemoryPool* arrow_pool) const { if (!physical_array || physical_array->type_id() != arrow::Type::MAP) { return Status::Invalid( @@ -646,14 +612,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, @@ -698,25 +660,4 @@ Result> DefaultSelectedKeysReadPlan::Materialize( return result; } -std::shared_ptr MapSharedShreddingFileReader::GetReaderMetrics() const { - return reader_->GetReaderMetrics(); -} - -void MapSharedShreddingFileReader::Close() { - reader_->Close(); -} - -Result MapSharedShreddingFileReader::GetPreviousBatchFileRowId( - uint64_t batch_row_id) const { - return reader_->GetPreviousBatchFileRowId(batch_row_id); -} - -Result MapSharedShreddingFileReader::GetNumberOfRows() const { - return reader_->GetNumberOfRows(); -} - -bool MapSharedShreddingFileReader::SupportPreciseBitmapSelection() const { - return reader_->SupportPreciseBitmapSelection(); -} - } // namespace paimon diff --git a/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h new file mode 100644 index 000000000..871c51c38 --- /dev/null +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h @@ -0,0 +1,49 @@ +/* + * 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 "arrow/api.h" +#include "paimon/common/data/shredding/map_shared_shredding_utils.h" +#include "paimon/common/data/shredding/shredding_read_plan.h" + +namespace paimon { + +/// Builds per-column read plans for shared-shredding MAP columns and selected-key MAP access. +class MapSharedShreddingReadPlanFactory { + public: + MapSharedShreddingReadPlanFactory() = delete; + ~MapSharedShreddingReadPlanFactory() = delete; + + static Result> CreateMapReadPlan( + const std::shared_ptr& logical_map_field, + const MapSharedShreddingFieldMeta& meta); + + static Result> CreateSharedSelectedKeysReadPlan( + const std::shared_ptr& selected_keys_field, + const MapSharedShreddingFieldMeta& meta); + + static Result> CreateDefaultSelectedKeysReadPlan( + const std::shared_ptr& file_map_field, + const std::shared_ptr& selected_keys_field); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp similarity index 74% rename from src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp rename to src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp index 80f5046b9..e02f0ffb0 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp +++ b/src/paimon/common/data/shredding/map_shared_shredding_read_plan_factory_test.cpp @@ -17,7 +17,7 @@ * under the License. */ -#include "paimon/common/data/shredding/map_shared_shredding_file_reader.h" +#include "paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h" #include #include @@ -33,6 +33,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/data/shredding/map_shredding_defs.h" +#include "paimon/common/data/shredding/shredding_file_reader.h" #include "paimon/common/fs/external_path_provider.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/append/append_only_writer.h" @@ -51,7 +52,7 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { -class MapSharedShreddingFileReaderTest : public ::testing::Test { +class MapSharedShreddingReadPlanFactoryTest : public ::testing::Test { public: void SetUp() override { pool_ = GetDefaultPool(); @@ -100,12 +101,12 @@ class MapSharedShreddingFileReaderTest : public ::testing::Test { .ValueOrDie(); } - std::unique_ptr WrapReader( + std::unique_ptr WrapReader( std::unique_ptr&& reader, const std::optional& selected_keys_str = std::nullopt) const { EXPECT_OK_AND_ASSIGN(auto c_file_schema, reader->GetFileSchema()); auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); - std::map> field_read_plans; + std::map> field_read_plans; for (const auto& field : file_schema->fields()) { auto metadata = std::const_pointer_cast(field->metadata()); if (!MapSharedShreddingUtils::HasShreddingMetadata(metadata)) { @@ -124,20 +125,22 @@ class MapSharedShreddingFileReaderTest : public ::testing::Test { EXPECT_TRUE(item_field); auto map_type = checked_pointer_cast(arrow::map( arrow::utf8(), arrow::field("value", item_field->type(), item_field->nullable()))); - std::shared_ptr logical_map_field = field->WithType(map_type); + std::shared_ptr logical_map_field = + arrow::field(field->name(), map_type, field->nullable()); if (selected_keys_str.has_value()) { logical_map_field = logical_map_field->WithMetadata(arrow::KeyValueMetadata::Make( {DataField::MAP_SELECTED_KEYS}, {selected_keys_str.value()})); } - EXPECT_OK_AND_ASSIGN(auto field_read_plan, MapFieldReadPlanFactory::CreateMapReadPlan( - logical_map_field, meta)); + EXPECT_OK_AND_ASSIGN( + auto field_read_plan, + MapSharedShreddingReadPlanFactory::CreateMapReadPlan(logical_map_field, meta)); field_read_plans.emplace(field->name(), std::move(field_read_plan)); } - return std::make_unique(std::move(reader), - std::move(field_read_plans), pool_); + return std::make_unique(std::move(reader), std::move(field_read_plans), + pool_); } - Result> CreateReader( + Result> CreateReader( std::shared_ptr physical_array = nullptr, std::shared_ptr physical_schema = nullptr, const std::optional& selected_keys = std::nullopt) const { @@ -236,20 +239,7 @@ class MapSharedShreddingFileReaderTest : public ::testing::Test { }; }; -TEST_F(MapSharedShreddingFileReaderTest, TestGetFileSchemaReturnsLogicalMapSchema) { - ASSERT_OK_AND_ASSIGN(auto reader, CreateReader()); - - ASSERT_OK_AND_ASSIGN(auto c_schema, reader->GetFileSchema()); - auto schema = arrow::ImportSchema(c_schema.get()).ValueOrDie(); - - ASSERT_TRUE(schema->Equals(logical_schema_, /*check_metadata=*/false)) - << "Expected:\n" - << logical_schema_->ToString() << "\nActual:\n" - << schema->ToString(); - ASSERT_FALSE(schema->field(1)->HasMetadata()); -} - -TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithoutOverflow) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestAllExistSelectedKeysWithoutOverflow) { ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr, /*selected_keys=*/"b")); @@ -271,7 +261,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithoutOverflow AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestAllExistSelectedKeysWithOverflow) { ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr, /*selected_keys=*/"a,c")); @@ -293,7 +283,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjection) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjection) { ASSERT_OK_AND_ASSIGN(auto physical_schema, PhysicalSchemaWithMetadata()); ASSERT_OK_AND_ASSIGN(auto physical_array, PhysicalArray()); auto mock_reader = std::make_unique( @@ -306,13 +296,13 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjection) { auto selected_field = arrow::field( "tags", selected_type, /*nullable=*/true, arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c,missing"})); - ASSERT_OK_AND_ASSIGN( - auto field_read_plan, - MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan(selected_field, TagsMeta())); - std::map> contexts; + ASSERT_OK_AND_ASSIGN(auto field_read_plan, + MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( + selected_field, TagsMeta())); + std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); - auto reader = std::make_unique(std::move(mock_reader), - std::move(contexts), pool_); + auto reader = + std::make_unique(std::move(mock_reader), std::move(contexts), pool_); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); @@ -333,7 +323,98 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjection) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionFromDefaultMap) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, 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, + MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( + selected_field, TagsMeta())); + ASSERT_OK_AND_ASSIGN(auto result, + field_read_plan->Assemble(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(MapSharedShreddingReadPlanFactoryTest, + 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, + MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan(selected_field, meta)); + ASSERT_OK_AND_ASSIGN(auto result, + field_read_plan->Assemble(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(MapSharedShreddingReadPlanFactoryTest, TestSelectedKeysStructProjectionFromDefaultMap) { auto map_type = checked_pointer_cast( arrow::map(arrow::utf8(), arrow::field("value", arrow::int64()))); auto file_schema = @@ -356,12 +437,12 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionFromDef arrow::field("tags", selected_type, /*nullable=*/true, arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,missing"})); ASSERT_OK_AND_ASSIGN(auto field_read_plan, - MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( + MapSharedShreddingReadPlanFactory::CreateDefaultSelectedKeysReadPlan( file_schema->field(1), selected_field)); - std::map> contexts; + std::map> contexts; contexts.emplace("tags", std::move(field_read_plan)); - auto reader = std::make_unique(std::move(mock_reader), - std::move(contexts), pool_); + auto reader = + std::make_unique(std::move(mock_reader), std::move(contexts), pool_); auto read_schema = ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field})); @@ -381,15 +462,15 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionFromDef AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestInvalidSelectedKeysStructProjection) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestInvalidSelectedKeysStructProjection) { auto file_map_field = arrow::field("tags", arrow::map(arrow::utf8(), arrow::int64())); auto mismatched_count_field = arrow::field("tags", arrow::struct_({arrow::field("a", arrow::int64())}), /*nullable=*/true, arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,b"})); - ASSERT_NOK_WITH_MSG(MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan( + ASSERT_NOK_WITH_MSG(MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( mismatched_count_field, TagsMeta()), "metadata size 2 does not match STRUCT field count 1"); - ASSERT_NOK_WITH_MSG(MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( + ASSERT_NOK_WITH_MSG(MapSharedShreddingReadPlanFactory::CreateDefaultSelectedKeysReadPlan( file_map_field, mismatched_count_field), "metadata size 2 does not match STRUCT field count 1"); @@ -397,15 +478,15 @@ TEST_F(MapSharedShreddingFileReaderTest, TestInvalidSelectedKeysStructProjection "tags", arrow::struct_({arrow::field("a", arrow::int64()), arrow::field("b", arrow::utf8())}), /*nullable=*/true, arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,b"})); - ASSERT_NOK_WITH_MSG(MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan( + ASSERT_NOK_WITH_MSG(MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( mismatched_type_field, TagsMeta()), "must have the same value type"); - ASSERT_NOK_WITH_MSG(MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( + ASSERT_NOK_WITH_MSG(MapSharedShreddingReadPlanFactory::CreateDefaultSelectedKeysReadPlan( file_map_field, mismatched_type_field), "must have the same value type"); } -TEST_F(MapSharedShreddingFileReaderTest, TestPartialExistSelectedKeys) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestPartialExistSelectedKeys) { ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr, /*selected_keys=*/"a,c,missing")); @@ -428,7 +509,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestPartialExistSelectedKeys) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestMissingSelectedKeysReadsWholeMap) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestMissingSelectedKeysReadsWholeMap) { ASSERT_OK_AND_ASSIGN(auto reader, CreateReader()); auto read_schema = ExportSchema(ReadSchema(std::nullopt)); ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr, @@ -448,7 +529,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestMissingSelectedKeysReadsWholeMap) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestSpecialSelectedKeys) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestSpecialSelectedKeys) { MapSharedShreddingFieldMeta meta; meta.name_to_id = {{"", 0}, {" ", 1}, {".", 2}, {"a", 3}}; meta.field_to_columns = {{0, {0}}, {1, {1}}, {2, {0}}, {3, {1}}}; @@ -501,7 +582,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestSpecialSelectedKeys) { ])"); } -TEST_F(MapSharedShreddingFileReaderTest, TestUnknownSelectedKeyReturnsEmptyMap) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestUnknownSelectedKeyReturnsEmptyMap) { ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr, /*selected_keys=*/"missing")); @@ -524,7 +605,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestUnknownSelectedKeyReturnsEmptyMap) AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingField) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestInvalidNullFieldMappingField) { ASSERT_OK_AND_ASSIGN(auto physical_schema, PhysicalSchemaWithMetadata()); std::string json = R"([ [1, [null, 10, null, null]] @@ -541,7 +622,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingField) { "__field_mapping cannot be null"); } -TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingFieldElement) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestInvalidNullFieldMappingFieldElement) { ASSERT_OK_AND_ASSIGN(auto physical_schema, PhysicalSchemaWithMetadata()); std::string json = R"([ [1, [[0, null], 10, null, null]] @@ -558,7 +639,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingFieldElement "__field_mapping element cannot be null"); } -TEST_F(MapSharedShreddingFileReaderTest, TestListValue) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestListValue) { std::shared_ptr logical_schema = arrow::schema({ arrow::field("id", arrow::int32()), arrow::field("tags", arrow::map(arrow::utf8(), arrow::list(arrow::int32()))), @@ -614,7 +695,7 @@ TEST_F(MapSharedShreddingFileReaderTest, TestListValue) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, TestOrcDictionaryEncodedStringValue) { std::shared_ptr logical_schema = arrow::schema({ arrow::field("id", arrow::int32()), arrow::field("tags", arrow::map(arrow::utf8(), arrow::utf8())), @@ -674,7 +755,67 @@ TEST_F(MapSharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) { AssertChunkedArrayEquals(expected, actual); } -TEST_F(MapSharedShreddingFileReaderTest, TestReadsRealFormatFile) { +TEST_F(MapSharedShreddingReadPlanFactoryTest, 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(MapSharedShreddingReadPlanFactoryTest, TestReadsRealFormatFile) { // TODO(lisizhuo.lsz): support other format auto options = options_; std::string format = "orc"; diff --git a/src/paimon/common/data/shredding/shredding_file_reader.cpp b/src/paimon/common/data/shredding/shredding_file_reader.cpp index 0f47350dd..1c2e34f57 100644 --- a/src/paimon/common/data/shredding/shredding_file_reader.cpp +++ b/src/paimon/common/data/shredding/shredding_file_reader.cpp @@ -84,7 +84,7 @@ Result ShreddingFileReader::NextBatchWithBitma 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->type_id() != arrow::Type::STRUCT) { + if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { return Status::Invalid("cannot cast batch to StructArray in ShreddingFileReader"); } auto struct_array = checked_pointer_cast(arrow_array); diff --git a/src/paimon/common/data/variant/generic_variant_test.cpp b/src/paimon/common/data/variant/generic_variant_test.cpp index 11386a28a..bc46a5730 100644 --- a/src/paimon/common/data/variant/generic_variant_test.cpp +++ b/src/paimon/common/data/variant/generic_variant_test.cpp @@ -19,6 +19,7 @@ #include "paimon/common/data/variant/generic_variant.h" +#include #include #include #include @@ -27,6 +28,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/variant/variant_builder.h" #include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/utils/math.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" @@ -350,6 +352,23 @@ TEST_F(GenericVariantTest, NonFiniteDoubleToJson) { ASSERT_EQ(json, "\"Infinity\""); } +TEST_F(GenericVariantTest, CanonicalizesFloatingPointNaN) { + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendFloat(FloatingPointFromBits(0xffc12345U))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string_view value, variant->Value()); + ASSERT_EQ(ToHex(value), "380000c07f"); + } + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendDouble(FloatingPointFromBits(0xfff8123456789abcULL))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string_view value, variant->Value()); + ASSERT_EQ(ToHex(value), "1c000000000000f87f"); + } +} + TEST_F(GenericVariantTest, GetTypeInfoReturnsHeaderBits) { // GetTypeInfo exposes the primitive header's type-info bits; 42 is encoded as an int1. auto v = FromJson("42"); diff --git a/src/paimon/common/data/variant/variant_access_utils.cpp b/src/paimon/common/data/variant/variant_access_utils.cpp index 78180d1ce..0238c086a 100644 --- a/src/paimon/common/data/variant/variant_access_utils.cpp +++ b/src/paimon/common/data/variant/variant_access_utils.cpp @@ -26,6 +26,7 @@ #include "fmt/format.h" #include "paimon/common/data/variant/variant_defs.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/string_utils.h" namespace paimon { @@ -61,7 +62,7 @@ std::vector SplitDescription(const std::string& description) { } bool HasAccessDescription(const std::shared_ptr& field) { - return GetDescription(field).rfind(VariantAccessUtils::kMetadataKey, 0) == 0; + return StringUtils::StartsWith(GetDescription(field), VariantAccessUtils::kMetadataKey); } } // namespace diff --git a/src/paimon/common/data/variant/variant_builder.cpp b/src/paimon/common/data/variant/variant_builder.cpp index b3cf7ee5d..51f704175 100644 --- a/src/paimon/common/data/variant/variant_builder.cpp +++ b/src/paimon/common/data/variant/variant_builder.cpp @@ -30,6 +30,7 @@ #include "fmt/format.h" #include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/utils/math.h" #include "rapidjson/error/en.h" #include "rapidjson/memorystream.h" #include "rapidjson/reader.h" @@ -339,8 +340,7 @@ Status VariantBuilder::AppendLong(int64_t l) { Status VariantBuilder::AppendDouble(double d) { PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8)); write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDouble); - int64_t bits; - memcpy(&bits, &d, sizeof(bits)); + const int64_t bits = CanonicalizeDoubleToLongBits(d); VariantBinaryUtil::WriteLong(bits, 8, write_buffer_.data(), write_pos_); write_pos_ += 8; return Status::OK(); @@ -409,8 +409,7 @@ Status VariantBuilder::AppendTimestampNtz(int64_t micros_since_epoch) { Status VariantBuilder::AppendFloat(float f) { PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 4)); write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kFloat); - int32_t bits; - memcpy(&bits, &f, sizeof(bits)); + const int32_t bits = CanonicalizeFloatToIntBits(f); VariantBinaryUtil::WriteLong(bits, 4, write_buffer_.data(), write_pos_); write_pos_ += 4; return Status::OK(); diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index bac4f16f7..f0b0f5611 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -46,6 +46,7 @@ const char Options::PAGE_SIZE[] = "page-size"; const char Options::PARTITION_DEFAULT_NAME[] = "partition.default-name"; const char Options::FILE_COMPRESSION[] = "file.compression"; const char Options::FILE_COMPRESSION_ZSTD_LEVEL[] = "file.compression.zstd-level"; +const char Options::FILE_BLOCK_SIZE[] = "file.block-size"; const char Options::MANIFEST_TARGET_FILE_SIZE[] = "manifest.target-file-size"; const char Options::MANIFEST_FORMAT[] = "manifest.format"; const char Options::MANIFEST_COMPRESSION[] = "manifest.compression"; @@ -59,6 +60,9 @@ 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::PREFETCH_IO_METRICS_ENABLED[] = "prefetch.io-metrics.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"; @@ -92,6 +96,12 @@ const char Options::DELETION_VECTOR_INDEX_FILE_TARGET_SIZE[] = "deletion-vector.index-file.target-size"; const char Options::DELETION_VECTOR_BITMAP64[] = "deletion-vectors.bitmap64"; const char Options::CHANGELOG_PRODUCER[] = "changelog-producer"; +const char Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE[] = "changelog-producer.row-deduplicate"; +const char Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS[] = + "changelog-producer.row-deduplicate-ignore-fields"; +const char Options::CHANGELOG_FILE_PREFIX[] = "changelog-file.prefix"; +const char Options::CHANGELOG_FILE_FORMAT[] = "changelog-file.format"; +const char Options::CHANGELOG_FILE_COMPRESSION[] = "changelog-file.compression"; const char Options::FORCE_LOOKUP[] = "force-lookup"; const char Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_DELETE[] = "partial-update.remove-record-on-delete"; @@ -100,6 +110,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 +158,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/executor/default_executor_test.cpp b/src/paimon/common/executor/default_executor_test.cpp index 91f2c94c2..07e78681b 100644 --- a/src/paimon/common/executor/default_executor_test.cpp +++ b/src/paimon/common/executor/default_executor_test.cpp @@ -125,6 +125,57 @@ TEST(DefaultExecutorTest, TestAddTaskAfterShutdownNowIgnored) { ASSERT_EQ(executed_count.load(), 0); } +TEST(DefaultExecutorTest, TestConcurrentShutdownNow) { + constexpr int32_t kShutdownThreadCount = 2; + constexpr int32_t kAttempts = 50; + for (int32_t attempt = 0; attempt < kAttempts; ++attempt) { + ASSERT_OK_AND_ASSIGN(auto executor, CreateDefaultExecutor(/*thread_count=*/4)); + std::atomic ready_shutdown_count = 0; + std::promise start_signal; + std::shared_future start_future = start_signal.get_future().share(); + std::vector shutdown_threads; + shutdown_threads.reserve(kShutdownThreadCount); + + for (int32_t thread_index = 0; thread_index < kShutdownThreadCount; ++thread_index) { + shutdown_threads.emplace_back([&]() { + ++ready_shutdown_count; + start_future.wait(); + executor->ShutdownNow(); + }); + } + for (int32_t retry = 0; retry < 100 && ready_shutdown_count.load() < kShutdownThreadCount; + ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + const int32_t ready_count_before_start = ready_shutdown_count.load(); + start_signal.set_value(); + for (std::thread& shutdown_thread : shutdown_threads) { + shutdown_thread.join(); + } + ASSERT_EQ(kShutdownThreadCount, ready_count_before_start); + } +} + +TEST(DefaultExecutorTest, TestDestroyFromWorkerThread) { + std::unique_ptr created = CreateDefaultExecutor(); + std::shared_ptr executor(std::move(created)); + std::shared_ptr task_executor = executor; + auto release = std::make_shared>(); + std::shared_future release_future = release->get_future().share(); + auto destroyed = std::make_shared>(); + std::future future = destroyed->get_future(); + + executor->Add([executor = std::move(task_executor), release_future, destroyed]() mutable { + release_future.wait(); + executor.reset(); + destroyed->set_value(); + }); + + executor.reset(); + release->set_value(); + ASSERT_EQ(std::future_status::ready, future.wait_for(std::chrono::seconds(5))); +} + TEST(DefaultExecutorTest, TestAddTaskFromMultipleThreads) { ASSERT_OK_AND_ASSIGN(auto executor, CreateDefaultExecutor(/*thread_count=*/4)); diff --git a/src/paimon/common/executor/executor.cpp b/src/paimon/common/executor/executor.cpp index cd3f699d3..45253fe1f 100644 --- a/src/paimon/common/executor/executor.cpp +++ b/src/paimon/common/executor/executor.cpp @@ -40,23 +40,26 @@ class DefaultExecutor : public Executor { uint32_t GetThreadNum() const override; private: - void WorkerThread(); + struct State { + std::queue> tasks; + std::mutex mutex; + std::condition_variable condition; + bool stop = false; + }; + + static void WorkerThread(std::shared_ptr state); void ShutdownInternal(bool wait_for_pending_tasks); private: uint32_t thread_count_; std::vector workers_; - std::queue> tasks_; - std::mutex queue_mutex_; - std::condition_variable condition_; - bool stop_ = false; - int32_t active_tasks_ = 0; + std::shared_ptr state_ = std::make_shared(); }; DefaultExecutor::DefaultExecutor(uint32_t thread_count) : thread_count_(thread_count) { assert(thread_count > 0); for (uint32_t i = 0; i < thread_count_; ++i) { - workers_.emplace_back(&DefaultExecutor::WorkerThread, this); + workers_.emplace_back(&DefaultExecutor::WorkerThread, state_); } } @@ -66,21 +69,25 @@ uint32_t DefaultExecutor::GetThreadNum() const { void DefaultExecutor::ShutdownInternal(bool wait_for_pending_tasks) { { - std::unique_lock lock(queue_mutex_); - if (stop_) { + std::unique_lock lock(state_->mutex); + if (state_->stop) { return; } - stop_ = true; + state_->stop = true; if (!wait_for_pending_tasks) { // Discard all pending tasks immediately. std::queue> empty; - tasks_.swap(empty); + state_->tasks.swap(empty); } - condition_.notify_all(); + state_->condition.notify_all(); } for (std::thread& worker : workers_) { if (worker.joinable()) { - worker.join(); + if (worker.get_id() == std::this_thread::get_id()) { + worker.detach(); + } else { + worker.join(); + } } } } @@ -100,38 +107,32 @@ void DefaultExecutor::Add(std::function func) { return; } { - std::unique_lock lock(queue_mutex_); - if (stop_) { + std::unique_lock lock(state_->mutex); + if (state_->stop) { return; } - tasks_.emplace(std::move(func)); + state_->tasks.emplace(std::move(func)); } - condition_.notify_one(); + state_->condition.notify_one(); } -void DefaultExecutor::WorkerThread() { +void DefaultExecutor::WorkerThread(std::shared_ptr state) { while (true) { std::function task; { - std::unique_lock lock(queue_mutex_); - condition_.wait(lock, [this] { return stop_ || !tasks_.empty(); }); - if (stop_ && tasks_.empty() && active_tasks_ == 0) { - condition_.notify_all(); + std::unique_lock lock(state->mutex); + state->condition.wait(lock, [&state] { return state->stop || !state->tasks.empty(); }); + if (state->stop && state->tasks.empty()) { + state->condition.notify_all(); return; } - if (!tasks_.empty()) { - task = std::move(tasks_.front()); - tasks_.pop(); - ++active_tasks_; + if (!state->tasks.empty()) { + task = std::move(state->tasks.front()); + state->tasks.pop(); } } if (task) { task(); - std::unique_lock lock(queue_mutex_); - --active_tasks_; - if (tasks_.empty() && active_tasks_ == 0) { - condition_.notify_all(); - } } } } diff --git a/src/paimon/common/factories/io_hook.cpp b/src/paimon/common/factories/io_hook.cpp index a0576b469..394dc8ea5 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 e0a2f68be..0c66381b9 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 9bbb1b342..653dc73b8 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 a97322597..2a5daa894 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 000000000..efaf921cc --- /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 d068d3c2b..ea45c9791 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,8 @@ 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(), + /*enable_io_metrics=*/false, pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/common/file_index/bitmap/bitmap_file_index.cpp b/src/paimon/common/file_index/bitmap/bitmap_file_index.cpp index 6df161436..43bb71b0e 100644 --- a/src/paimon/common/file_index/bitmap/bitmap_file_index.cpp +++ b/src/paimon/common/file_index/bitmap/bitmap_file_index.cpp @@ -103,24 +103,18 @@ Result> BitmapFileIndex::CreateWriter( "invalid schema for BitmapFileIndexWriter, supposed to have single " "field."); } - auto arrow_field = arrow_schema->field(0); - return BitmapFileIndexWriter::Create(arrow_schema, arrow_field->name(), options_, pool); + return BitmapFileIndexWriter::Create(arrow_schema->field(0), options_, pool); } Result> BitmapFileIndexWriter::Create( - const std::shared_ptr& arrow_schema, const std::string& field_name, - const std::map& options, const std::shared_ptr& pool) { + const std::shared_ptr& field, const std::map& options, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(int8_t version, OptionsUtils::GetValueFromMap(options, BitmapFileIndex::VERSION, BitmapFileIndex::VERSION_2)); - auto arrow_field = arrow_schema->GetFieldByName(field_name); - if (!arrow_field) { - return Status::Invalid( - fmt::format("field {} not in arrow_schema for BitmapFileIndexWriter", field_name)); - } - auto struct_type = arrow::struct_({arrow_field}); + std::shared_ptr struct_type = arrow::struct_({field}); return std::shared_ptr( - new BitmapFileIndexWriter(version, struct_type, arrow_field->type(), options, pool)); + new BitmapFileIndexWriter(version, struct_type, field->type(), options, pool)); } BitmapFileIndexWriter::BitmapFileIndexWriter(int8_t version, @@ -137,15 +131,7 @@ BitmapFileIndexWriter::BitmapFileIndexWriter(int8_t version, Status BitmapFileIndexWriter::AddBatch(::ArrowArray* batch) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, arrow::ImportArray(batch, struct_type_)); - if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("invalid batch for BitmapFileIndexWriter, expected a struct array"); - } auto struct_array = checked_pointer_cast(arrow_array); - if (struct_array->num_fields() != 1) { - return Status::Invalid( - "invalid batch for BitmapFileIndexWriter, expected a struct array with exactly one " - "field"); - } PAIMON_ASSIGN_OR_RAISE( std::vector array_values, LiteralConverter::ConvertLiteralsFromArray(*(struct_array->field(0)), /*own_data=*/true)); diff --git a/src/paimon/common/file_index/bitmap/bitmap_file_index.h b/src/paimon/common/file_index/bitmap/bitmap_file_index.h index 9cf15ddc3..6aa0f4169 100644 --- a/src/paimon/common/file_index/bitmap/bitmap_file_index.h +++ b/src/paimon/common/file_index/bitmap/bitmap_file_index.h @@ -71,7 +71,7 @@ class PAIMON_EXPORT BitmapFileIndex : public FileIndexer { class BitmapFileIndexWriter : public FileIndexWriter { public: static Result> Create( - const std::shared_ptr& arrow_schema, const std::string& field_name, + const std::shared_ptr& field, const std::map& options, const std::shared_ptr& pool); Status AddBatch(::ArrowArray* batch) override; diff --git a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.cpp b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.cpp index ca33c696f..86d07d63b 100644 --- a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.cpp +++ b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.cpp @@ -19,10 +19,19 @@ #include "paimon/common/file_index/bloomfilter/bloom_filter_file_index.h" #include +#include +#include #include #include +#include +#include "arrow/c/bridge.h" #include "fmt/format.h" +#include "paimon/common/predicate/literal_converter.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/math.h" +#include "paimon/common/utils/options_utils.h" #include "paimon/fs/file_system.h" #include "paimon/memory/bytes.h" #include "paimon/predicate/literal.h" @@ -31,7 +40,8 @@ namespace paimon { class MemoryPool; -BloomFilterFileIndex::BloomFilterFileIndex(const std::map& options) {} +BloomFilterFileIndex::BloomFilterFileIndex(const std::map& options) + : options_(options) {} Result> BloomFilterFileIndex::CreateReader( ::ArrowSchema* c_arrow_schema, int32_t start, int32_t length, const std::shared_ptr& input_stream, @@ -58,15 +68,77 @@ Result> BloomFilterFileIndex::CreateReader( return BloomFilterFileIndexReader::Create(arrow_type, bytes); } +Result> BloomFilterFileIndex::CreateWriter( + ::ArrowSchema* c_arrow_schema, const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, + arrow::ImportSchema(c_arrow_schema)); + if (arrow_schema->num_fields() != 1) { + return Status::Invalid( + "invalid schema for BloomFilterFileIndexWriter, supposed to have single field."); + } + return BloomFilterFileIndexWriter::Create(arrow_schema->field(0), options_, pool); +} + +Result> BloomFilterFileIndexWriter::Create( + const std::shared_ptr& field, const std::map& options, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(FastHash::HashFunction hash_function, + FastHash::GetHashFunction(field->type())); + PAIMON_ASSIGN_OR_RAISE( + int32_t items, OptionsUtils::GetValueFromMap(options, BloomFilterFileIndex::kItems, + BloomFilterFileIndex::kDefaultItems)); + PAIMON_ASSIGN_OR_RAISE( + double fpp, OptionsUtils::GetValueFromMap(options, BloomFilterFileIndex::kFpp, + BloomFilterFileIndex::kDefaultFpp)); + std::shared_ptr struct_type = arrow::struct_({field}); + PAIMON_ASSIGN_OR_RAISE(BloomFilter64 filter, BloomFilter64::Create(items, fpp, pool)); + return std::shared_ptr( + new BloomFilterFileIndexWriter(struct_type, hash_function, std::move(filter), pool)); +} + +BloomFilterFileIndexWriter::BloomFilterFileIndexWriter( + const std::shared_ptr& struct_type, + const FastHash::HashFunction& hash_function, BloomFilter64&& filter, + const std::shared_ptr& pool) + : struct_type_(struct_type), + hash_function_(hash_function), + filter_(std::move(filter)), + pool_(pool) {} + +Status BloomFilterFileIndexWriter::AddBatch(::ArrowArray* batch) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(batch, struct_type_)); + std::shared_ptr struct_array = + checked_pointer_cast(array); + PAIMON_ASSIGN_OR_RAISE( + std::vector values, + LiteralConverter::ConvertLiteralsFromArray(*struct_array->field(0), /*own_data=*/false)); + for (const Literal& value : values) { + if (!value.IsNull()) { + filter_.AddHash(hash_function_(value)); + } + } + return Status::OK(); +} + +Result> BloomFilterFileIndexWriter::SerializedBytes() const { + constexpr int32_t kHeaderLength = sizeof(int32_t); + const int32_t bit_set_length = filter_.GetBitSet().ByteLength(); + PAIMON_UNIQUE_PTR bytes = + Bytes::AllocateBytes(kHeaderLength + bit_set_length, pool_.get()); + const int32_t num_hash_functions = ToBigEndian(filter_.GetNumHashFunctions()); + std::memcpy(bytes->data(), &num_hash_functions, sizeof(num_hash_functions)); + filter_.GetBitSet().ToByteArray(kHeaderLength, bit_set_length, bytes->data()); + return bytes; +} + Result> BloomFilterFileIndexReader::Create( const std::shared_ptr& arrow_type, const std::shared_ptr& bytes) { - // compatible with java, little endian - const char* data = bytes->data(); - auto num_hash_functions = - static_cast((static_cast(static_cast(data[0])) << 24) | - (static_cast(static_cast(data[1])) << 16) | - (static_cast(static_cast(data[2])) << 8) | - static_cast(static_cast(data[3]))); + // Compatible with Java's big-endian numHashFunctions header. + int32_t big_endian_num_hash_functions; + std::memcpy(&big_endian_num_hash_functions, bytes->data(), + sizeof(big_endian_num_hash_functions)); + const int32_t num_hash_functions = FromBigEndian(big_endian_num_hash_functions); PAIMON_ASSIGN_OR_RAISE(FastHash::HashFunction hash_function, FastHash::GetHashFunction(arrow_type)); auto bit_set = std::make_unique(bytes, /*offset=*/sizeof(int32_t)); diff --git a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.h b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.h index 8152c2204..62d082c76 100644 --- a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.h +++ b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index.h @@ -25,10 +25,10 @@ #include "arrow/c/bridge.h" #include "paimon/common/file_index/bloomfilter/fast_hash.h" -#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/bloom_filter64.h" #include "paimon/file_index/file_index_reader.h" #include "paimon/file_index/file_index_result.h" +#include "paimon/file_index/file_index_writer.h" #include "paimon/file_index/file_indexer.h" #include "paimon/result.h" namespace paimon { @@ -55,11 +55,37 @@ class BloomFilterFileIndex : public FileIndexer { const std::shared_ptr& pool) const override; Result> CreateWriter( - ::ArrowSchema* arrow_schema, const std::shared_ptr& pool) const override { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_type, - arrow::ImportType(arrow_schema)); - return Status::NotImplemented("do not support index writer in bloom filter"); - } + ::ArrowSchema* arrow_schema, const std::shared_ptr& pool) const override; + + static constexpr int32_t kDefaultItems = 1000000; + static constexpr double kDefaultFpp = 0.1; + static constexpr char kItems[] = "items"; + static constexpr char kFpp[] = "fpp"; + + private: + std::map options_; +}; + +class BloomFilterFileIndexWriter : public FileIndexWriter { + public: + static Result> Create( + const std::shared_ptr& field, + const std::map& options, const std::shared_ptr& pool); + + Status AddBatch(::ArrowArray* batch) override; + + Result> SerializedBytes() const override; + + private: + BloomFilterFileIndexWriter(const std::shared_ptr& struct_type, + const FastHash::HashFunction& hash_function, BloomFilter64&& filter, + const std::shared_ptr& pool); + + private: + std::shared_ptr struct_type_; + FastHash::HashFunction hash_function_; + BloomFilter64 filter_; + std::shared_ptr pool_; }; class BloomFilterFileIndexReader : public FileIndexReader { diff --git a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index_test.cpp b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index_test.cpp index 209f0b409..23fadce28 100644 --- a/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index_test.cpp +++ b/src/paimon/common/file_index/bloomfilter/bloom_filter_file_index_test.cpp @@ -18,10 +18,16 @@ #include "paimon/common/file_index/bloomfilter/bloom_filter_file_index.h" +#include +#include +#include #include #include +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" @@ -49,21 +55,60 @@ class BloomFilterIndexReaderTest : public ::testing::Test { return c_schema; } + Result> WriteIndex( + const std::shared_ptr& data_type, const std::string& json, + const std::map& options) const { + const std::shared_ptr schema = + arrow::schema({arrow::field("f0", data_type)}); + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); + BloomFilterFileIndex file_index(options); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + file_index.CreateWriter(&c_schema, pool_)); + ::ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array)); + return writer->SerializedBytes(); + } + private: std::shared_ptr pool_; }; +TEST_F(BloomFilterIndexReaderTest, TestWriterRejectsInvalidOptionsAndType) { + auto create_writer = [&](const std::shared_ptr& type, + const std::map& options) { + BloomFilterFileIndex file_index(options); + return file_index.CreateWriter(CreateArrowSchema(type).get(), pool_); + }; + ASSERT_NOK_WITH_MSG(create_writer(arrow::int32(), {{"items", "0"}}), + "items must be greater than 0"); + ASSERT_NOK_WITH_MSG(create_writer(arrow::int32(), {{"fpp", "1"}}), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(create_writer(arrow::boolean(), {}), + "bloom filter index does not support BOOLEAN"); +} + TEST_F(BloomFilterIndexReaderTest, TestStringType) { - // data: "a", "b", "" - std::vector index_bytes = {0, 0, 0, 6, 0, 32, 32, 3, 208, 32, 0, 64, 73, 16, 201}; - auto input_stream = std::make_shared( - reinterpret_cast(index_bytes.data()), index_bytes.size()); + // Java writer output for data: "a", "b", "" with items=10 and fpp=0.02. + const std::vector expected = {0, 0, 0, 6, 0, 32, 32, 3, 208, 32, 0, 64, 73, 16, 201}; + ASSERT_OK_AND_ASSIGN( + PAIMON_UNIQUE_PTR bytes, + WriteIndex(arrow::utf8(), R"([["a"], ["b"], [""]])", {{"items", "10"}, {"fpp", "0.02"}})); + ASSERT_EQ(expected.size(), bytes->size()); + ASSERT_EQ(0, std::memcmp(expected.data(), bytes->data(), expected.size())); + + std::shared_ptr input_stream = + std::make_shared(bytes->data(), bytes->size()); BloomFilterFileIndex file_index({}); ASSERT_OK_AND_ASSIGN( auto reader, file_index.CreateReader(CreateArrowSchema(arrow::utf8()).get(), - /*start=*/0, /*length=*/index_bytes.size(), input_stream, pool_)); + /*start=*/0, /*length=*/bytes->size(), input_stream, pool_)); ASSERT_TRUE(reader); ASSERT_TRUE(reader->VisitEqual(Literal(FieldType::STRING, "a", 1)).value()->IsRemain().value()); ASSERT_TRUE(reader->VisitEqual(Literal(FieldType::STRING, "b", 1)).value()->IsRemain().value()); diff --git a/src/paimon/common/file_index/bloomfilter/fast_hash.cpp b/src/paimon/common/file_index/bloomfilter/fast_hash.cpp index d98dc476e..ef63f0fc1 100644 --- a/src/paimon/common/file_index/bloomfilter/fast_hash.cpp +++ b/src/paimon/common/file_index/bloomfilter/fast_hash.cpp @@ -27,6 +27,7 @@ #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/file_index/file_index_result.h" @@ -58,17 +59,11 @@ Result FastHash::GetHashFunction( }); case FieldType::FLOAT: return HashFunction([](const Literal& literal) -> int64_t { - auto raw_value = literal.GetValue(); - int32_t bits = 0; - std::memcpy(&bits, &raw_value, sizeof(raw_value)); - return GetLongHash(bits); + return GetLongHash(CanonicalizeFloatToIntBits(literal.GetValue())); }); case FieldType::DOUBLE: return HashFunction([](const Literal& literal) -> int64_t { - auto raw_value = literal.GetValue(); - int64_t bits; - std::memcpy(&bits, &raw_value, sizeof(raw_value)); - return GetLongHash(bits); + return GetLongHash(CanonicalizeDoubleToLongBits(literal.GetValue())); }); case FieldType::TIMESTAMP: { auto ts_type = checked_pointer_cast(arrow_type); diff --git a/src/paimon/common/file_index/bloomfilter/fast_hash_test.cpp b/src/paimon/common/file_index/bloomfilter/fast_hash_test.cpp index 680d591bd..6ffbc7964 100644 --- a/src/paimon/common/file_index/bloomfilter/fast_hash_test.cpp +++ b/src/paimon/common/file_index/bloomfilter/fast_hash_test.cpp @@ -18,11 +18,14 @@ #include "paimon/common/file_index/bloomfilter/fast_hash.h" +#include +#include #include #include #include #include "gtest/gtest.h" +#include "paimon/common/utils/math.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/file_index/file_index_result.h" @@ -164,4 +167,22 @@ TEST_F(FastHashTest, TestCompatibleWithJava) { } } +TEST_F(FastHashTest, TestNaNCompatibleWithJava) { + const auto float_nan = FloatingPointFromBits(0x7fc12345U); + const auto negative_float_nan = FloatingPointFromBits(0xffc54321U); + ASSERT_TRUE(std::isnan(float_nan)); + ASSERT_TRUE(std::isnan(negative_float_nan)); + ASSERT_OK_AND_ASSIGN(auto float_hash_function, FastHash::GetHashFunction(arrow::float32())); + CheckResult(float_hash_function, {Literal(float_nan), Literal(negative_float_nan)}, + {0x67c27c6d9936ae63, 0x67c27c6d9936ae63}); + + const auto double_nan = FloatingPointFromBits(0x7ff8123456789abcULL); + const auto negative_double_nan = FloatingPointFromBits(0xfff8abcdef012345ULL); + ASSERT_TRUE(std::isnan(double_nan)); + ASSERT_TRUE(std::isnan(negative_double_nan)); + ASSERT_OK_AND_ASSIGN(auto double_hash_function, FastHash::GetHashFunction(arrow::float64())); + CheckResult(double_hash_function, {Literal(double_nan), Literal(negative_double_nan)}, + {0x13d2d3f2cc0e846e, 0x13d2d3f2cc0e846e}); +} + } // namespace paimon::test diff --git a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.cpp b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.cpp index 7ce53e9a6..4cb98cdca 100644 --- a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.cpp +++ b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.cpp @@ -18,13 +18,23 @@ #include "paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h" +#include #include #include #include #include +#include +#include +#include +#include +#include "arrow/c/bridge.h" #include "fmt/format.h" #include "paimon/common/file_index/bsi/bit_slice_index_roaring_bitmap.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/predicate/literal_converter.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/field_type_utils.h" @@ -53,6 +63,101 @@ class MemoryPool; BitSliceIndexBitmapFileIndex::BitSliceIndexBitmapFileIndex( const std::map& options) {} +Result> BitSliceIndexBitmapFileIndex::CreateWriter( + ::ArrowSchema* c_arrow_schema, const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, + arrow::ImportSchema(c_arrow_schema)); + if (arrow_schema->num_fields() != 1) { + return Status::Invalid( + "invalid schema for BitSliceIndexBitmapFileIndexWriter, supposed to have single " + "field."); + } + const std::shared_ptr field = arrow_schema->field(0); + PAIMON_ASSIGN_OR_RAISE(ValueMapperType value_mapper, GetValueMapper(field->type())); + return std::make_shared(field, value_mapper, pool); +} + +BitSliceIndexBitmapFileIndexWriter::BitSliceIndexBitmapFileIndexWriter( + const std::shared_ptr& field, + const BitSliceIndexBitmapFileIndex::ValueMapperType& value_mapper, + const std::shared_ptr& pool) + : struct_type_(arrow::struct_({field})), + field_name_(field->name()), + value_mapper_(value_mapper), + pool_(pool) {} + +Status BitSliceIndexBitmapFileIndexWriter::AddBatch(::ArrowArray* batch) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(batch, struct_type_)); + auto struct_array = checked_pointer_cast(array); + if (struct_array->length() > static_cast(std::numeric_limits::max()) - + static_cast(values_.size())) { + return Status::Invalid("bsi index row count exceeds the supported int32 range"); + } + PAIMON_ASSIGN_OR_RAISE( + std::vector literals, + LiteralConverter::ConvertLiteralsFromArray(*struct_array->field(0), /*own_data=*/false)); + values_.reserve(values_.size() + literals.size()); + for (const Literal& literal : literals) { + if (literal.IsNull()) { + values_.emplace_back(std::nullopt); + continue; + } + PAIMON_ASSIGN_OR_RAISE(int64_t value, value_mapper_(literal)); + if (value == std::numeric_limits::min()) { + return Status::Invalid( + fmt::format("bsi index does not support INT64_MIN for field '{}'", field_name_)); + } + values_.emplace_back(value); + if (value < 0) { + const int64_t absolute_value = SafeAbs(value); + negative_min_ = std::min(negative_min_, absolute_value); + negative_max_ = std::max(negative_max_, absolute_value); + } else { + positive_min_ = std::min(positive_min_, value); + positive_max_ = std::max(positive_max_, value); + } + } + return Status::OK(); +} + +Result> BitSliceIndexBitmapFileIndexWriter::SerializedBytes() const { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr positive, + BitSliceIndexRoaringBitmap::Appender::Create(positive_min_, positive_max_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr negative, + BitSliceIndexRoaringBitmap::Appender::Create(negative_min_, negative_max_)); + for (size_t i = 0; i < values_.size(); ++i) { + if (!values_[i]) { + continue; + } + const int64_t value = values_[i].value(); + if (value < 0) { + PAIMON_RETURN_NOT_OK(negative->Append(static_cast(i), SafeAbs(value))); + } else { + PAIMON_RETURN_NOT_OK(positive->Append(static_cast(i), value)); + } + } + + MemorySegmentOutputStream output_stream(MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_); + output_stream.SetOrder(ByteOrder::PAIMON_BIG_ENDIAN); + output_stream.WriteValue(BitSliceIndexBitmapFileIndex::VERSION_1); + output_stream.WriteValue(static_cast(values_.size())); + const bool has_positive = positive->IsNotEmpty(); + output_stream.WriteValue(has_positive); + if (has_positive) { + output_stream.WriteBytes(positive->Serialize(pool_)); + } + const bool has_negative = negative->IsNotEmpty(); + output_stream.WriteValue(has_negative); + if (has_negative) { + output_stream.WriteBytes(negative->Serialize(pool_)); + } + return MemorySegmentUtils::CopyToBytes(output_stream.Segments(), /*offset=*/0, + /*num_bytes=*/output_stream.CurrentSize(), pool_.get()); +} + Result> BitSliceIndexBitmapFileIndex::CreateReader( ::ArrowSchema* c_arrow_schema, int32_t start, int32_t length, const std::shared_ptr& input_stream, diff --git a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h index e661ea2a3..d7dfe092f 100644 --- a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h +++ b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h @@ -22,15 +22,15 @@ #include #include #include +#include #include #include #include -#include "arrow/c/bridge.h" #include "arrow/type.h" -#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/file_index/file_index_reader.h" #include "paimon/file_index/file_index_result.h" +#include "paimon/file_index/file_index_writer.h" #include "paimon/file_index/file_indexer.h" #include "paimon/predicate/literal.h" #include "paimon/result.h" @@ -54,14 +54,12 @@ class BitSliceIndexBitmapFileIndex : public FileIndexer { const std::shared_ptr& pool) const override; Result> CreateWriter( - ::ArrowSchema* arrow_schema, const std::shared_ptr& pool) const override { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_type, - arrow::ImportType(arrow_schema)); - return Status::NotImplemented("do not support index writer in bsi"); - } + ::ArrowSchema* arrow_schema, const std::shared_ptr& pool) const override; using ValueMapperType = std::function(const Literal& literal)>; + static constexpr int8_t VERSION_1 = 1; + private: static Result GetValueMapper( const std::shared_ptr& arrow_type); @@ -74,9 +72,29 @@ class BitSliceIndexBitmapFileIndex : public FileIndexer { } return static_cast(literal.GetValue()); } +}; + +class BitSliceIndexBitmapFileIndexWriter : public FileIndexWriter { + public: + BitSliceIndexBitmapFileIndexWriter( + const std::shared_ptr& field, + const BitSliceIndexBitmapFileIndex::ValueMapperType& value_mapper, + const std::shared_ptr& pool); + + Status AddBatch(::ArrowArray* batch) override; + + Result> SerializedBytes() const override; private: - static constexpr int8_t VERSION_1 = 1; + std::shared_ptr struct_type_; + std::string field_name_; + BitSliceIndexBitmapFileIndex::ValueMapperType value_mapper_; + std::vector> values_; + int64_t positive_min_ = 0; + int64_t positive_max_ = 0; + int64_t negative_min_ = 0; + int64_t negative_max_ = 0; + std::shared_ptr pool_; }; class BitSliceIndexBitmapFileIndexReader diff --git a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp index 128c1a69c..760434e19 100644 --- a/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp +++ b/src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp @@ -19,9 +19,13 @@ #include "paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.h" #include +#include #include +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" @@ -60,6 +64,24 @@ class BitSliceIndexBitmapIndexReaderTest : public ::testing::Test { << ", expected=" << RoaringBitmap32::From(expected).ToString(); } + Result> WriteIndex(const std::shared_ptr& data_type, + const std::string& json) const { + const std::shared_ptr schema = + arrow::schema({arrow::field("f0", data_type)}); + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); + BitSliceIndexBitmapFileIndex file_index({}); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + file_index.CreateWriter(&c_schema, pool_)); + ::ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array)); + return writer->SerializedBytes(); + } + private: std::shared_ptr pool_; }; @@ -75,48 +97,75 @@ TEST_F(BitSliceIndexBitmapIndexReaderTest, TestMix) { 0, 0, 0, 2, 58, 48, 0, 0, 1, 0, 0, 0, 0, 0, 2, 0, 16, 0, 0, 0, 3, 0, 4, 0, 5, 0, 0, 0, 0, 2, 58, 48, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 16, 0, 0, 0, 5, 0, 58, 48, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 16, 0, 0, 0, 3, 0, 4, 0}; + const auto check_mix_reader = [this](const std::shared_ptr& reader) { + // test equal + CheckResult(reader->VisitEqual(Literal(2)).value(), {1, 7}); + CheckResult(reader->VisitEqual(Literal(-2)).value(), {3, 4}); + CheckResult(reader->VisitEqual(Literal(100)).value(), {}); + + // test not equal + CheckResult(reader->VisitNotEqual(Literal(2)).value(), {0, 3, 4, 5, 8, 9}); + CheckResult(reader->VisitNotEqual(Literal(-2)).value(), {0, 1, 5, 7, 8, 9}); + CheckResult(reader->VisitNotEqual(Literal(100)).value(), {0, 1, 3, 4, 5, 7, 8, 9}); + + // test in + CheckResult(reader->VisitIn({Literal(-1), Literal(1), Literal(2), Literal(3)}).value(), + {0, 1, 5, 7}); + + // test not in + CheckResult(reader->VisitNotIn({Literal(-1), Literal(1), Literal(2), Literal(3)}).value(), + {3, 4, 8, 9}); + + // test null + CheckResult(reader->VisitIsNull().value(), {2, 6, 10}); + + // test not null + CheckResult(reader->VisitIsNotNull().value(), {0, 1, 3, 4, 5, 7, 8, 9}); + + // test less than + CheckResult(reader->VisitLessThan(Literal(2)).value(), {0, 3, 4, 5, 8}); + CheckResult(reader->VisitLessOrEqual(Literal(2)).value(), {0, 1, 3, 4, 5, 7, 8}); + CheckResult(reader->VisitLessThan(Literal(-1)).value(), {3, 4}); + CheckResult(reader->VisitLessOrEqual(Literal(-1)).value(), {3, 4, 5}); + + // test greater than + CheckResult(reader->VisitGreaterThan(Literal(-2)).value(), {0, 1, 5, 7, 8, 9}); + CheckResult(reader->VisitGreaterOrEqual(Literal(-2)).value(), {0, 1, 3, 4, 5, 7, 8, 9}); + CheckResult(reader->VisitGreaterThan(Literal(2)).value(), {9}); + CheckResult(reader->VisitGreaterOrEqual(Literal(2)).value(), {1, 7, 9}); + }; + + ASSERT_OK_AND_ASSIGN( + PAIMON_UNIQUE_PTR written_bytes, + WriteIndex(arrow::int32(), + R"([[1], [2], [null], [-2], [-2], [-1], [null], [2], [0], [5], [null]])")); + auto written_stream = + std::make_shared(written_bytes->data(), written_bytes->size()); + BitSliceIndexBitmapFileIndex file_index({}); + ASSERT_OK_AND_ASSIGN(auto written_reader, + file_index.CreateReader(CreateArrowSchema(arrow::int32()).get(), + /*start=*/0, /*length=*/written_bytes->size(), + written_stream, pool_)); + check_mix_reader(written_reader); + + // Reading the Java-produced fixture below. C++ and Java may choose different valid portable + // roaring containers for the same bitmap, so their complete byte streams need not be identical. auto input_stream = std::make_shared(index_bytes.data(), index_bytes.size()); - BitSliceIndexBitmapFileIndex file_index({}); ASSERT_OK_AND_ASSIGN( - auto reader, + auto java_bytes_reader, file_index.CreateReader(CreateArrowSchema(arrow::int32()).get(), /*start=*/0, /*length=*/index_bytes.size(), input_stream, pool_)); - // test equal - CheckResult(reader->VisitEqual(Literal(2)).value(), {1, 7}); - CheckResult(reader->VisitEqual(Literal(-2)).value(), {3, 4}); - CheckResult(reader->VisitEqual(Literal(100)).value(), {}); - - // test not equal - CheckResult(reader->VisitNotEqual(Literal(2)).value(), {0, 3, 4, 5, 8, 9}); - CheckResult(reader->VisitNotEqual(Literal(-2)).value(), {0, 1, 5, 7, 8, 9}); - CheckResult(reader->VisitNotEqual(Literal(100)).value(), {0, 1, 3, 4, 5, 7, 8, 9}); - - // test in - CheckResult(reader->VisitIn({Literal(-1), Literal(1), Literal(2), Literal(3)}).value(), - {0, 1, 5, 7}); - - // test not in - CheckResult(reader->VisitNotIn({Literal(-1), Literal(1), Literal(2), Literal(3)}).value(), - {3, 4, 8, 9}); - - // test null - CheckResult(reader->VisitIsNull().value(), {2, 6, 10}); - - // test not null - CheckResult(reader->VisitIsNotNull().value(), {0, 1, 3, 4, 5, 7, 8, 9}); + check_mix_reader(java_bytes_reader); +} - // test less than - CheckResult(reader->VisitLessThan(Literal(2)).value(), {0, 3, 4, 5, 8}); - CheckResult(reader->VisitLessOrEqual(Literal(2)).value(), {0, 1, 3, 4, 5, 7, 8}); - CheckResult(reader->VisitLessThan(Literal(-1)).value(), {3, 4}); - CheckResult(reader->VisitLessOrEqual(Literal(-1)).value(), {3, 4, 5}); +TEST_F(BitSliceIndexBitmapIndexReaderTest, TestWriterRejectsInt64MinAndUnsupportedType) { + ASSERT_NOK_WITH_MSG(WriteIndex(arrow::int64(), R"([[-9223372036854775808]])"), + "bsi index does not support INT64_MIN for field 'f0'"); - // test greater than - CheckResult(reader->VisitGreaterThan(Literal(-2)).value(), {0, 1, 5, 7, 8, 9}); - CheckResult(reader->VisitGreaterOrEqual(Literal(-2)).value(), {0, 1, 3, 4, 5, 7, 8, 9}); - CheckResult(reader->VisitGreaterThan(Literal(2)).value(), {9}); - CheckResult(reader->VisitGreaterOrEqual(Literal(2)).value(), {1, 7, 9}); + BitSliceIndexBitmapFileIndex file_index({}); + ASSERT_NOK_WITH_MSG(file_index.CreateWriter(CreateArrowSchema(arrow::boolean()).get(), pool_), + "BitSliceIndexBitmapFileIndex only support"); } TEST_F(BitSliceIndexBitmapIndexReaderTest, TestPositiveOnly) { diff --git a/src/paimon/common/file_index/file_index_format.cpp b/src/paimon/common/file_index/file_index_format.cpp index 855008450..fab5c7a7b 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 7851d57e1..40989f7e9 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/file_index/rangebitmap/range_bitmap_file_index.cpp b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.cpp index 5d0543bef..88e081bd1 100644 --- a/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.cpp +++ b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.cpp @@ -60,18 +60,12 @@ Result> RangeBitmapFileIndex::CreateWriter( return Status::Invalid( "invalid schema for RangeBitmapFileIndexWriter, supposed to have single field."); } - const auto arrow_field = arrow_schema_ptr->field(0); - return RangeBitmapFileIndexWriter::Create(arrow_schema_ptr, arrow_field->name(), options_, - pool); + return RangeBitmapFileIndexWriter::Create(arrow_schema_ptr->field(0), options_, pool); } Result> RangeBitmapFileIndexWriter::Create( - const std::shared_ptr& arrow_schema, const std::string& field_name, - const std::map& options, const std::shared_ptr& pool) { - const auto field = arrow_schema->GetFieldByName(field_name); - if (!field) { - return Status::Invalid(fmt::format("Field not found in schema: {}", field_name)); - } + const std::shared_ptr& field, const std::map& options, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(FieldType field_type, FieldTypeUtils::ConvertToFieldType(field->type()->id())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr shared_key_factory, @@ -82,27 +76,18 @@ Result> RangeBitmapFileIndexWriter:: chunk_size_it != options.end()) { PAIMON_ASSIGN_OR_RAISE(parsed_chunk_size, MemorySize::ParseBytes(chunk_size_it->second)); } - auto struct_type = arrow::struct_({field}); + std::shared_ptr struct_type = arrow::struct_({field}); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr appender_ptr, RangeBitmap::Appender::Create(shared_key_factory, parsed_chunk_size, pool)); - return std::make_shared( - struct_type, field->type(), options, pool, shared_key_factory, std::move(appender_ptr)); + return std::make_shared(struct_type, pool, shared_key_factory, + std::move(appender_ptr)); } Status RangeBitmapFileIndexWriter::AddBatch(::ArrowArray* batch) { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(batch, struct_type_)); - if (!array || array->type_id() != arrow::Type::STRUCT) { - return Status::Invalid( - "invalid batch for RangeBitmapFileIndexWriter, expected a struct array"); - } auto struct_array = checked_pointer_cast(array); - if (struct_array->num_fields() != 1) { - return Status::Invalid( - "invalid batch for RangeBitmapFileIndexWriter, expected a struct array with exactly " - "one field"); - } PAIMON_ASSIGN_OR_RAISE(std::vector array_values, LiteralConverter::ConvertLiteralsFromArray(*(struct_array->field(0)), /*own_data=*/true)); @@ -117,13 +102,9 @@ Result> RangeBitmapFileIndexWriter::SerializedBytes() c } RangeBitmapFileIndexWriter::RangeBitmapFileIndexWriter( - const std::shared_ptr& struct_type, - const std::shared_ptr& arrow_type, - const std::map& options, const std::shared_ptr& pool, + const std::shared_ptr& struct_type, const std::shared_ptr& pool, const std::shared_ptr& key_factory, std::unique_ptr appender) : struct_type_(struct_type), - arrow_type_(arrow_type), - options_(options), pool_(pool), key_factory_(key_factory), appender_(std::move(appender)) {} diff --git a/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.h b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.h index fc99cac33..64289a536 100644 --- a/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.h +++ b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index.h @@ -61,15 +61,13 @@ class PAIMON_EXPORT RangeBitmapFileIndex final : public FileIndexer { class RangeBitmapFileIndexWriter final : public FileIndexWriter { public: static Result> Create( - const std::shared_ptr& arrow_schema, const std::string& field_name, + const std::shared_ptr& field, const std::map& options, const std::shared_ptr& pool); Status AddBatch(::ArrowArray* batch) override; Result> SerializedBytes() const override; RangeBitmapFileIndexWriter(const std::shared_ptr& struct_type, - const std::shared_ptr& arrow_type, - const std::map& options, const std::shared_ptr& pool, const std::shared_ptr& key_factory, std::unique_ptr appender); @@ -78,8 +76,6 @@ class RangeBitmapFileIndexWriter final : public FileIndexWriter { /// @note struct_type_ contains only one field with arrow_type_, used for import from C /// interface. std::shared_ptr struct_type_; - std::shared_ptr arrow_type_; - std::map options_; std::shared_ptr pool_; std::shared_ptr key_factory_; std::unique_ptr appender_; diff --git a/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index_test.cpp b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index_test.cpp index e14797849..6157ee020 100644 --- a/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index_test.cpp +++ b/src/paimon/common/file_index/rangebitmap/range_bitmap_file_index_test.cpp @@ -123,17 +123,15 @@ Result> RangeBitmapFileIndexTest::Cr std::shared_ptr arrow_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Finish(&arrow_array)); // Wrap in StructArray (single field) as required by RangeBitmapFileIndexWriter - arrow::FieldVector fields = {arrow::field("test_field", arrow_type)}; + auto field = arrow::field("test_field", arrow_type); + arrow::FieldVector fields = {field}; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr struct_array, arrow::StructArray::Make({arrow_array}, fields)); auto c_array = std::make_unique<::ArrowArray>(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, c_array.get())); - // Create schema for the field - const auto schema = arrow::schema({arrow::field("test_field", arrow_type)}); // Create writer - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr writer, - RangeBitmapFileIndexWriter::Create(schema, "test_field", options, pool_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + RangeBitmapFileIndexWriter::Create(field, options, pool_)); // Add the batch PAIMON_RETURN_NOT_OK(writer->AddBatch(c_array.get())); // Get serialized payload diff --git a/src/paimon/common/file_index/rangebitmap/range_bitmap_io_test.cpp b/src/paimon/common/file_index/rangebitmap/range_bitmap_io_test.cpp index e2bb867ed..fd8e0c888 100644 --- a/src/paimon/common/file_index/rangebitmap/range_bitmap_io_test.cpp +++ b/src/paimon/common/file_index/rangebitmap/range_bitmap_io_test.cpp @@ -65,18 +65,16 @@ class RangeBitmapIoTest : public ::testing::Test { PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Finish(&arrow_array)); // Wrap in StructArray - arrow::FieldVector fields = {arrow::field("test_field", arrow_type)}; + const std::shared_ptr field = arrow::field("test_field", arrow_type); + arrow::FieldVector fields = {field}; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr struct_array, arrow::StructArray::Make({arrow_array}, fields)); auto c_array = std::make_unique<::ArrowArray>(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, c_array.get())); - // Create schema - const auto schema = arrow::schema({arrow::field("test_field", arrow_type)}); - // Create writer and write data PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, - RangeBitmapFileIndexWriter::Create(schema, "test_field", {}, pool_)); + RangeBitmapFileIndexWriter::Create(field, {}, pool_)); PAIMON_RETURN_NOT_OK(writer->AddBatch(c_array.get())); return writer->SerializedBytes(); diff --git a/src/paimon/common/fs/object_store_file_system.cpp b/src/paimon/common/fs/object_store_file_system.cpp index ae44835e2..26cbe2aa4 100644 --- a/src/paimon/common/fs/object_store_file_system.cpp +++ b/src/paimon/common/fs/object_store_file_system.cpp @@ -30,6 +30,7 @@ #include "paimon/common/utils/path_util.h" namespace paimon { + namespace { constexpr int64_t kMaxReadAheadMemory = 64LL * 1024LL * 1024LL; diff --git a/src/paimon/common/fs/object_store_file_system.h b/src/paimon/common/fs/object_store_file_system.h index a1313b851..a40ef15c3 100644 --- a/src/paimon/common/fs/object_store_file_system.h +++ b/src/paimon/common/fs/object_store_file_system.h @@ -19,7 +19,12 @@ #pragma once +#include + +#include +#include #include +#include #include #include #include @@ -31,6 +36,71 @@ namespace paimon { +class ObjectStoreFileSystemUtils { + public: + static inline int64_t ParseModificationTime(const std::string& value) { + std::tm parsed_time{}; + const char* current = strptime(value.c_str(), "%Y-%m-%dT%H:%M:%S", &parsed_time); + if (current != nullptr) { + int32_t milliseconds = 0; + int32_t fraction_digits = 0; + if (*current == '.') { + ++current; + const char* fraction_begin = current; + while (std::isdigit(static_cast(*current))) { + if (fraction_digits < 3) { + milliseconds = milliseconds * 10 + (*current - '0'); + } + ++fraction_digits; + ++current; + } + if (current == fraction_begin) { + current = nullptr; + } + while (fraction_digits < 3) { + milliseconds *= 10; + ++fraction_digits; + } + } + + int32_t timezone_offset_seconds = 0; + bool valid_timezone = false; + if (current != nullptr && *current == 'Z' && current[1] == '\0') { + valid_timezone = true; + } else if (current != nullptr && (*current == '+' || *current == '-') && + std::isdigit(static_cast(current[1])) && + std::isdigit(static_cast(current[2])) && current[3] == ':' && + std::isdigit(static_cast(current[4])) && + std::isdigit(static_cast(current[5])) && current[6] == '\0') { + int32_t hours = (current[1] - '0') * 10 + current[2] - '0'; + int32_t minutes = (current[4] - '0') * 10 + current[5] - '0'; + if (hours <= 23 && minutes <= 59) { + timezone_offset_seconds = (hours * 60 + minutes) * 60; + if (*current == '-') { + timezone_offset_seconds = -timezone_offset_seconds; + } + valid_timezone = true; + } + } + + if (valid_timezone) { + errno = 0; + time_t seconds = timegm(&parsed_time); + if (seconds != static_cast(-1) || errno != EOVERFLOW) { + int64_t utc_seconds = static_cast(seconds) - timezone_offset_seconds; + return utc_seconds * 1000 + milliseconds; + } + } + } + + time_t seconds = curl_getdate(value.c_str(), nullptr); + if (seconds == static_cast(-1)) { + return FileStatus::kUnknownModificationTime; + } + return static_cast(seconds) * 1000; + } +}; + struct ObjectStorePath { std::string bucket; std::string key; @@ -39,7 +109,7 @@ struct ObjectStorePath { struct ObjectMetadata { std::string key; int64_t size = 0; - int64_t modification_time = 0; + int64_t modification_time = FileStatus::kUnknownModificationTime; }; struct ListObjectsResult { diff --git a/src/paimon/common/fs/object_store_file_system_test.cpp b/src/paimon/common/fs/object_store_file_system_test.cpp index 824060f80..6b8b56503 100644 --- a/src/paimon/common/fs/object_store_file_system_test.cpp +++ b/src/paimon/common/fs/object_store_file_system_test.cpp @@ -33,6 +33,17 @@ namespace paimon::test { namespace { +TEST(ObjectStoreFileSystemTest, TestParseModificationTime) { + ASSERT_EQ(1704067200000, + ObjectStoreFileSystemUtils::ParseModificationTime("2024-01-01T00:00:00.000Z")); + ASSERT_EQ(1704067200123, + ObjectStoreFileSystemUtils::ParseModificationTime("2024-01-01T08:00:00.123+08:00")); + ASSERT_EQ(1704067200000, + ObjectStoreFileSystemUtils::ParseModificationTime("Mon, 01 Jan 2024 00:00:00 GMT")); + ASSERT_EQ(FileStatus::kUnknownModificationTime, + ObjectStoreFileSystemUtils::ParseModificationTime("not-a-timestamp")); +} + using Range = std::pair; class MockObjectStoreClient : public ObjectStoreClient { diff --git a/src/paimon/common/global_index/btree/key_serializer.cpp b/src/paimon/common/global_index/btree/key_serializer.cpp index 464f37ea4..4b66d1d1f 100644 --- a/src/paimon/common/global_index/btree/key_serializer.cpp +++ b/src/paimon/common/global_index/btree/key_serializer.cpp @@ -27,6 +27,7 @@ #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/math.h" #include "paimon/common/utils/preconditions.h" #include "paimon/common/utils/var_length_int_utils.h" #include "paimon/data/decimal.h" @@ -164,19 +165,13 @@ Result> KeySerializer::SerializeKey( case FieldType::FLOAT: { MemorySliceOutput output(4, pool); output.Reset(); - auto fvalue = literal.GetValue(); - int32_t ivalue; - memcpy(&ivalue, &fvalue, sizeof(float)); - output.WriteValue(ivalue); + output.WriteValue(CanonicalizeFloatToIntBits(literal.GetValue())); return output.ToSlice().CopyBytes(pool); } case FieldType::DOUBLE: { MemorySliceOutput output(8, pool); output.Reset(); - auto dvalue = literal.GetValue(); - int64_t ivalue; - memcpy(&ivalue, &dvalue, sizeof(double)); - output.WriteValue(ivalue); + output.WriteValue(CanonicalizeDoubleToLongBits(literal.GetValue())); return output.ToSlice().CopyBytes(pool); } case FieldType::STRING: { diff --git a/src/paimon/common/global_index/btree/key_serializer_test.cpp b/src/paimon/common/global_index/btree/key_serializer_test.cpp index e36e72ed6..61322fde4 100644 --- a/src/paimon/common/global_index/btree/key_serializer_test.cpp +++ b/src/paimon/common/global_index/btree/key_serializer_test.cpp @@ -19,7 +19,11 @@ #include "paimon/common/global_index/btree/key_serializer.h" +#include +#include + #include "gtest/gtest.h" +#include "paimon/common/utils/math.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" #include "paimon/testing/utils/testharness.h" @@ -208,6 +212,30 @@ TEST_F(KeySerializerTest, SerializeAndDeserializeAllTypes) { } } +TEST_F(KeySerializerTest, CanonicalizesFloatingPointNaN) { + const auto float_nan = FloatingPointFromBits(0xffc12345U); + const auto canonical_float_nan = FloatingPointFromBits(kCanonicalFloatNaNBits); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr float_bytes, + KeySerializer::SerializeKey(Literal(float_nan), arrow::float32(), pool_.get())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr canonical_float_bytes, + KeySerializer::SerializeKey(Literal(canonical_float_nan), arrow::float32(), pool_.get())); + ASSERT_EQ(std::string(float_bytes->data(), float_bytes->size()), + std::string(canonical_float_bytes->data(), canonical_float_bytes->size())); + + const auto double_nan = FloatingPointFromBits(0xfff8123456789abcULL); + const auto canonical_double_nan = FloatingPointFromBits(kCanonicalDoubleNaNBits); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr double_bytes, + KeySerializer::SerializeKey(Literal(double_nan), arrow::float64(), pool_.get())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr canonical_double_bytes, + KeySerializer::SerializeKey(Literal(canonical_double_nan), arrow::float64(), pool_.get())); + ASSERT_EQ(std::string(double_bytes->data(), double_bytes->size()), + std::string(canonical_double_bytes->data(), canonical_double_bytes->size())); +} + TEST_F(KeySerializerTest, RejectsMalformedSerializedKeys) { auto wrap = [this](const std::string& value) { return MemorySlice::Wrap(std::make_shared(value, pool_.get())); diff --git a/src/paimon/common/global_index/global_index_result.cpp b/src/paimon/common/global_index/global_index_result.cpp index f329b0361..d2b94f216 100644 --- a/src/paimon/common/global_index/global_index_result.cpp +++ b/src/paimon/common/global_index/global_index_result.cpp @@ -22,6 +22,7 @@ #include "fmt/format.h" #include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/global_index/bitmap_scored_global_index_result.h" #include "paimon/io/byte_array_input_stream.h" @@ -37,8 +38,8 @@ void WriteBitmapAndScores(const RoaringBitmap64* bitmap, const std::vectorWriteBytes(bitmap_bytes); out->WriteValue(scores.size()); - for (auto score : scores) { - out->WriteValue(score); + for (float score : scores) { + out->WriteValue(CanonicalizeFloatingPoint(score)); } } diff --git a/src/paimon/common/global_index/global_index_result_test.cpp b/src/paimon/common/global_index/global_index_result_test.cpp index 73c6d05ed..3179e6710 100644 --- a/src/paimon/common/global_index/global_index_result_test.cpp +++ b/src/paimon/common/global_index/global_index_result_test.cpp @@ -19,9 +19,12 @@ #include "paimon/global_index/global_index_result.h" +#include +#include #include #include "gtest/gtest.h" +#include "paimon/common/utils/math.h" #include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/global_index/bitmap_scored_global_index_result.h" #include "paimon/testing/utils/testharness.h" @@ -144,6 +147,29 @@ TEST_F(GlobalIndexResultTest, TestSerializeAndDeserializeWithScore) { serialize_bytes->data() + serialize_bytes->size())); } +TEST_F(GlobalIndexResultTest, TestSerializeCanonicalizesNaNScore) { + auto pool = GetDefaultPool(); + const auto payload_nan = FloatingPointFromBits(0xffc12345U); + const auto canonical_nan = FloatingPointFromBits(kCanonicalFloatNaNBits); + auto index_result = std::make_shared( + RoaringBitmap64::From({1}), std::vector{payload_nan}); + auto canonical_index_result = std::make_shared( + RoaringBitmap64::From({1}), std::vector{canonical_nan}); + + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR serialized, + GlobalIndexResult::Serialize(index_result, pool)); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR canonical_serialized, + GlobalIndexResult::Serialize(canonical_index_result, pool)); + ASSERT_EQ(*serialized, *canonical_serialized); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr deserialized, + GlobalIndexResult::Deserialize(serialized->data(), serialized->size(), pool)); + auto scored_result = std::dynamic_pointer_cast(deserialized); + ASSERT_TRUE(scored_result); + ASSERT_TRUE(std::isnan(scored_result->GetScores()[0])); +} + TEST_F(GlobalIndexResultTest, TestInvalidSerialize) { auto pool = GetDefaultPool(); auto result = std::make_shared(std::vector({1, 3, 5, 100})); 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 000000000..bc6d1f1ab --- /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 000000000..9b87ca429 --- /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 000000000..bd185095d --- /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 f899d46ce..6fafefb5c 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 000000000..4d7c180a3 --- /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 9ccbf2608..015b082c3 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 d4a61854f..0b14dc33b 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 4e6063706..0a5dd5748 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 5355f72ba..2d0d274a7 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 61fbfe305..69c207ce9 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/lookup/lookup_store_factory.cpp b/src/paimon/common/lookup/lookup_store_factory.cpp index c5742e12a..d6cb8d400 100644 --- a/src/paimon/common/lookup/lookup_store_factory.cpp +++ b/src/paimon/common/lookup/lookup_store_factory.cpp @@ -36,7 +36,8 @@ Result> LookupStoreFactory::BfGenerator(int64_t row if (row_count <= 0 || !options.LookupCacheBloomFilterEnabled()) { return std::shared_ptr(); } - auto bloom_filter = BloomFilter::Create(row_count, options.GetLookupCacheBloomFilterFpp()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bloom_filter, + BloomFilter::Create(row_count, options.GetLookupCacheBloomFilterFpp())); MemorySegment memory_segment = MemorySegment::AllocateHeapMemory(bloom_filter->ByteLength(), pool); PAIMON_RETURN_NOT_OK(bloom_filter->SetMemorySegment(memory_segment)); diff --git a/src/paimon/common/predicate/literal.cpp b/src/paimon/common/predicate/literal.cpp index 3b2bcc0e6..d679c2ccf 100644 --- a/src/paimon/common/predicate/literal.cpp +++ b/src/paimon/common/predicate/literal.cpp @@ -18,7 +18,6 @@ #include "paimon/predicate/literal.h" -#include #include #include #include @@ -29,6 +28,7 @@ #include "fmt/format.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/math.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" #include "paimon/status.h" @@ -63,9 +63,9 @@ class Literal::Impl { case FieldType::BIGINT: return std::hash{}(value_.BigIntVal); case FieldType::FLOAT: - return std::hash{}(value_.FloatVal); + return std::hash{}(CanonicalizeFloatingPoint(value_.FloatVal)); case FieldType::DOUBLE: - return std::hash{}(value_.DoubleVal); + return std::hash{}(CanonicalizeFloatingPoint(value_.DoubleVal)); case FieldType::STRING: case FieldType::BINARY: return std::hash{}(std::string_view(value_.Buffer, size_)); diff --git a/src/paimon/common/predicate/literal_converter.cpp b/src/paimon/common/predicate/literal_converter.cpp index 194107687..102348d70 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/blob_fallback_batch_reader.cpp b/src/paimon/common/reader/blob_fallback_batch_reader.cpp index 6ca3106b3..21db0ade6 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader.cpp @@ -31,6 +31,7 @@ #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/table/special_fields.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" @@ -238,11 +239,9 @@ Result> BlobFallbackBatchReader::AssembleRowIdRun( break; } } - if (pieces.size() == 1 && pieces[0]->offset() == 0) { + if (pieces.size() == 1) { return pieces[0]; } - // Concatenate flattens non-zero offsets left by Slice, so the exported batch honors the - // zero-offset BatchReader contract. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr concat_array, arrow::Concatenate(pieces, arrow_pool_.get())); return concat_array; @@ -303,11 +302,9 @@ Result> BlobFallbackBatchReader::AssembleColumn( } run_start = run_end; } - if (pieces.size() == 1 && pieces[0]->offset() == 0) { + if (pieces.size() == 1) { return pieces[0]; } - // Concatenate flattens non-zero offsets left by Slice, so the exported batch honors the - // zero-offset BatchReader contract. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr concat_array, arrow::Concatenate(pieces, arrow_pool_.get())); return concat_array; @@ -360,12 +357,14 @@ Result BlobFallbackBatchReader::NextBatch() { AssembleColumn(field_idx, group_choice, group_chunks)); columns.push_back(std::move(column)); } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr target_array, + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::StructArray::Make(columns, read_schema_->fields())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_array, + ArrowUtils::NormalizeArrayOffsets(array, arrow_pool_.get())); std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportArray(*target_array, c_array.get(), c_schema.get())); + arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); return std::make_pair(std::move(c_array), std::move(c_schema)); } diff --git a/src/paimon/common/reader/delegating_prefetch_reader.h b/src/paimon/common/reader/delegating_prefetch_reader.h index 3cbcb08bd..78396e200 100644 --- a/src/paimon/common/reader/delegating_prefetch_reader.h +++ b/src/paimon/common/reader/delegating_prefetch_reader.h @@ -45,7 +45,7 @@ class DelegatingPrefetchReader : public FileBatchReader { } std::shared_ptr GetReaderMetrics() const override { - return GetReader()->GetReaderMetrics(); + return prefetch_reader_->GetReaderMetrics(); } Result> GetFileSchema() const override { diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp new file mode 100644 index 000000000..ed7fa304d --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -0,0 +1,369 @@ +/* + * 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/reader/late_materializing_file_batch_reader.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" +#include "arrow/c/bridge.h" +#include "arrow/memory_pool.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/predicate/predicate_validator.h" +#include "paimon/common/reader/reader_utils.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/predicate/predicate_utils.h" +#include "paimon/status.h" + +namespace paimon { + +Result> LateMaterializingFileBatchReader::Create( + std::unique_ptr inner, std::shared_ptr pool) { + // The reader's own compaction allocations go through an arrow pool; bridge the paimon pool + // once here so the accounting matches the rest of the read path. + if (pool == nullptr) { + return Status::Invalid("pool could not be nullptr."); + } + if (inner == nullptr) { + return Status::Invalid("inner could not be nullptr."); + } + auto* prefetch_inner = dynamic_cast(inner.get()); + std::shared_ptr arrow_pool = GetArrowPool(pool); + auto reader = + std::unique_ptr(new LateMaterializingFileBatchReader( + std::move(inner), prefetch_inner, std::move(arrow_pool))); + return reader; +} + +Result LateMaterializingFileBatchReader::NextBatch() { + if (state_ == kInit) { + // SetReadSchema has not been called: read with the file schema, matching the + // FileBatchReader contract for schema-less reads. + state_ = kNoLatMat; + } + if (state_ == kProbing) { + PAIMON_RETURN_NOT_OK(ReadAndFilterProbeData()); + if (matched_bitmap_.IsEmpty()) { + state_ = kEOF; + } else { + // payload pass reads only the matched rows (matched_bitmap_ is non-empty here). + PAIMON_RETURN_NOT_OK( + SetInnerReadSchema(payload_schema_, /*predicate=*/nullptr, matched_bitmap_)); + state_ = kRunning; + } + } + + if (state_ == kNoLatMat) { + return inner_->NextBatch(); + } else if (state_ == kRunning) { + return ReadPayloadBatch(); + } else if (state_ == kEOF) { + return MakeEofBatch(); + } + return Status::Invalid("invalid state when calling NextBatch: " + std::to_string(state_)); +} + +Result LateMaterializingFileBatchReader::FilterProbeBatch( + const std::shared_ptr& array, + const std::shared_ptr& bound_filter) { + // TODO(zhouhonfeng.zhf): use arrow::compute::Filter instead of PredicateFilter + PAIMON_ASSIGN_OR_RAISE(std::vector results, bound_filter->Test(*array)); + if (results.size() != static_cast(array->length())) { + return Status::Invalid( + fmt::format("predicate result size {} does not match probe batch length {}", + results.size(), array->length())); + } + // batch-local offsets of the rows passing both the predicate and the selection + RoaringBitmap32 batch_matched; + for (int64_t i = 0; i < array->length(); ++i) { + if (!results[static_cast(i)]) { + continue; + } + // map batch offset to file row id + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + inner_->GetPreviousBatchFileRowId(static_cast(i))); + if (selection_ && !selection_->Contains(static_cast(file_row))) { + continue; + } + batch_matched.Add(static_cast(i)); + matched_bitmap_.Add(file_row); + } + return batch_matched; +} + +Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { + matched_bitmap_ = RoaringBitmap32(); + probe_cursor_ = 0; + arrow::ArrayVector probe_arrays; + while (true) { + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch batch, inner_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + 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(RoaringBitmap32 batch_matched, + FilterProbeBatch(array, probe_filter_)); + // Compact each probe batch down to its matched rows so probe_data_ aligns row-for-row + // (ascending file order) with matched_bitmap_ and the later payload output. + if (!batch_matched.IsEmpty()) { + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector matched_slices, + ReaderUtils::GenerateFilteredArrayVector(array, batch_matched)); + probe_arrays.insert(probe_arrays.end(), std::make_move_iterator(matched_slices.begin()), + std::make_move_iterator(matched_slices.end())); + } + } + + std::shared_ptr probe_array; + if (probe_arrays.empty()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + probe_array, arrow::MakeEmptyArray(arrow::struct_(probe_schema_->fields()))); + } else { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(probe_array, + arrow::Concatenate(probe_arrays, arrow_pool_.get())); + } + probe_data_ = arrow::internal::checked_pointer_cast(probe_array); + return Status::OK(); +} + +Result LateMaterializingFileBatchReader::ReadPayloadBatch() { + while (true) { + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatchWithBitmap batch_with_bitmap, + inner_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + state_ = kEOF; + if (probe_cursor_ != probe_data_->length()) { + return Status::Invalid( + fmt::format("probe cursor {} does not match probe data length {}", + probe_cursor_, probe_data_->length())); + } + return MakeEofBatch(); + } + auto& [batch, bitmap] = batch_with_bitmap; + if (bitmap.IsEmpty()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + return Status::Invalid("inner read bitmap is empty."); + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + + // Generate the valid bitmap and row_mapping_ + RoaringBitmap32 valid; + row_mapping_.clear(); + for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { + auto offset = static_cast(*it); + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(offset)); + if (!matched_bitmap_.Contains(file_row)) { + continue; + } + valid.Add(static_cast(offset)); + row_mapping_.push_back(file_row); + } + if (valid.IsEmpty()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + continue; + } + + // Compact the payload superset down to the matched rows (ascending file row order). + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices, + ReaderUtils::GenerateFilteredArrayVector(payload_array, valid)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, + arrow::Concatenate(payload_slices, arrow_pool_.get())); + + auto card = static_cast(valid.Cardinality()); + if (probe_cursor_ + card > probe_data_->length()) { + return Status::Invalid( + fmt::format("probe cache underflow: cursor {} + {} exceeds probe rows {}", + probe_cursor_, card, probe_data_->length())); + } + std::shared_ptr probe_selected = probe_data_->Slice(probe_cursor_, card); + PAIMON_ASSIGN_OR_RAISE( + probe_selected, ArrowUtils::NormalizeArrayOffsets(probe_selected, arrow_pool_.get())); + probe_cursor_ += card; + + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch assembled, + AssembleFullBatch(payload_compacted, probe_selected)); + return assembled; + } +} + +Result LateMaterializingFileBatchReader::AssembleFullBatch( + const std::shared_ptr& payload_array, + const std::shared_ptr& probe_array) { + auto payload_struct = arrow::internal::checked_pointer_cast(payload_array); + auto probe_struct = arrow::internal::checked_pointer_cast(probe_array); + arrow::ArrayVector children; + children.reserve(full_schema_->num_fields()); + for (const auto& field : full_schema_->fields()) { + std::shared_ptr col = payload_struct->GetFieldByName(field->name()); + if (!col) { + col = probe_struct->GetFieldByName(field->name()); + } + if (!col) { + return Status::Invalid( + fmt::format("field {} missing in both payload and probe columns", field->name())); + } + PAIMON_ASSIGN_OR_RAISE(col, ArrowUtils::NormalizeArrayOffsets(col, arrow_pool_.get())); + children.push_back(std::move(col)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr full_struct, + arrow::StructArray::Make(children, full_schema_->fields())); + std::unique_ptr<::ArrowArray> c_array = std::make_unique<::ArrowArray>(); + std::unique_ptr<::ArrowSchema> c_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*full_struct, c_array.get(), c_schema.get())); + return std::make_pair(std::move(c_array), std::move(c_schema)); +} + +Status LateMaterializingFileBatchReader::SetInnerReadSchema( + const std::shared_ptr& read_schema, const std::shared_ptr& predicate, + const std::optional& selection) { + ::ArrowSchema c_read_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_read_schema, predicate, selection)); + return Status::OK(); +} + +Status LateMaterializingFileBatchReader::SetReadSchema( + ::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + Reset(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(full_schema_, arrow::ImportSchema(read_schema)); + predicate_ = predicate; + selection_ = selection_bitmap; + if (predicate_ != nullptr) { + std::set probe_names; + PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names)); + arrow::FieldVector probe_fields; + arrow::FieldVector payload_fields; + for (const auto& field : full_schema_->fields()) { + if (probe_names.count(field->name()) > 0) { + probe_fields.push_back(field); + } else { + payload_fields.push_back(field); + } + } + // probing only pays off when the predicate fields are a strict subset of the read schema + if (!probe_fields.empty() && !payload_fields.empty()) { + probe_schema_ = arrow::schema(probe_fields, full_schema_->metadata()); + payload_schema_ = arrow::schema(payload_fields, full_schema_->metadata()); + PAIMON_RETURN_NOT_OK(PredicateValidator::ValidatePredicateWithSchema( + *probe_schema_, predicate_, /*validate_field_idx=*/false)); + std::map name_to_idx; + for (int32_t i = 0; i < probe_schema_->num_fields(); ++i) { + name_to_idx.emplace(probe_schema_->field(i)->name(), i); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr bound_predicate, + PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx)); + probe_filter_ = std::dynamic_pointer_cast(bound_predicate); + if (!probe_filter_) { + return Status::Invalid("failed to bind predicate to probe schema"); + } + } + } + + if (predicate_ == nullptr || probe_schema_ == nullptr) { + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(full_schema_, predicate_, selection_)); + state_ = kNoLatMat; + } else { + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(probe_schema_, predicate_, selection_)); + state_ = kProbing; + } + return Status::OK(); +} + +Result LateMaterializingFileBatchReader::GetPreviousBatchFileRowId( + uint64_t batch_row_id) const { + if (state_ == kNoLatMat) { + return inner_->GetPreviousBatchFileRowId(batch_row_id); + } + // In kRunning the emitted batch is compacted/reassembled, so row ids come from row_mapping_ + // instead of the inner reader. + if (batch_row_id >= row_mapping_.size()) { + return Status::Invalid( + fmt::format("batch_row_id {} is out of range, last batch row count is {}", batch_row_id, + row_mapping_.size())); + } + return row_mapping_[batch_row_id]; +} + +Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("SeekToRow")); + PAIMON_RETURN_NOT_OK(prefetch_reader->SeekToRow(row_number)); + if (state_ == kRunning || state_ == kEOF) { + if (matched_bitmap_.IsEmpty()) { + state_ = kEOF; + return Status::OK(); + } + int64_t cursor = 0; + for (auto it = matched_bitmap_.Begin(); it != matched_bitmap_.End(); ++it) { + if (static_cast(*it) >= row_number) { + break; + } + ++cursor; + } + probe_cursor_ = cursor; + // a seek after EOF re-activates payload reading + state_ = kRunning; + } + return Status::OK(); +} + +Status LateMaterializingFileBatchReader::SetReadRanges( + const std::vector>& read_ranges) { + if (prefetch_inner_ == nullptr) { + // Only the format reader can act on this hint, and the PrefetchFileBatchReader contract + // lets an implementation that cannot honor it ignore the hint. + return Status::OK(); + } + return prefetch_inner_->SetReadRanges(read_ranges); +} + +void LateMaterializingFileBatchReader::Reset() { + state_ = kInit; + matched_bitmap_ = RoaringBitmap32(); + probe_data_.reset(); + probe_cursor_ = 0; + row_mapping_.clear(); + probe_schema_.reset(); + payload_schema_.reset(); + full_schema_.reset(); + probe_filter_.reset(); + predicate_.reset(); + selection_.reset(); + probe_cursor_ = 0; + row_mapping_.clear(); +} + +} // namespace paimon diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h new file mode 100644 index 000000000..231625db6 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -0,0 +1,184 @@ +/* + * 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 +#include + +#include "fmt/format.h" +#include "paimon/reader/prefetch_file_batch_reader.h" + +namespace paimon { + +class PredicateFilter; + +// For convenience, we abbreviate `Later Materializing` as `LatMat`. +// This reader is installed below the prefetch layer (see +// AbstractSplitRead::CreateFileBatchReader) and performs probe/payload two-phase reads when a +// predicate is pushed down through SetReadSchema; without a predicate it is a plain passthrough. +class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { + public: + static Result> Create( + std::unique_ptr inner, std::shared_ptr pool); + + Result NextBatch() override; + + std::shared_ptr GetReaderMetrics() const override { + return inner_->GetReaderMetrics(); + }; + + void Close() override { + Reset(); + inner_->Close(); + } + + Result> GetFileSchema() const override { + return inner_->GetFileSchema(); + } + + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + + Result GetNumberOfRows() const override { + return inner_->GetNumberOfRows(); + } + + bool SupportPreciseBitmapSelection() const override { + // When probe_schema_ or payload_schema_ is null, lat-mat does not take effect. + // Here we simply pass through the inner reader's support. + return inner_->SupportPreciseBitmapSelection(); + } + + Status SeekToRow(uint64_t row_number) override; + + Result GetNextRowToRead() const override { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GetNextRowToRead")); + return prefetch_reader->GetNextRowToRead(); + } + + Result>> GenReadRanges( + bool* need_prefetch) const override { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GenReadRanges")); + return prefetch_reader->GenReadRanges(need_prefetch); + } + + Status SetReadRanges(const std::vector>& read_ranges) override; + + Result>> PreBufferRange() override { + // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the + // probe data, consider read the payload data as well. + if (prefetch_inner_ == nullptr) { + return std::vector>{}; + } + return prefetch_inner_->PreBufferRange(); + } + + private: + LateMaterializingFileBatchReader(std::unique_ptr inner, + PrefetchFileBatchReader* prefetch_inner, + std::shared_ptr arrow_pool) + : inner_(std::move(inner)), + prefetch_inner_(prefetch_inner), + arrow_pool_(std::move(arrow_pool)) {} + + /// Reset the state of the late materializing reader, does not close inner reader. + void Reset(); + + enum LatMatState { + kInit, + kProbing, // schema is set, probing is in progress + kNoLatMat, // no need to late materialization + kRunning, // Lat-mat is enabled and the payload reader is reading data + kEOF + }; + + /// Read the probe projection once (whole file) and evaluating the predicate batch by batch. + /// This function updates matched_bitmap_ and probe_data_. + /// TODO(zhouhongfeng.zhf): Read the probe data batch by batch to save memory. + Status ReadAndFilterProbeData(); + + Result FilterProbeBatch(const std::shared_ptr& array, + const std::shared_ptr& bound_filter); + + /// Read one payload batch with bitmap (matched rows only) + Result ReadPayloadBatch(); + + /// Combine the compacted payload columns and the selected probe columns into a single struct + /// array following full_schema_'s field order. + Result AssembleFullBatch( + const std::shared_ptr& payload_array, + const std::shared_ptr& probe_array); + + Status SetInnerReadSchema(const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::optional& selection); + + /// Returns the inner reader's prefetch interface, or an error when the format reader does not + /// implement it (avro and blob do not). + Result GetPrefetchReaderOrRaise( + std::string_view function_name) const { + if (prefetch_inner_ == nullptr) { + return Status::NotImplemented( + fmt::format("format reader is not a prefetch reader, function {} not supported", + function_name)); + } + return prefetch_inner_; + } + + /// The probe/payload logic needs nothing beyond FileBatchReader, so the inner reader is held as + /// the base type: parquet and orc readers implement PrefetchFileBatchReader, while avro and + /// blob readers only implement FileBatchReader. The prefetch-only methods are rejected for the + /// latter; AbstractSplitRead::CreateFileBatchReader keeps those formats out of the prefetch + /// layer so nothing calls them. + std::unique_ptr inner_; + /// Non-owning view of inner_ when it implements the prefetch interface, nullptr otherwise. + /// inner_ is never reassigned, so the cast is resolved once in Create(). + PrefetchFileBatchReader* prefetch_inner_ = nullptr; + std::shared_ptr arrow_pool_; + LatMatState state_ = kInit; + std::shared_ptr full_schema_; + // projection holding only the predicate fields; nullptr when probing is not applicable + std::shared_ptr probe_schema_; + // projection holding the payload (non-probe) fields; nullptr when probing is not applicable + std::shared_ptr payload_schema_; + std::shared_ptr predicate_; + // predicate bound to probe_schema_'s field indices; null when probing is not applicable + std::shared_ptr probe_filter_; + std::optional selection_; + // the probe_data_ is sliced and compacted with the matched_bitmap_ + std::shared_ptr probe_data_; + RoaringBitmap32 matched_bitmap_; + // read cursor into probe_data_ for the payload phase + int64_t probe_cursor_ = 0; + // to support GetPreviousBatchFileRowId + std::vector row_mapping_; +}; + +} // namespace paimon diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp new file mode 100644 index 000000000..075fbe107 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -0,0 +1,670 @@ +/* + * 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/reader/late_materializing_file_batch_reader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/builder_nested.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" +#include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/read_ahead_cache.h" +#include "paimon/executor.h" +#include "paimon/format/reader_builder.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/reader/prefetch_file_batch_reader.h" +#include "paimon/status.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/roaring_bitmap32.h" + +namespace paimon::test { + +class LateMaterializingFileBatchReaderTest : public ::testing::Test { + public: + void SetUp() override { + k_field_ = arrow::field("k", arrow::int64()); + v_field_ = arrow::field("v", arrow::utf8()); + full_fields_ = {k_field_, v_field_}; + full_type_ = arrow::struct_(full_fields_); + } + + // Build a struct array with column k (int64, values = ks) and column v (utf8, "v_"). + std::shared_ptr BuildData(const std::vector& ks) { + arrow::StructBuilder builder( + full_type_, arrow::default_memory_pool(), + {std::make_shared(), std::make_shared()}); + auto* k_builder = checked_cast(builder.field_builder(0)); + auto* v_builder = checked_cast(builder.field_builder(1)); + for (size_t i = 0; i < ks.size(); ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(k_builder->Append(ks[i]).ok()); + EXPECT_TRUE(v_builder->Append("v_" + std::to_string(i)).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + struct Row { + int64_t k; + std::string v; + uint64_t file_row; + }; + + // Drive the reader through NextBatchWithBitmap to EOF, decoding the full-schema output rows. + Result> Collect(LateMaterializingFileBatchReader* reader) { + std::vector rows; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + break; + } + auto& [batch, bitmap] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + auto struct_array = arrow::internal::checked_pointer_cast(array); + EXPECT_EQ(bitmap.Cardinality(), static_cast(struct_array->length())); + auto k_array = arrow::internal::checked_pointer_cast( + struct_array->GetFieldByName("k")); + if (!k_array) { + return Status::Invalid("output batch missing k column"); + } + // v is only present when it belongs to the read schema (payload projection). + auto v_array = arrow::internal::checked_pointer_cast( + struct_array->GetFieldByName("v")); + for (int64_t i = 0; i < struct_array->length(); ++i) { + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + rows.push_back(Row{k_array->Value(i), + v_array ? v_array->GetString(i) : std::string(), file_row}); + } + } + return rows; + } + + Status SetReadSchema(LateMaterializingFileBatchReader* reader, + const std::shared_ptr& schema, + const std::shared_ptr& predicate, + const std::optional& selection) { + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); + return reader->SetReadSchema(&c_schema, predicate, selection); + } + + // Collect all output rows as a single concatenated struct array (for schema/nested checks), + // reusing the shared collector so the batch-offset and bitmap contracts are checked too. + Result> CollectStruct(FileBatchReader* reader) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr chunked, + ReadResultCollector::CollectResult(reader)); + if (chunked == nullptr) { + return std::shared_ptr(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr combined, + arrow::Concatenate(chunked->chunks())); + return arrow::internal::checked_pointer_cast(combined); + } + + // Build a struct with 5 columns [a:int64, b:utf8, c:int64, d:utf8, e:int64], each carrying a + // distinct value pattern so any column reordering is detected. + std::shared_ptr BuildMultiFieldData(int32_t n) { + auto type = + arrow::struct_({arrow::field("a", arrow::int64()), arrow::field("b", arrow::utf8()), + arrow::field("c", arrow::int64()), arrow::field("d", arrow::utf8()), + arrow::field("e", arrow::int64())}); + arrow::StructBuilder builder( + type, arrow::default_memory_pool(), + {std::make_shared(), std::make_shared(), + std::make_shared(), std::make_shared(), + std::make_shared()}); + auto* a = checked_cast(builder.field_builder(0)); + auto* b = checked_cast(builder.field_builder(1)); + auto* c = checked_cast(builder.field_builder(2)); + auto* d = checked_cast(builder.field_builder(3)); + auto* e = checked_cast(builder.field_builder(4)); + for (int32_t i = 0; i < n; ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(a->Append(i).ok()); + EXPECT_TRUE(b->Append("b_" + std::to_string(i)).ok()); + EXPECT_TRUE(c->Append(static_cast(i) * 100).ok()); + EXPECT_TRUE(d->Append("d_" + std::to_string(i)).ok()); + EXPECT_TRUE(e->Append(static_cast(i) * 10000).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + // Build a struct with a nested payload column [k:int64, arr:list, tag:utf8]. + std::shared_ptr BuildNestedData(int32_t n) { + auto type = arrow::struct_({arrow::field("k", arrow::int64()), + arrow::field("arr", arrow::list(arrow::int64())), + arrow::field("tag", arrow::utf8())}); + auto arr_value_builder = std::make_shared(); + arrow::StructBuilder builder( + type, arrow::default_memory_pool(), + {std::make_shared(), + std::make_shared(arrow::default_memory_pool(), arr_value_builder), + std::make_shared()}); + auto* k = checked_cast(builder.field_builder(0)); + auto* arr = checked_cast(builder.field_builder(1)); + auto* arr_values = checked_cast(arr->value_builder()); + auto* tag = checked_cast(builder.field_builder(2)); + for (int32_t i = 0; i < n; ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(k->Append(i).ok()); + EXPECT_TRUE(arr->Append().ok()); + EXPECT_TRUE(arr_values->Append(i).ok()); + EXPECT_TRUE(arr_values->Append(i + 1).ok()); + EXPECT_TRUE(tag->Append("t_" + std::to_string(i)).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + protected: + std::shared_ptr k_field_; + std::shared_ptr v_field_; + arrow::FieldVector full_fields_; + std::shared_ptr full_type_; +}; + +// No predicate: the reader must pass through the inner reader unchanged (all rows, all columns). +TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, + std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); + for (int64_t i = 0; i < 5; ++i) { + EXPECT_EQ(rows[i].k, i); + EXPECT_EQ(rows[i].v, "v_" + std::to_string(i)); + EXPECT_EQ(rows[i].file_row, static_cast(i)); + } +} + +// The predicate references every projected column, so the payload set is empty: no late +// materialization, plain pass-through. +TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + // read schema is just {k}; the predicate on k covers all columns -> payload empty + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(0l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema({k_field_}), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); + for (size_t idx = 0; idx < rows.size(); ++idx) { + EXPECT_EQ(rows[idx].k, static_cast(idx)); + // v is outside the read schema, so the pass-through output must not carry it + EXPECT_EQ(rows[idx].v, ""); + EXPECT_EQ(rows[idx].file_row, static_cast(idx)); + } +} + +// Contiguous matched subset spanning multiple batches. +TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(4l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); // k = 5..9 + for (size_t idx = 0; idx < rows.size(); ++idx) { + int64_t expected = 5 + static_cast(idx); + EXPECT_EQ(rows[idx].k, expected); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows[idx].file_row, static_cast(expected)); + } +} + +// Scattered (alternating) matched rows: predicate matches every other row. +TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { + // k = 0,1,0,1,... ; predicate k == 1 matches all odd file rows. + std::vector ks; + for (int i = 0; i < 12; ++i) { + ks.push_back(i % 2); + } + auto data = BuildData(ks); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(1l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 6u); // odd rows 1,3,5,7,9,11 + for (size_t idx = 0; idx < rows.size(); ++idx) { + uint64_t expected_row = 2 * idx + 1; + EXPECT_EQ(rows[idx].k, 1); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected_row)); + EXPECT_EQ(rows[idx].file_row, expected_row); + } +} + +// The selection bitmap further restricts the matched rows: matched must be a subset of selection. +TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { + std::vector ks; + for (int i = 0; i < 12; ++i) { + ks.push_back(i % 2); + } + auto data = BuildData(ks); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(1l)); + // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} + RoaringBitmap32 selection; + selection.Add(1); + selection.Add(5); + selection.Add(9); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, + std::optional(selection))); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 3u); + const std::vector expected_rows = {1u, 5u, 9u}; + for (size_t idx = 0; idx < rows.size(); ++idx) { + EXPECT_EQ(rows[idx].k, 1); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected_rows[idx])); + EXPECT_EQ(rows[idx].file_row, expected_rows[idx]); + } +} + +// No matched rows: the reader returns EOF immediately. +TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(100l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_TRUE(rows.empty()); +} + +// SeekToRow during payload emission must re-align the probe cursor so probe/payload stay matched. +TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(5l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + // First payload batch triggers the probe scan; matched rows are 5..9. + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap first, reader->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(first)); + ReaderUtils::ReleaseReadBatch(std::move(first.first)); + + // Seek forward to file row 8: subsequent output must be exactly rows 8 and 9, correctly paired. + ASSERT_OK(reader->SeekToRow(8)); + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 2u); + EXPECT_EQ(rows[0].k, 8); + EXPECT_EQ(rows[0].v, "v_8"); + EXPECT_EQ(rows[0].file_row, 8u); + EXPECT_EQ(rows[1].k, 9); + EXPECT_EQ(rows[1].v, "v_9"); + EXPECT_EQ(rows[1].file_row, 9u); +} + +// SetReadRanges must be cached and re-forwarded to the inner reader across the probe/payload +// schema switches (SetReadSchema resets the inner reader's ranges). +TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + auto* mock_ptr = mock.get(); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(2l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + std::vector> ranges = {{0, 8}}; + ASSERT_OK(reader->SetReadRanges(ranges)); + + // Drive to EOF; this performs the probe pass and the payload schema switch. + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 6u); // k = 2..7 + for (size_t idx = 0; idx < rows.size(); ++idx) { + int64_t expected = 2 + static_cast(idx); + EXPECT_EQ(rows[idx].k, expected); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows[idx].file_row, static_cast(expected)); + } + + // The inner reader must have received the cached ranges again after the payload switch. + // SetReadSchema (invoked on the payload switch) clears the inner reader's ranges, so the + // cached ranges still being present at the end proves LM re-forwarded them after the switch. + ASSERT_EQ(mock_ptr->GetReadRanges(), ranges); +} + +// SetReadSchema is re-entrant: a second call with a different predicate resets probe state. +TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + + auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(7l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate1, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector rows1, Collect(reader.get())); + ASSERT_EQ(rows1.size(), 2u); // k = 8,9 + for (size_t idx = 0; idx < rows1.size(); ++idx) { + int64_t expected = 8 + static_cast(idx); + EXPECT_EQ(rows1[idx].k, expected); + EXPECT_EQ(rows1[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows1[idx].file_row, static_cast(expected)); + } + + auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(3l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate2, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector rows2, Collect(reader.get())); + ASSERT_EQ(rows2.size(), 3u); // k = 0,1,2 + for (size_t idx = 0; idx < rows2.size(); ++idx) { + EXPECT_EQ(rows2[idx].k, static_cast(idx)); + EXPECT_EQ(rows2[idx].v, "v_" + std::to_string(idx)); + EXPECT_EQ(rows2[idx].file_row, static_cast(idx)); + } +} + +// Forwarded metadata accessors should reflect the inner reader. +TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { + auto data = BuildData({0, 1, 2, 3}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + + ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); + EXPECT_EQ(num_rows, 4u); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> c_file_schema, reader->GetFileSchema()); + auto import_result = arrow::ImportType(c_file_schema.get()); + ASSERT_TRUE(import_result.ok()); + EXPECT_TRUE(import_result.ValueOrDie()->Equals(full_type_)); +} + +// With many columns and a predicate over two non-adjacent probe columns, the output must keep the +// full read-schema field order (and each probe/payload column's values must not be scrambled). +TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { + auto data = BuildMultiFieldData(10); + auto type = data->type(); + auto mock = std::make_unique(data, type, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} + auto pred_a = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); + auto pred_c = + PredicateBuilder::LessThan(/*field_index=*/2, "c", FieldType::BIGINT, Literal(700l)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({pred_a, pred_c})); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_TRUE(result); + // a >= 3 and c(=i*100) < 700 -> i in {3,4,5,6} + ASSERT_EQ(result->length(), 4); + // output field order must equal the requested full schema order + ASSERT_EQ(result->num_fields(), 5); + auto out_type = arrow::internal::checked_pointer_cast(result->type()); + EXPECT_EQ(out_type->field(0)->name(), "a"); + EXPECT_EQ(out_type->field(1)->name(), "b"); + EXPECT_EQ(out_type->field(2)->name(), "c"); + EXPECT_EQ(out_type->field(3)->name(), "d"); + EXPECT_EQ(out_type->field(4)->name(), "e"); + auto a = arrow::internal::checked_pointer_cast(result->GetFieldByName("a")); + auto b = arrow::internal::checked_pointer_cast(result->GetFieldByName("b")); + auto c = arrow::internal::checked_pointer_cast(result->GetFieldByName("c")); + auto d = arrow::internal::checked_pointer_cast(result->GetFieldByName("d")); + auto e = arrow::internal::checked_pointer_cast(result->GetFieldByName("e")); + const std::vector expected = {3, 4, 5, 6}; + for (size_t j = 0; j < expected.size(); ++j) { + int64_t i = expected[j]; + EXPECT_EQ(a->Value(j), i); + EXPECT_EQ(b->GetString(j), "b_" + std::to_string(i)); + EXPECT_EQ(c->Value(j), i * 100); + EXPECT_EQ(d->GetString(j), "d_" + std::to_string(i)); + EXPECT_EQ(e->Value(j), i * 10000); + } +} + +// A nested (list) payload column must round-trip unchanged for the matched rows. +TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { + auto data = BuildNestedData(8); + auto type = data->type(); + auto mock = std::make_unique(data, type, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + // probe = {k}; payload = {arr (list), tag} + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 3); // k = 5,6,7 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto arr = + arrow::internal::checked_pointer_cast(result->GetFieldByName("arr")); + auto tag = + arrow::internal::checked_pointer_cast(result->GetFieldByName("tag")); + ASSERT_TRUE(k && arr && tag); + for (int64_t j = 0; j < result->length(); ++j) { + int64_t i = 5 + j; + EXPECT_EQ(k->Value(j), i); + EXPECT_EQ(tag->GetString(j), "t_" + std::to_string(i)); + auto sub = arrow::internal::checked_pointer_cast(arr->value_slice(j)); + ASSERT_EQ(sub->length(), 2); + EXPECT_EQ(sub->Value(0), i); + EXPECT_EQ(sub->Value(1), i + 1); + } +} + +// The late-materialization reader must work correctly as an inner reader driven by +// PrefetchFileBatchReaderImpl (schema broadcast, range dispatch, seek, row-id tracking). +TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + /*enable_io_metrics=*/false, GetDefaultPool())); + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 6); // k = 4..9 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto v = arrow::internal::checked_pointer_cast(result->GetFieldByName("v")); + ASSERT_TRUE(k && v); + for (int64_t j = 0; j < result->length(); ++j) { + EXPECT_EQ(k->Value(j), 4 + j); + EXPECT_EQ(v->GetString(j), "v_" + std::to_string(4 + j)); + } + impl->Close(); +} + +// Re-setting the read schema on the prefetch impl (which re-broadcasts to the inner LM readers and +// re-plans ranges) must reset the probe state and produce correct results for the new predicate. +TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + /*enable_io_metrics=*/false, GetDefaultPool())); + + auto full_schema = arrow::schema(full_fields_); + auto predicate1 = + PredicateBuilder::GreaterThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(6l)); + ::ArrowSchema c_schema1; + ASSERT_TRUE(arrow::ExportSchema(*full_schema, &c_schema1).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema1, predicate1, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result1, CollectStruct(impl.get())); + ASSERT_TRUE(result1); + ASSERT_EQ(result1->length(), 3); // k = 7,8,9 + auto k1 = + arrow::internal::checked_pointer_cast(result1->GetFieldByName("k")); + auto v1 = + arrow::internal::checked_pointer_cast(result1->GetFieldByName("v")); + ASSERT_TRUE(k1 && v1); + for (int64_t j = 0; j < result1->length(); ++j) { + EXPECT_EQ(k1->Value(j), 7 + j); + EXPECT_EQ(v1->GetString(j), "v_" + std::to_string(7 + j)); + } + + auto predicate2 = + PredicateBuilder::LessThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(3l)); + ::ArrowSchema c_schema2; + ASSERT_TRUE(arrow::ExportSchema(*full_schema, &c_schema2).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema2, predicate2, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result2, CollectStruct(impl.get())); + ASSERT_TRUE(result2); + ASSERT_EQ(result2->length(), 3); // k = 0,1,2 + auto k2 = + arrow::internal::checked_pointer_cast(result2->GetFieldByName("k")); + auto v2 = + arrow::internal::checked_pointer_cast(result2->GetFieldByName("v")); + for (int64_t j = 0; j < result2->length(); ++j) { + EXPECT_EQ(k2->Value(j), j); + EXPECT_EQ(v2->GetString(j), "v_" + std::to_string(j)); + } + impl->Close(); +} + +// With multiple parallel inner readers and per-batch ranges, the prefetch impl dispatches disjoint +// ranges to each LM reader and drives them via EnsureReaderPosition/SeekToRow. The merged output +// must still be exactly the matched rows in ascending file order. +TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSeek) { + std::vector ks; + for (int i = 0; i < 20; ++i) { + ks.push_back(i); + } + auto data = BuildData(ks); + // Per-batch ranges (the mock's default) let the impl split work across the parallel readers, + // and each range-honoring reader only reads its assigned slice. + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + /*enable_io_metrics=*/false, GetDefaultPool())); + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 15); // k = 5..19 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto v = arrow::internal::checked_pointer_cast(result->GetFieldByName("v")); + ASSERT_TRUE(k && v); + for (int64_t j = 0; j < result->length(); ++j) { + EXPECT_EQ(k->Value(j), 5 + j); + EXPECT_EQ(v->GetString(j), "v_" + std::to_string(5 + j)); + } + impl->Close(); +} + +// When the predicate's field type does not match the probe schema, the +// ValidatePredicateWithSchema check must fail with a clear error instead +// of silently producing incorrect results. +TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + // k is int64 in the schema, but the predicate claims FieldType::INT (int32). + auto predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::INT, Literal(10)); + ASSERT_NOK_WITH_MSG( + SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt), + "mismatches"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h new file mode 100644 index 000000000..5c7be5077 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -0,0 +1,70 @@ +/* + * 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/reader/late_materializing_file_batch_reader.h" +#include "paimon/format/reader_builder.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/reader/prefetch_file_batch_reader.h" +#include "paimon/result.h" + +namespace paimon { + +class LateMaterializingReaderBuilder : public ReaderBuilder { + public: + LateMaterializingReaderBuilder(std::unique_ptr inner, + std::shared_ptr pool) + : inner_(std::move(inner)), pool_(std::move(pool)) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + pool_ = pool; + inner_->WithMemoryPool(pool); + return this; + } + + ReaderBuilder* WithCache(const std::shared_ptr& cache) override { + inner_->WithCache(cache); + return this; + } + + ReaderBuilder* WithReadHints(const std::optional& hints) override { + inner_->WithReadHints(hints); + return this; + } + + Result> Build( + const std::shared_ptr& stream) const override { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, + inner_->Build(stream)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(format_reader), pool_)); + return std::unique_ptr(std::move(reader)); + } + + private: + std::unique_ptr inner_; + std::shared_ptr pool_; +}; + +} // namespace paimon 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 c44651790..5285a2d80 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "arrow/array/array_base.h" @@ -30,11 +31,13 @@ #include "paimon/common/io/cache_input_stream.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.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/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; @@ -42,8 +45,149 @@ class Schema; namespace paimon { +struct PrefetchMetricsState { + std::atomic read_ranges_total{0}; + std::atomic read_ranges_after_bitmap{0}; + std::atomic seek_count{0}; + std::atomic produced_batches{0}; + std::atomic consumed_batches{0}; + std::atomic discarded_batches{0}; + std::atomic errors{0}; + std::atomic adaptive_disabled_count{0}; + std::atomic queue_full_count{0}; + std::atomic queue_depth{0}; + std::atomic queue_depth_max{0}; + std::atomic enabled{false}; + std::shared_ptr histograms = std::make_shared(); +}; + +struct PrefetchIoMetricsState { + std::atomic read_requests{0}; + std::atomic read_requested_bytes{0}; + std::atomic read_physical_bytes{0}; + std::atomic read_failed{0}; + std::atomic read_latency_count{0}; + std::atomic read_latency_sum_us{0}; + std::atomic async_requests{0}; + std::atomic async_requested_bytes{0}; + std::atomic async_physical_bytes{0}; + std::atomic async_completed{0}; + std::atomic async_failed{0}; + std::atomic async_pending{0}; + std::atomic async_latency_count{0}; + std::atomic async_latency_sum_us{0}; +}; + namespace { +// Metrics do not synchronize reader state. A concurrently collected snapshot may be approximate. +constexpr std::memory_order kMetricsMemoryOrder = std::memory_order_relaxed; + +uint64_t ElapsedMicros(const std::chrono::steady_clock::time_point& start) { + return static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); +} + +void UpdateMax(std::atomic* target, uint64_t value) { + uint64_t current = target->load(kMetricsMemoryOrder); + while (current < value && !target->compare_exchange_weak(current, value, kMetricsMemoryOrder, + kMetricsMemoryOrder)) { + } +} + +void RecordLatency(uint64_t latency_us, std::atomic* count, + std::atomic* sum_us) { + count->fetch_add(1, kMetricsMemoryOrder); + sum_us->fetch_add(latency_us, kMetricsMemoryOrder); +} + +class MetricsInputStream : public InputStream { + public: + MetricsInputStream(const std::shared_ptr& stream, + const std::shared_ptr& metrics) + : stream_(stream), metrics_(metrics) {} + + MetricsInputStream(std::unique_ptr&& stream, + const std::shared_ptr& metrics) + : stream_(std::move(stream)), metrics_(metrics) {} + + Status Seek(int64_t offset, SeekOrigin origin) override { + return stream_->Seek(offset, origin); + } + + Result GetPos() const override { + return stream_->GetPos(); + } + + Result Read(char* buffer, int64_t size) override { + return RecordRead([&]() { return stream_->Read(buffer, size); }, size); + } + + Result Read(char* buffer, int64_t size, int64_t offset) override { + return RecordRead([&]() { return stream_->Read(buffer, size, offset); }, size); + } + + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + metrics_->async_requests.fetch_add(1, kMetricsMemoryOrder); + metrics_->async_requested_bytes.fetch_add(static_cast(std::max(0, size)), + kMetricsMemoryOrder); + metrics_->async_pending.fetch_add(1, kMetricsMemoryOrder); + std::shared_ptr metrics = metrics_; + const auto start = std::chrono::steady_clock::now(); + stream_->ReadAsync( + buffer, size, offset, + [metrics, size, start, callback = std::move(callback)](Status status) mutable { + metrics->async_pending.fetch_sub(1, kMetricsMemoryOrder); + if (status.ok()) { + metrics->async_completed.fetch_add(1, kMetricsMemoryOrder); + metrics->async_physical_bytes.fetch_add( + static_cast(std::max(0, size)), kMetricsMemoryOrder); + } else { + metrics->async_failed.fetch_add(1, kMetricsMemoryOrder); + } + RecordLatency(ElapsedMicros(start), &metrics->async_latency_count, + &metrics->async_latency_sum_us); + callback(status); + }); + } + + Status Close() override { + return stream_->Close(); + } + + Result GetUri() const override { + return stream_->GetUri(); + } + + Result Length() const override { + return stream_->Length(); + } + + private: + template + Result RecordRead(ReadFunction&& read, int64_t size) { + metrics_->read_requests.fetch_add(1, kMetricsMemoryOrder); + metrics_->read_requested_bytes.fetch_add(static_cast(std::max(0, size)), + kMetricsMemoryOrder); + const auto start = std::chrono::steady_clock::now(); + Result result = read(); + if (result.ok()) { + metrics_->read_physical_bytes.fetch_add( + static_cast(std::max(0, result.value())), kMetricsMemoryOrder); + } else { + metrics_->read_failed.fetch_add(1, kMetricsMemoryOrder); + } + RecordLatency(ElapsedMicros(start), &metrics_->read_latency_count, + &metrics_->read_latency_sum_us); + return result; + } + + std::shared_ptr stream_; + std::shared_ptr metrics_; +}; + std::pair ComputeBatchSliceByReadRange( const std::vector& global_row_ids, const std::pair& read_range) { auto begin_it = @@ -60,7 +204,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, bool enable_io_metrics, const std::shared_ptr& pool) { if (prefetch_max_parallel_num == 0) { return Status::Invalid("prefetch max parallel num should be greater than 0."); @@ -81,24 +225,35 @@ Result> PrefetchFileBatchReaderImpl return Status::Invalid("executor should not be nullptr."); } + std::shared_ptr io_metrics; + if (enable_io_metrics) { + io_metrics = std::make_shared(); + } 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))); + if (io_metrics) { + input_stream = std::make_shared(input_stream, io_metrics); + } cache = std::make_shared(input_stream, cache_config, pool); } std::vector>>> futures; for (uint32_t i = 0; i < prefetch_max_parallel_num; i++) { - futures.push_back( - Via(executor.get(), - [&fs, &data_file_path, data_file_size, &reader_builder, - &cache]() -> Result> { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr input_stream, - fs->Open(FileStatus(data_file_path, data_file_size))); - auto cache_input_stream = - std::make_shared(std::move(input_stream), cache); - return reader_builder->Build(cache_input_stream); - })); + futures.push_back(Via( + executor.get(), + [&fs, &data_file_path, data_file_size, &reader_builder, &cache, + io_metrics]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr input_stream, + fs->Open(FileStatus(data_file_path, data_file_size))); + if (io_metrics) { + input_stream = + std::make_unique(std::move(input_stream), io_metrics); + } + auto cache_input_stream = + std::make_shared(std::move(input_stream), cache); + return reader_builder->Build(cache_input_stream); + })); } std::vector> readers; for (auto& file_batch_reader : CollectAll(futures)) { @@ -121,7 +276,7 @@ Result> PrefetchFileBatchReaderImpl auto reader = std::unique_ptr(new PrefetchFileBatchReaderImpl( readers, batch_size, prefetch_queue_capacity, enable_adaptive_prefetch_strategy, executor, - cache, prefetch_cache_mode)); + cache, io_metrics, pool)); 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. @@ -134,14 +289,17 @@ 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& io_metrics, + const std::shared_ptr& pool) : readers_(std::move(readers)), batch_size_(batch_size), executor_(executor), cache_(cache), - cache_mode_(cache_mode), + arrow_pool_(GetArrowPool(pool)), prefetch_queue_capacity_(prefetch_queue_capacity), - enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy) { + enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy), + prefetch_metrics_(std::make_shared()), + io_metrics_(io_metrics) { for (size_t i = 0; i < readers_.size(); i++) { prefetch_queues_.emplace_back(std::make_unique>()); readers_pos_.emplace_back(std::make_unique>(0)); @@ -158,6 +316,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,12 +333,16 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema( Status PrefetchFileBatchReaderImpl::RefreshReadRanges() { PAIMON_RETURN_NOT_OK(CleanUp()); + if (cache_) { + cache_->Reset(); + } return RefreshReadRangesAfterCleanUp(); } Status PrefetchFileBatchReaderImpl::RefreshReadRangesAfterCleanUp() { bool need_prefetch; PAIMON_ASSIGN_OR_RAISE(auto read_ranges, readers_[0]->GenReadRanges(&need_prefetch)); + const bool format_requested_prefetch = need_prefetch; if (!enable_adaptive_prefetch_strategy_) { need_prefetch = true; @@ -189,8 +354,17 @@ Status PrefetchFileBatchReaderImpl::RefreshReadRangesAfterCleanUp() { } } + if (format_requested_prefetch && !need_prefetch) { + prefetch_metrics_->adaptive_disabled_count.fetch_add(1, kMetricsMemoryOrder); + } need_prefetch_ = need_prefetch; - PAIMON_RETURN_NOT_OK(SetReadRanges(FilterReadRanges(read_ranges, selection_bitmap_))); + prefetch_metrics_->enabled.store(need_prefetch_, kMetricsMemoryOrder); + prefetch_metrics_->read_ranges_total.fetch_add(read_ranges.size(), kMetricsMemoryOrder); + std::vector> filtered_ranges = + FilterReadRanges(read_ranges, selection_bitmap_); + prefetch_metrics_->read_ranges_after_bitmap.fetch_add(filtered_ranges.size(), + kMetricsMemoryOrder); + PAIMON_RETURN_NOT_OK(SetReadRanges(filtered_ranges)); return Status::OK(); } @@ -264,6 +438,7 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { if (batch == std::nullopt) { break; } + prefetch_metrics_->discarded_batches.fetch_add(1, kMetricsMemoryOrder); ReaderUtils::ReleaseReadBatch(std::move(batch.value().batch.first)); } } @@ -289,40 +464,20 @@ Status PrefetchFileBatchReaderImpl::CleanUp() { current_batch_global_row_ids_.clear(); read_ranges_freshed_ = false; clean_prefetch_queue(); + prefetch_metrics_->queue_depth.store(0, kMetricsMemoryOrder); for (size_t i = 0; i < readers_pos_.size(); i++) { readers_pos_[i]->store(0); 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 +487,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()); @@ -365,6 +525,7 @@ void PrefetchFileBatchReaderImpl::Workloop() { } if (prefetch_queues_[reader_idx]->size() >= prefetch_queue_capacity_) { // queue is full, skip + prefetch_metrics_->queue_full_count.fetch_add(1, kMetricsMemoryOrder); continue; } if (readers_pos_[reader_idx]->load() != std::numeric_limits::max()) { @@ -423,7 +584,9 @@ std::optional> PrefetchFileBatchReaderImpl::GetCur Status PrefetchFileBatchReaderImpl::EnsureReaderPosition( size_t reader_idx, const std::pair& current_read_range) const { uint64_t pos = std::max(readers_pos_[reader_idx]->load(), current_read_range.first); - if (readers_[reader_idx]->GetNextRowToRead() != pos) { + PAIMON_ASSIGN_OR_RAISE(uint64_t next_row_to_read, readers_[reader_idx]->GetNextRowToRead()); + if (next_row_to_read != pos) { + prefetch_metrics_->seek_count.fetch_add(1, kMetricsMemoryOrder); return readers_[reader_idx]->SeekToRow(pos); } return Status::OK(); @@ -454,6 +617,7 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( global_row_ids.push_back(global_row_id); } if (global_row_ids.empty()) { + prefetch_metrics_->discarded_batches.fetch_add(1, kMetricsMemoryOrder); ReaderUtils::ReleaseReadBatch(std::move(read_batch)); return Status::OK(); } @@ -471,6 +635,7 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( FindReadRangeContaining(reader_idx, global_row_ids[0]); if (owner_range == std::nullopt) { readers_pos_[reader_idx]->store(global_row_ids[0]); + prefetch_metrics_->discarded_batches.fetch_add(1, kMetricsMemoryOrder); ReaderUtils::ReleaseReadBatch(std::move(read_batch)); return Status::OK(); } @@ -481,29 +646,42 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( } else if (slice_end < c_array->length) { // partially out of range, data before read_range.second has been effectively consumed readers_pos_[reader_idx]->store(read_range.second); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr src_array, + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(c_array.get(), c_schema.get())); - auto array = src_array->Slice(0, slice_end); + std::shared_ptr sliced_array = array->Slice(/*offset=*/0, slice_end); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr normalized_array, + ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportArray(*array, c_array.get(), c_schema.get())); - bitmap.RemoveRange(slice_end, src_array->length()); + arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); + bitmap.RemoveRange(slice_end, array->length()); global_row_ids = std::vector(global_row_ids.begin(), global_row_ids.begin() + slice_end); } else { // all within the range, data before readers_[reader_idx]->GetNextRowToRead() has been // effectively consumed - readers_pos_[reader_idx]->store(readers_[reader_idx]->GetNextRowToRead()); + PAIMON_ASSIGN_OR_RAISE(uint64_t next_row_to_read, + readers_[reader_idx]->GetNextRowToRead()); + readers_pos_[reader_idx]->store(next_row_to_read); } if (bitmap.IsEmpty()) { + prefetch_metrics_->discarded_batches.fetch_add(1, kMetricsMemoryOrder); ReaderUtils::ReleaseReadBatch(std::move(read_batch)); return Status::OK(); } prefetch_queue->push( {read_range, std::move(read_batch_with_bitmap), std::move(global_row_ids)}); + prefetch_metrics_->produced_batches.fetch_add(1, kMetricsMemoryOrder); + const uint64_t queue_depth = + prefetch_metrics_->queue_depth.fetch_add(1, kMetricsMemoryOrder) + 1; + UpdateMax(&prefetch_metrics_->queue_depth_max, queue_depth); } else { std::pair eof_range; PAIMON_ASSIGN_OR_RAISE(eof_range, EofRange()); prefetch_queue->push({eof_range, std::move(read_batch_with_bitmap), {}}); + const uint64_t queue_depth = + prefetch_metrics_->queue_depth.fetch_add(1, kMetricsMemoryOrder) + 1; + UpdateMax(&prefetch_metrics_->queue_depth_max, queue_depth); readers_pos_[reader_idx]->store(std::numeric_limits::max()); } return Status::OK(); @@ -534,8 +712,11 @@ Status PrefetchFileBatchReaderImpl::DoReadBatch(size_t reader_idx) { FileBatchReader* reader = readers_[reader_idx].get(); PAIMON_RETURN_NOT_OK(EnsureReaderPosition(reader_idx, read_range)); - PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap read_batch_with_bitmap, - reader->NextBatchWithBitmap()); + const auto read_start = std::chrono::steady_clock::now(); + Result read_result = reader->NextBatchWithBitmap(); + prefetch_metrics_->histograms->ObserveHistogram(PrefetchMetrics::READER_READ_LATENCY_US, + ElapsedMicros(read_start)); + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap read_batch_with_bitmap, std::move(read_result)); return HandleReadResult(reader_idx, read_range, std::move(read_batch_with_bitmap)); } @@ -549,6 +730,7 @@ Result PrefetchFileBatchReaderImpl::NextBatchW std::make_unique(&PrefetchFileBatchReaderImpl::Workloop, this); } + const auto wait_start = std::chrono::steady_clock::now(); while (true) { PAIMON_RETURN_NOT_OK(GetReadStatus()); if (is_shutdown_) { @@ -587,6 +769,10 @@ Result PrefetchFileBatchReaderImpl::NextBatchW cv_.notify_one(); } current_batch_global_row_ids_ = std::move(prefetch_batch.value().global_row_ids); + prefetch_metrics_->consumed_batches.fetch_add(1, kMetricsMemoryOrder); + prefetch_metrics_->queue_depth.fetch_sub(1, kMetricsMemoryOrder); + prefetch_metrics_->histograms->ObserveHistogram( + PrefetchMetrics::CONSUMER_WAIT_LATENCY_US, ElapsedMicros(wait_start)); return std::move(prefetch_batch).value().batch; } } @@ -597,6 +783,8 @@ Result PrefetchFileBatchReaderImpl::NextBatchW return Status::Invalid("peek batch not suppose to be nullptr"); } current_batch_global_row_ids_.clear(); + prefetch_metrics_->histograms->ObserveHistogram( + PrefetchMetrics::CONSUMER_WAIT_LATENCY_US, ElapsedMicros(wait_start)); return BatchReader::MakeEofBatchWithBitmap(); } if (value_count == prefetch_queues_.size()) { @@ -622,7 +810,69 @@ Status PrefetchFileBatchReaderImpl::SeekToRow(uint64_t row_number) { } std::shared_ptr PrefetchFileBatchReaderImpl::GetReaderMetrics() const { - return MetricsImpl::CollectReadMetrics(readers_); + auto result = std::make_shared(); + if (need_prefetch_) { + result->Merge(MetricsImpl::CollectReadMetrics(readers_)); + } else if (!readers_.empty()) { + result->Merge(readers_[0]->GetReaderMetrics()); + } + + auto set_prefetch_counter = [&result](const char* name, const std::atomic& value) { + result->SetCounter(name, value.load(kMetricsMemoryOrder)); + }; + set_prefetch_counter(PrefetchMetrics::READ_RANGES_TOTAL, prefetch_metrics_->read_ranges_total); + set_prefetch_counter(PrefetchMetrics::READ_RANGES_AFTER_BITMAP, + prefetch_metrics_->read_ranges_after_bitmap); + set_prefetch_counter(PrefetchMetrics::SEEK_COUNT, prefetch_metrics_->seek_count); + set_prefetch_counter(PrefetchMetrics::PRODUCED_BATCHES, prefetch_metrics_->produced_batches); + set_prefetch_counter(PrefetchMetrics::CONSUMED_BATCHES, prefetch_metrics_->consumed_batches); + set_prefetch_counter(PrefetchMetrics::DISCARDED_BATCHES, prefetch_metrics_->discarded_batches); + set_prefetch_counter(PrefetchMetrics::ERRORS, prefetch_metrics_->errors); + set_prefetch_counter(PrefetchMetrics::ADAPTIVE_DISABLED_COUNT, + prefetch_metrics_->adaptive_disabled_count); + set_prefetch_counter(PrefetchMetrics::QUEUE_FULL_COUNT, prefetch_metrics_->queue_full_count); + result->SetGauge(PrefetchMetrics::ENABLED, + prefetch_metrics_->enabled.load(kMetricsMemoryOrder) ? 1.0 : 0.0); + result->SetGauge(PrefetchMetrics::PARALLELISM, + prefetch_metrics_->enabled.load(kMetricsMemoryOrder) + ? static_cast(parallel_num_) + : 1.0); + result->SetGauge(PrefetchMetrics::QUEUE_DEPTH, + static_cast(prefetch_metrics_->queue_depth.load(kMetricsMemoryOrder))); + result->SetGauge( + PrefetchMetrics::QUEUE_DEPTH_MAX, + static_cast(prefetch_metrics_->queue_depth_max.load(kMetricsMemoryOrder))); + result->Merge(prefetch_metrics_->histograms); + if (cache_) { + // PR #209 owns the read-ahead cache metrics. Keep collecting its file-level + // hit/miss counters without defining another C++ metrics surface here. + std::shared_ptr cache_metrics = std::make_shared(); + cache_->CollectMetrics(&cache_metrics); + result->Merge(cache_metrics); + } + + if (!io_metrics_) { + return result; + } + auto set_io_counter = [&result](const char* name, const std::atomic& value) { + result->SetCounter(name, value.load(kMetricsMemoryOrder)); + }; + set_io_counter(PrefetchIoMetrics::READ_REQUESTS, io_metrics_->read_requests); + set_io_counter(PrefetchIoMetrics::READ_REQUESTED_BYTES, io_metrics_->read_requested_bytes); + set_io_counter(PrefetchIoMetrics::READ_PHYSICAL_BYTES, io_metrics_->read_physical_bytes); + set_io_counter(PrefetchIoMetrics::READ_FAILED, io_metrics_->read_failed); + set_io_counter(PrefetchIoMetrics::READ_LATENCY_COUNT, io_metrics_->read_latency_count); + set_io_counter(PrefetchIoMetrics::READ_LATENCY_SUM_US, io_metrics_->read_latency_sum_us); + set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTS, io_metrics_->async_requests); + set_io_counter(PrefetchIoMetrics::ASYNC_REQUESTED_BYTES, io_metrics_->async_requested_bytes); + set_io_counter(PrefetchIoMetrics::ASYNC_PHYSICAL_BYTES, io_metrics_->async_physical_bytes); + set_io_counter(PrefetchIoMetrics::ASYNC_COMPLETED, io_metrics_->async_completed); + set_io_counter(PrefetchIoMetrics::ASYNC_FAILED, io_metrics_->async_failed); + set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_COUNT, io_metrics_->async_latency_count); + set_io_counter(PrefetchIoMetrics::ASYNC_LATENCY_SUM_US, io_metrics_->async_latency_sum_us); + result->SetGauge(PrefetchIoMetrics::ASYNC_PENDING, + static_cast(io_metrics_->async_pending.load(kMetricsMemoryOrder))); + return result; } Result> PrefetchFileBatchReaderImpl::GetFileSchema() const { @@ -650,12 +900,15 @@ Result PrefetchFileBatchReaderImpl::GetNumberOfRows() const { return readers_[0]->GetNumberOfRows(); } -uint64_t PrefetchFileBatchReaderImpl::GetNextRowToRead() const { +Result PrefetchFileBatchReaderImpl::GetNextRowToRead() const { assert(false); return -1; } void PrefetchFileBatchReaderImpl::SetReadStatus(const Status& status) { + if (!status.ok()) { + prefetch_metrics_->errors.fetch_add(1, kMetricsMemoryOrder); + } std::unique_lock lock(rw_mutex_); read_status_ = status; } @@ -676,7 +929,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 78cfbb5f9..c99323698 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -35,23 +35,30 @@ #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; +namespace arrow { +class MemoryPool; +} // namespace arrow + namespace paimon { +class MemoryPool; class ReaderBuilder; class FileSystem; class Executor; class Predicate; class Metrics; +struct PrefetchMetricsState; +struct PrefetchIoMetricsState; class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { public: @@ -60,8 +67,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, + bool enable_io_metrics, const std::shared_ptr& pool); ~PrefetchFileBatchReaderImpl() override; @@ -80,7 +87,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { Status SeekToRow(uint64_t row_number) override; Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; Result GetNumberOfRows() const override; - uint64_t GetNextRowToRead() const override; + Result GetNextRowToRead() const override; void Close() override; Status SetReadRanges(const std::vector>& read_ranges) override; @@ -114,7 +121,8 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { 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& io_metrics, + const std::shared_ptr& pool); Status CleanUp(); void Workloop(); @@ -143,7 +151,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 +169,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { std::condition_variable cv_; std::shared_ptr executor_; std::shared_ptr cache_; - PrefetchCacheMode cache_mode_; + std::unique_ptr arrow_pool_; mutable std::shared_mutex rw_mutex_; std::unique_ptr background_thread_; @@ -174,5 +181,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { const uint32_t prefetch_queue_capacity_; const bool enable_adaptive_prefetch_strategy_; int32_t parallel_num_; + std::shared_ptr prefetch_metrics_; + std::shared_ptr io_metrics_; }; } // namespace paimon 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 b3828f7e5..514e7e49d 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,19 +27,20 @@ #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" #include "paimon/format/format_writer.h" #include "paimon/fs/file_system_factory.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/metrics.h" #include "paimon/predicate/predicate_builder.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 paimon::test { @@ -113,9 +114,65 @@ class ControlledMockFormatReaderBuilder : public ReaderBuilder { mutable std::atomic build_count_{0}; }; +class FailingInputStream : public MockInputStream { + public: + Result Read(char* buffer, int64_t size, int64_t offset) override { + return Status::IOError("injected synchronous read failure"); + } +}; + +class FailingFileSystem : public MockFileSystem { + public: + Result> Open(const std::string& path) const override { + return std::make_unique(); + } +}; + +class IoReadingMockFileBatchReader : public MockFileBatchReader { + public: + IoReadingMockFileBatchReader(const std::shared_ptr& data, + const std::shared_ptr& schema, + int32_t read_batch_size, + const std::shared_ptr& input_stream) + : MockFileBatchReader(data, schema, read_batch_size), input_stream_(input_stream) {} + + Result NextBatchWithBitmap() override { + char value = 0; + PAIMON_ASSIGN_OR_RAISE(int64_t read_size, input_stream_->Read(&value, 1, 0)); + (void)read_size; + return MockFileBatchReader::NextBatchWithBitmap(); + } + + private: + std::shared_ptr input_stream_; +}; + +class IoReadingMockFormatReaderBuilder : public ReaderBuilder { + public: + IoReadingMockFormatReaderBuilder(const std::shared_ptr& data, + const std::shared_ptr& schema, + int32_t read_batch_size) + : data_(data), schema_(schema), read_batch_size_(read_batch_size) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + return this; + } + + Result> Build( + const std::shared_ptr& input_stream) const override { + return std::make_unique(data_, schema_, read_batch_size_, + input_stream); + } + + private: + std::shared_ptr data_; + std::shared_ptr schema_; + int32_t read_batch_size_ = 0; +}; + struct TestParam { std::string file_format; - PrefetchCacheMode cache_mode; + bool read_ahead_cache_enabled; }; class PrefetchFileBatchReaderImplTest : public ::testing::Test, @@ -194,7 +251,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 +266,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(), + /*enable_io_metrics=*/true, 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 +333,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 +351,16 @@ 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(), /*enable_io_metrics=*/false, GetDefaultPool())); + if (prefetch_max_parallel_num == 1) { + ASSERT_NOK( + reader->GetReaderMetrics()->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); + } ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(auto array_and_row_ids, CollectResultAndRowIds(reader.get())); @@ -323,14 +377,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(), /*enable_io_metrics=*/true, 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, @@ -345,6 +398,30 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLimits) { // test metrics auto read_metrics = reader->GetReaderMetrics(); ASSERT_TRUE(read_metrics); + ASSERT_OK_AND_ASSIGN(double prefetch_enabled, read_metrics->GetGauge(PrefetchMetrics::ENABLED)); + ASSERT_OK_AND_ASSIGN(uint64_t produced_batches, + read_metrics->GetCounter(PrefetchMetrics::PRODUCED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t consumed_batches, + read_metrics->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t read_ranges, + read_metrics->GetCounter(PrefetchMetrics::READ_RANGES_TOTAL)); + ASSERT_EQ(prefetch_enabled, 1.0); + ASSERT_GT(produced_batches, 0); + ASSERT_EQ(consumed_batches, 8); + ASSERT_GT(read_ranges, 0); + ASSERT_OK(read_metrics->GetGauge(PrefetchMetrics::QUEUE_DEPTH)); + ASSERT_OK_AND_ASSIGN(uint64_t async_completed, + read_metrics->GetCounter(PrefetchIoMetrics::ASYNC_COMPLETED)); + ASSERT_OK_AND_ASSIGN(uint64_t async_failed, + read_metrics->GetCounter(PrefetchIoMetrics::ASYNC_FAILED)); + ASSERT_OK_AND_ASSIGN(uint64_t async_latency_count, + read_metrics->GetCounter(PrefetchIoMetrics::ASYNC_LATENCY_COUNT)); + ASSERT_EQ(async_latency_count, async_completed + async_failed); + + std::shared_ptr second_metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t second_consumed_batches, + second_metrics->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_EQ(second_consumed_batches, consumed_batches); } TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithoutInitializeReadRanges) { @@ -353,20 +430,49 @@ 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(), /*enable_io_metrics=*/false, GetDefaultPool())); // simulate read limits, only read 8 batches ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "prefetch reader read ranges are not initialized"); reader->Close(); } +TEST_F(PrefetchFileBatchReaderImplTest, TestFailedIoMetrics) { + auto data_array = PrepareArray(10); + IoReadingMockFormatReaderBuilder reader_builder(data_array, data_type_, /*read_batch_size=*/10); + auto failing_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN( + auto reader, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &reader_builder, failing_fs, + /*prefetch_max_parallel_num=*/1, /*batch_size=*/10, + /*prefetch_batch_count=*/2, /*enable_adaptive_prefetch_strategy=*/false, executor_, + /*initialize_read_ranges=*/true, /*read_ahead_cache_enabled=*/false, CacheConfig(), + /*enable_io_metrics=*/true, GetDefaultPool())); + + ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(), "injected synchronous read failure"); + std::shared_ptr metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t read_requests, + metrics->GetCounter(PrefetchIoMetrics::READ_REQUESTS)); + ASSERT_OK_AND_ASSIGN(uint64_t failed_requests, + metrics->GetCounter(PrefetchIoMetrics::READ_FAILED)); + ASSERT_OK_AND_ASSIGN(uint64_t physical_bytes, + metrics->GetCounter(PrefetchIoMetrics::READ_PHYSICAL_BYTES)); + ASSERT_EQ(read_requests, 1); + ASSERT_EQ(failed_requests, 1); + ASSERT_EQ(physical_bytes, 0); + ASSERT_OK_AND_ASSIGN(uint64_t latency_count, + metrics->GetCounter(PrefetchIoMetrics::READ_LATENCY_COUNT)); + ASSERT_OK(metrics->GetCounter(PrefetchIoMetrics::READ_LATENCY_SUM_US)); + ASSERT_EQ(latency_count, read_requests); +} + TEST_F(PrefetchFileBatchReaderImplTest, FilterReadRangesWithoutBitmap) { std::vector> read_ranges = { {0, 1000}, {1000, 2000}, {2000, 3000}, {3000, 4000}, {4000, 5000}, @@ -430,14 +536,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_OK(prefetch_reader->RefreshReadRanges()); std::vector> read_ranges_0 = {{0, 30}, {90, 101}}; @@ -460,17 +565,24 @@ 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(), /*enable_io_metrics=*/false, GetDefaultPool())); ASSERT_FALSE(reader->NeedPrefetch()); + std::shared_ptr metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(double enabled, metrics->GetGauge(PrefetchMetrics::ENABLED)); + ASSERT_OK_AND_ASSIGN(double parallelism, metrics->GetGauge(PrefetchMetrics::PARALLELISM)); + ASSERT_OK_AND_ASSIGN(uint64_t adaptive_disabled_count, + metrics->GetCounter(PrefetchMetrics::ADAPTIVE_DISABLED_COUNT)); + ASSERT_EQ(enabled, 0.0); + ASSERT_EQ(parallelism, 1.0); + ASSERT_EQ(adaptive_disabled_count, 1); } TEST_F(PrefetchFileBatchReaderImplTest, SetReadRanges) { @@ -478,14 +590,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); ASSERT_FALSE(prefetch_reader->need_prefetch_); prefetch_reader->need_prefetch_ = true; @@ -522,14 +633,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->need_prefetch_ = true; @@ -539,43 +649,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, /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->Workloop(); @@ -589,14 +679,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->is_shutdown_ = true; @@ -608,14 +697,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); prefetch_reader->read_ranges_in_group_ = {{}}; @@ -627,14 +715,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(), /*enable_io_metrics=*/false, 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 +735,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(), /*enable_io_metrics=*/false, 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 +780,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(), /*enable_io_metrics=*/false, GetDefaultPool())); auto prefetch_reader = dynamic_cast(reader.get()); for (int32_t i = 0; i < prefetch_max_parallel_num; i++) { @@ -717,6 +802,9 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) { ASSERT_FALSE(prefetch_reader->is_shutdown_); ASSERT_NOK(prefetch_reader->GetReadStatus()); ASSERT_FALSE(HasValue(prefetch_reader->prefetch_queues_)); + std::shared_ptr metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t prefetch_errors, metrics->GetCounter(PrefetchMetrics::ERRORS)); + ASSERT_GT(prefetch_errors, 0); // call NextBatch again, will still return error status auto batch_result2 = reader->NextBatchWithBitmap(); @@ -730,14 +818,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(), /*enable_io_metrics=*/false, 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 +837,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(), /*enable_io_metrics=*/false, 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; @@ -765,10 +851,28 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) { auto expected_array = std::make_shared(data_array); ASSERT_TRUE(array_and_row_ids.first->Equals(expected_array)); + std::shared_ptr eof_metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t produced_batches, + eof_metrics->GetCounter(PrefetchMetrics::PRODUCED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t consumed_batches, + eof_metrics->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t discarded_batches, + eof_metrics->GetCounter(PrefetchMetrics::DISCARDED_BATCHES)); + // continue to call NextBatch() after reading eof ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, reader->NextBatchWithBitmap()); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_TRUE(BatchReader::IsEofBatch(batch_with_bitmap)); + std::shared_ptr repeated_eof_metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t produced_batches_after, + repeated_eof_metrics->GetCounter(PrefetchMetrics::PRODUCED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t consumed_batches_after, + repeated_eof_metrics->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t discarded_batches_after, + repeated_eof_metrics->GetCounter(PrefetchMetrics::DISCARDED_BATCHES)); + ASSERT_EQ(produced_batches_after, produced_batches); + ASSERT_EQ(consumed_batches_after, consumed_batches); + ASSERT_EQ(discarded_batches_after, discarded_batches); } TEST_F(PrefetchFileBatchReaderImplTest, TestCreateReaderWithoutNextBatch) { @@ -776,14 +880,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(), /*enable_io_metrics=*/false, GetDefaultPool())); } TEST_F(PrefetchFileBatchReaderImplTest, TestInvalidCase) { @@ -797,16 +900,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(), + /*enable_io_metrics=*/false, 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(), + /*enable_io_metrics=*/false, GetDefaultPool())); } { ASSERT_NOK(PrefetchFileBatchReaderImpl::Create( @@ -814,33 +917,33 @@ 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(), /*enable_io_metrics=*/false, + 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(), + /*enable_io_metrics=*/false, 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(), + /*enable_io_metrics=*/false, 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(), /*enable_io_metrics=*/false, GetDefaultPool())); ASSERT_NOK_WITH_MSG(reader->SeekToRow(/*row_number=*/101), "not support seek to row for prefetch reader"); } @@ -850,7 +953,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 +969,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; @@ -876,6 +979,18 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom auto expected_array = std::make_shared(expected_array_vector); ASSERT_TRUE(expected_array->Equals(array_and_row_ids.first)); ASSERT_EQ(expected_row_ids, array_and_row_ids.second); + + std::shared_ptr metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t sync_requests, + metrics->GetCounter(PrefetchIoMetrics::READ_REQUESTS)); + ASSERT_OK_AND_ASSIGN(uint64_t async_requests, + metrics->GetCounter(PrefetchIoMetrics::ASYNC_REQUESTS)); + ASSERT_OK_AND_ASSIGN(uint64_t sync_physical_bytes, + metrics->GetCounter(PrefetchIoMetrics::READ_PHYSICAL_BYTES)); + ASSERT_OK_AND_ASSIGN(uint64_t async_physical_bytes, + metrics->GetCounter(PrefetchIoMetrics::ASYNC_PHYSICAL_BYTES)); + ASSERT_GT(sync_requests + async_requests, 0); + ASSERT_GT(sync_physical_bytes + async_physical_bytes, 0); } /// There are three stripes: [0,30), [30,60), [60,90). Each stripe has 3 row groups. @@ -883,7 +998,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 +1015,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())); @@ -931,8 +1046,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithBitmap) { /*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())); + /*read_ahead_cache_enabled=*/true, CacheConfig(), + /*enable_io_metrics=*/false, 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 +1061,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 +1074,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())); @@ -976,12 +1091,28 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestRowMapping) { ASSERT_EQ(reader->GetPreviousBatchFileRowId(i).value(), 70 + i); } + std::shared_ptr metrics_before_schema = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t consumed_before_schema, + metrics_before_schema->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t ranges_before_schema, + metrics_before_schema->GetCounter(PrefetchMetrics::READ_RANGES_TOTAL)); + // Set read schema again std::unique_ptr c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok()); predicate = PredicateBuilder::Between(/*field_index=*/1, /*field_name=*/"f1", FieldType::BIGINT, Literal(30l), Literal(49l)); ASSERT_OK(reader->SetReadSchema(c_schema.get(), predicate, std::nullopt)); + std::shared_ptr metrics_after_schema = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t consumed_after_schema, + metrics_after_schema->GetCounter(PrefetchMetrics::CONSUMED_BATCHES)); + ASSERT_OK_AND_ASSIGN(uint64_t ranges_after_schema, + metrics_after_schema->GetCounter(PrefetchMetrics::READ_RANGES_TOTAL)); + ASSERT_OK_AND_ASSIGN(double queue_depth_after_schema, + metrics_after_schema->GetGauge(PrefetchMetrics::QUEUE_DEPTH)); + ASSERT_EQ(consumed_after_schema, consumed_before_schema); + ASSERT_GT(ranges_after_schema, ranges_before_schema); + ASSERT_EQ(queue_depth_after_schema, 0.0); ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); ASSERT_OK_AND_ASSIGN(batch, diff --git a/src/paimon/common/sst/sst_file_io_test.cpp b/src/paimon/common/sst/sst_file_io_test.cpp index eac9a8a89..54d4b1b13 100644 --- a/src/paimon/common/sst/sst_file_io_test.cpp +++ b/src/paimon/common/sst/sst_file_io_test.cpp @@ -93,7 +93,7 @@ TEST_P(SstFileIOTest, TestSimple) { fs_->Create(index_path, /*overwrite=*/false)); // write data - auto bf = BloomFilter::Create(30, 0.01); + ASSERT_OK_AND_ASSIGN(std::shared_ptr bf, BloomFilter::Create(30, 0.01)); auto seg_for_bf = MemorySegment::AllocateHeapMemory(bf->ByteLength(), pool_.get()); ASSERT_OK(bf->SetMemorySegment(seg_for_bf)); auto writer = std::make_shared(out, bf, 50, factory, pool_); @@ -234,7 +234,7 @@ TEST_F(SstFileIOTest, TestIOException) { CHECK_HOOK_STATUS(out_result.status(), i); std::shared_ptr out = std::move(out_result).value(); - auto bf = BloomFilter::Create(30, 0.01); + ASSERT_OK_AND_ASSIGN(std::shared_ptr bf, BloomFilter::Create(30, 0.01)); MemorySegment seg_for_bf = MemorySegment::AllocateHeapMemory(bf->ByteLength(), pool_.get()); ASSERT_OK(bf->SetMemorySegment(seg_for_bf)); auto writer = std::make_shared(out, bf, 50, factory, pool_); diff --git a/src/paimon/common/sst/sst_file_writer.cpp b/src/paimon/common/sst/sst_file_writer.cpp index ec736e33c..f2b3b2dee 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 74b95b19c..8908f9f0c 100644 --- a/src/paimon/common/table/special_fields.h +++ b/src/paimon/common/table/special_fields.h @@ -21,6 +21,7 @@ #include #include #include +#include #include "arrow/type_fwd.h" #include "paimon/common/types/data_field.h" @@ -66,13 +67,20 @@ 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; } return field_name == SequenceNumber().Name() || field_name == ValueKind().Name() || field_name == RowKind().Name() || field_name == RowId().Name() || - field_name == IndexScore().Name(); + field_name == IndexScore().Name() || field_name == RealtimeOffset().Name(); } // TODO(xinyu.lxy): add a func to complete row-tracking fields diff --git a/src/paimon/common/table/special_fields_test.cpp b/src/paimon/common/table/special_fields_test.cpp index 68e805fd6..58a025ba2 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_TRUE(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 9b6d3c903..2bf5d73cb 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 33308ab69..d04e5ade5 100644 --- a/src/paimon/common/types/data_type_json_parser.cpp +++ b/src/paimon/common/types/data_type_json_parser.cpp @@ -19,12 +19,12 @@ #include "paimon/common/types/data_type_json_parser.h" -#include #include #include #include #include #include +#include #include #include #include @@ -33,6 +33,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 +149,7 @@ enum class Keyword : int32_t { ROW, BLOB, VARIANT, + VECTOR, // NULL is keyword in c++ NULL_, RAW, @@ -197,6 +199,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 +252,7 @@ class TokenParser { Result> ParseDoubleType(); Result> ParseTimestampType(); Result> ParseTimestampLtzType(); + Result> ParseVectorType(); Result ParseOptionalPrecision(int32_t default_precision); private: @@ -326,10 +330,7 @@ std::vector Tokenize(const std::string& chars) { builder.clear(); cursor = ConsumeIdentifier(chars, cursor, builder); auto token = builder.str(); - auto normalized_token = token; - std::transform(normalized_token.begin(), normalized_token.end(), - normalized_token.begin(), - [](unsigned char c) { return std::toupper(c); }); + std::string normalized_token = StringUtils::ToUpperCase(token); if (Keywords().find(normalized_token) != Keywords().end()) { tokens.emplace_back(TokenType::KEYWORD, cursor, normalized_token); } else { @@ -526,6 +527,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 +610,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 +687,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 +711,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 92cb5d55f..2134236a7 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 5026db020..e5dfbc21e 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 9568c3ff0..d6eacdc1f 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 000000000..9c55c165b --- /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 707e888f7..1ba4c153b 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" @@ -37,12 +38,21 @@ namespace paimon { namespace { -bool HasNonZeroOffset(const std::shared_ptr& data) { +bool NeedsNormalization(const std::shared_ptr& data) { if (data->offset != 0) { return true; } + if (data->type->id() == arrow::Type::STRUCT) { + for (const auto& child : data->child_data) { + // StructArray::Slice(0, length) shortens only the parent ArrayData. Its children may + // still describe the full unsliced arrays, which cannot be imported as a RecordBatch. + if (child->length != data->length) { + return true; + } + } + } for (const auto& child : data->child_data) { - if (HasNonZeroOffset(child)) { + if (NeedsNormalization(child)) { return true; } } @@ -160,6 +170,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) { @@ -211,7 +243,7 @@ Result> RebaseFixedWidth( /// proportional to the number of rows instead of the number of value bytes. Result> RebaseToZeroOffset( const std::shared_ptr& data, arrow::MemoryPool* pool) { - if (!HasNonZeroOffset(data)) { + if (!NeedsNormalization(data)) { return data; } // An empty array may not carry the buffers the layouts below slice. @@ -233,6 +265,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,16 +354,44 @@ 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; } } +uint64_t ArrowUtils::GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t result = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + result += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + result += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + result += GetArrayMemoryUsage(data->dictionary); + } + return result; +} + bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type) { 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 +425,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); @@ -400,15 +468,13 @@ Result> ArrowUtils::NormalizeRecordBatchOffs arrow::ArrayVector normalized_columns; for (int32_t i = 0; i < record_batch->num_columns(); ++i) { const std::shared_ptr& column = record_batch->column(i); - if (!HasNonZeroOffset(column->data())) { + if (!NeedsNormalization(column->data())) { continue; } if (normalized_columns.empty()) { normalized_columns = record_batch->columns(); } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_data, - RebaseToZeroOffset(column->data(), pool)); - normalized_columns[i] = arrow::MakeArray(normalized_data); + PAIMON_ASSIGN_OR_RAISE(normalized_columns[i], NormalizeArrayOffsets(column, pool)); } if (normalized_columns.empty()) { return record_batch; @@ -417,6 +483,13 @@ Result> ArrowUtils::NormalizeRecordBatchOffs std::move(normalized_columns)); } +Result> ArrowUtils::NormalizeArrayOffsets( + const std::shared_ptr& array, arrow::MemoryPool* pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_data, + RebaseToZeroOffset(array->data(), pool)); + return arrow::MakeArray(normalized_data); +} + Result ArrowUtils::GetCompressionType(const std::string& compression) { std::string normalized = StringUtils::ToLowerCase(compression); if (normalized.empty() || normalized == "none") { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index 326b3889e..59db09815 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include "arrow/api.h" @@ -48,6 +49,8 @@ class PAIMON_EXPORT ArrowUtils { // avoid subsequent multi-threading problems. static void TraverseArray(const std::shared_ptr& array); + static uint64_t GetArrayMemoryUsage(const std::shared_ptr& data); + static Result> RemoveFieldFromStructArray( const std::shared_ptr& struct_array, const std::string& field_name); @@ -57,6 +60,11 @@ class PAIMON_EXPORT ArrowUtils { static Result> NormalizeRecordBatchOffsets( const std::shared_ptr& record_batch, arrow::MemoryPool* pool); + /// Returns an Array with zero offsets. Struct children are also sliced to the parent's visible + /// range so the result can be exported and imported as a RecordBatch. + static Result> NormalizeArrayOffsets( + const std::shared_ptr& array, arrow::MemoryPool* pool); + static bool EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type); diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp b/src/paimon/common/utils/arrow/arrow_utils_test.cpp index 3680291c9..2aad598b5 100644 --- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp @@ -20,6 +20,7 @@ #include "paimon/common/utils/arrow/arrow_utils.h" #include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" #include "paimon/common/types/data_field.h" @@ -249,6 +250,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); @@ -465,6 +502,33 @@ TEST(ArrowUtilsTest, TestNormalizeRecordBatchOffsets) { ASSERT_EQ(unchanged_batch.get(), normalized_batch.get()); } +TEST(ArrowUtilsTest, TestNormalizeArrayOffsetsSlicesZeroOffsetStructChildren) { + std::shared_ptr ints = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2, 3]").ValueOrDie(); + std::shared_ptr texts = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c", "d"])") + .ValueOrDie(); + std::shared_ptr array = + arrow::StructArray::Make({ints, texts}, std::vector{"i", "s"}).ValueOrDie(); + std::shared_ptr sliced = array->Slice(/*offset=*/0, /*length=*/2); + ASSERT_EQ(0, sliced->offset()); + ASSERT_EQ(4, sliced->data()->child_data[0]->length); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr normalized, + ArrowUtils::NormalizeArrayOffsets(sliced, arrow::default_memory_pool())); + ASSERT_TRUE(normalized->Equals(sliced)); + ASSERT_EQ(0, normalized->offset()); + ASSERT_EQ(2, normalized->data()->child_data[0]->length); + ASSERT_EQ(2, normalized->data()->child_data[1]->length); + + ::ArrowArray c_array = {}; + ::ArrowSchema c_schema = {}; + ASSERT_TRUE(arrow::ExportArray(*normalized, &c_array, &c_schema).ok()); + std::shared_ptr batch = + arrow::ImportRecordBatch(&c_array, &c_schema).ValueOrDie(); + ASSERT_EQ(2, batch->num_rows()); +} + namespace { /// A buffer that rebasing must expose as a view into the source. @@ -523,6 +587,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 +825,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/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 7e8986be5..1e9695332 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -24,12 +24,31 @@ #include #include +#include "arrow/c/abi.h" +#include "arrow/c/helpers.h" #include "arrow/memory_pool.h" #include "arrow/status.h" #include "fmt/format.h" #include "paimon/memory/memory_pool.h" namespace paimon { +namespace { + +struct ArrowArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleaseArrowArray(ArrowArray* array) { + std::unique_ptr data( + static_cast(array->private_data)); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); +} + +} // namespace class ArrowMemPoolAdaptor : public arrow::MemoryPool { public: @@ -107,4 +126,26 @@ std::unique_ptr GetArrowPool(const std::shared_ptr(pool); } +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release) { + return Status::Invalid("cannot retain Arrow array memory pool"); + } + if (!arrow_pool) { + ArrowArrayRelease(array); + return Status::Invalid("cannot retain Arrow array memory pool"); + } + std::unique_ptr data; + try { + data = std::make_unique( + ArrowArrayPrivateData{array->release, array->private_data, arrow_pool}); + } catch (const std::bad_alloc&) { + ArrowArrayRelease(array); + return Status::OutOfMemory("failed to retain Arrow array memory pool"); + } + array->private_data = data.release(); + array->release = ReleaseArrowArray; + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 96b59e3e8..214bb4509 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -23,11 +23,17 @@ #include "arrow/memory_pool.h" #include "paimon/memory/memory_pool.h" +#include "paimon/status.h" #include "paimon/visibility.h" +struct ArrowArray; + namespace paimon { PAIMON_EXPORT std::unique_ptr GetArrowPool( const std::shared_ptr& pool); +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool); + } // namespace paimon 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 000000000..e5cc8396f --- /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 000000000..0031a5de7 --- /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 000000000..1cce6853e --- /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 43ec7d40d..ec773e32c 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 63aa7549e..2e96ad6f1 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 974c4d7af..18c421780 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/bloom_filter.cpp b/src/paimon/common/utils/bloom_filter.cpp index d66420f47..ba9fa4b44 100644 --- a/src/paimon/common/utils/bloom_filter.cpp +++ b/src/paimon/common/utils/bloom_filter.cpp @@ -25,29 +25,35 @@ namespace paimon { -int32_t BloomFilter::OptimalNumOfBits(int64_t expect_entries, double fpp) { - if (expect_entries <= 0 || fpp <= 0.0 || fpp >= 1.0) { +int32_t BloomFilter::OptimalNumOfBits(int64_t expected_entries, double fpp) { + if (expected_entries <= 0 || fpp <= 0.0 || fpp >= 1.0) { return 0; } - double result = -static_cast(expect_entries) * log(fpp) / (log(2) * log(2)); + double result = -static_cast(expected_entries) * log(fpp) / (log(2) * log(2)); if (result > INT32_MAX) return INT32_MAX; if (result < 0) return 0; return static_cast(result); } -int32_t BloomFilter::OptimalNumOfHashFunctions(int64_t expect_entries, int64_t bit_size) { - if (expect_entries <= 0) { +int32_t BloomFilter::OptimalNumOfHashFunctions(int64_t expected_entries, int64_t bit_size) { + if (expected_entries <= 0) { return 1; } - double ratio = static_cast(bit_size) / static_cast(expect_entries); + double ratio = static_cast(bit_size) / static_cast(expected_entries); double result = ratio * std::log(2.0); return std::max(1, static_cast(std::round(result))); } -std::shared_ptr BloomFilter::Create(int64_t expect_entries, double fpp) { +Result> BloomFilter::Create(int64_t expected_entries, double fpp) { + if (expected_entries <= 0) { + return Status::Invalid("expected entries must be greater than 0 for bloom filter"); + } + if (!std::isfinite(fpp) || fpp <= 0.0 || fpp >= 1.0) { + return Status::Invalid("fpp must be greater than 0 and less than 1 for bloom filter"); + } auto bytes = - static_cast(ceil(BloomFilter::OptimalNumOfBits(expect_entries, fpp) / 8.0)); - return std::make_shared(expect_entries, bytes); + static_cast(ceil(BloomFilter::OptimalNumOfBits(expected_entries, fpp) / 8.0)); + return std::make_shared(expected_entries, bytes); } BloomFilter::BloomFilter(int64_t expected_entries, int32_t byte_length) diff --git a/src/paimon/common/utils/bloom_filter.h b/src/paimon/common/utils/bloom_filter.h index 8f1b32711..61e27af59 100644 --- a/src/paimon/common/utils/bloom_filter.h +++ b/src/paimon/common/utils/bloom_filter.h @@ -22,7 +22,7 @@ #include #include "paimon/common/utils/bit_set.h" -#include "paimon/memory/bytes.h" +#include "paimon/result.h" #include "paimon/visibility.h" namespace paimon { @@ -30,9 +30,9 @@ namespace paimon { /// Bloom filter based on MemorySegment. class PAIMON_EXPORT BloomFilter { public: - static int32_t OptimalNumOfBits(int64_t expect_entries, double fpp); - static int32_t OptimalNumOfHashFunctions(int64_t expect_entries, int64_t bit_size); - static std::shared_ptr Create(int64_t expect_entries, double fpp); + static int32_t OptimalNumOfBits(int64_t expected_entries, double fpp); + static int32_t OptimalNumOfHashFunctions(int64_t expected_entries, int64_t bit_size); + static Result> Create(int64_t expected_entries, double fpp); public: BloomFilter(int64_t expected_entries, int32_t byte_length); diff --git a/src/paimon/common/utils/bloom_filter64.cpp b/src/paimon/common/utils/bloom_filter64.cpp index f2685a13f..7a975161c 100644 --- a/src/paimon/common/utils/bloom_filter64.cpp +++ b/src/paimon/common/utils/bloom_filter64.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include #include "paimon/memory/bytes.h" @@ -47,17 +49,42 @@ bool BloomFilter64::BitSet::Get(int32_t index) const { } int32_t BloomFilter64::BitSet::BitSize() const { - return (bytes_->size() - offset_) * BloomFilter64::BYTE_SIZE; + return ByteLength() * BloomFilter64::BYTE_SIZE; } -BloomFilter64::BloomFilter64(int64_t items, double fpp, const std::shared_ptr& pool) - : pool_(pool) { - auto nb = static_cast(-items * std::log(fpp) / (std::log(2) * std::log(2))); - num_bits_ = nb + (BloomFilter64::BYTE_SIZE - (nb % BloomFilter64::BYTE_SIZE)); - num_hash_functions_ = std::max( - 1, static_cast(std::round(static_cast(num_bits_) / items * std::log(2)))); - auto bytes = std::make_shared(num_bits_ / BloomFilter64::BYTE_SIZE, pool_.get()); - bit_set_ = std::make_unique(bytes, /*offset=*/0); +int32_t BloomFilter64::BitSet::ByteLength() const { + return static_cast(bytes_->size() - offset_); +} + +void BloomFilter64::BitSet::ToByteArray(int32_t offset, int32_t length, char* bytes) const { + assert(bytes); + assert(offset >= 0); + assert(length >= 0); + assert(static_cast(offset_ + length) <= bytes_->size()); + std::memcpy(bytes + offset, bytes_->data() + offset_, length); +} + +Result BloomFilter64::Create(int64_t items, double fpp, + const std::shared_ptr& pool) { + if (items <= 0) { + return Status::Invalid("items must be greater than 0 for bloom filter"); + } + if (!std::isfinite(fpp) || fpp <= 0.0 || fpp >= 1.0) { + return Status::Invalid("fpp must be greater than 0 and less than 1 for bloom filter"); + } + const double log_two = std::log(2); + const double estimated_bits = -static_cast(items) * std::log(fpp) / (log_two * log_two); + if (estimated_bits > std::numeric_limits::max() - BYTE_SIZE) { + return Status::Invalid("bloom filter size exceeds the supported range"); + } + const auto num_bits_without_padding = static_cast(estimated_bits); + const int32_t num_bits = + num_bits_without_padding + (BYTE_SIZE - (num_bits_without_padding % BYTE_SIZE)); + const int32_t num_hash_functions = std::max( + 1, static_cast(std::round(static_cast(num_bits) / items * log_two))); + auto bytes = std::make_shared(num_bits / BYTE_SIZE, pool.get()); + auto bit_set = std::make_unique(bytes, /*offset=*/0); + return BloomFilter64(num_hash_functions, std::move(bit_set), pool); } BloomFilter64::BloomFilter64(int32_t num_hash_functions, std::unique_ptr&& bit_set) @@ -65,6 +92,13 @@ BloomFilter64::BloomFilter64(int32_t num_hash_functions, std::unique_ptr num_hash_functions_(num_hash_functions), bit_set_(std::move(bit_set)) {} +BloomFilter64::BloomFilter64(int32_t num_hash_functions, std::unique_ptr&& bit_set, + const std::shared_ptr& pool) + : num_bits_(bit_set->BitSize()), + num_hash_functions_(num_hash_functions), + pool_(pool), + bit_set_(std::move(bit_set)) {} + void BloomFilter64::AddHash(int64_t hash64) { auto hash1 = static_cast(hash64); auto hash2 = static_cast(static_cast(hash64) >> 32); diff --git a/src/paimon/common/utils/bloom_filter64.h b/src/paimon/common/utils/bloom_filter64.h index 3ba4e86f9..2440b3033 100644 --- a/src/paimon/common/utils/bloom_filter64.h +++ b/src/paimon/common/utils/bloom_filter64.h @@ -22,6 +22,7 @@ #include #include "paimon/memory/bytes.h" +#include "paimon/result.h" #include "paimon/visibility.h" namespace paimon { @@ -31,7 +32,9 @@ class MemoryPool; /// Bloom filter 64 handle 64 bits hash. class PAIMON_EXPORT BloomFilter64 { public: - BloomFilter64(int64_t items, double fpp, const std::shared_ptr& pool); + static Result Create(int64_t items, double fpp, + const std::shared_ptr& pool); + class BitSet; BloomFilter64(int32_t num_hash_functions, std::unique_ptr&& bit_set); @@ -54,6 +57,8 @@ class PAIMON_EXPORT BloomFilter64 { void Set(int32_t index); bool Get(int32_t index) const; int32_t BitSize() const; + int32_t ByteLength() const; + void ToByteArray(int32_t offset, int32_t length, char* bytes) const; private: static constexpr int8_t MASK = 0x07; @@ -64,9 +69,11 @@ class PAIMON_EXPORT BloomFilter64 { }; private: + BloomFilter64(int32_t num_hash_functions, std::unique_ptr&& bit_set, + const std::shared_ptr& pool); + static constexpr int32_t BYTE_SIZE = 8; - private: int32_t num_bits_ = -1; int32_t num_hash_functions_ = -1; std::shared_ptr pool_; diff --git a/src/paimon/common/utils/bloom_filter64_test.cpp b/src/paimon/common/utils/bloom_filter64_test.cpp index aa548fb7d..6f60fd018 100644 --- a/src/paimon/common/utils/bloom_filter64_test.cpp +++ b/src/paimon/common/utils/bloom_filter64_test.cpp @@ -28,13 +28,14 @@ #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(BloomFilter64Test, TestSimple) { int32_t items = 10000; auto pool = GetDefaultPool(); - BloomFilter64 bloom_filter(items, 0.02, pool); + ASSERT_OK_AND_ASSIGN(BloomFilter64 bloom_filter, BloomFilter64::Create(items, 0.02, pool)); std::mt19937_64 engine(std::random_device{}()); // NOLINT(whitespace/braces) std::uniform_int_distribution distribution(std::numeric_limits::min(), std::numeric_limits::max()); @@ -61,6 +62,39 @@ TEST(BloomFilter64Test, TestSimple) { ASSERT_TRUE(static_cast(false_positives) / num < 0.03); } +TEST(BloomFilter64Test, TestInvalidItemsAndFpp) { + std::shared_ptr pool = GetDefaultPool(); + + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/0, /*fpp=*/0.1, pool), + "items must be greater than 0"); + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/-1, /*fpp=*/0.1, pool), + "items must be greater than 0"); + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/100, /*fpp=*/0.0, pool), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/100, /*fpp=*/-0.1, pool), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/100, /*fpp=*/1.0, pool), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter64::Create(/*items=*/100, /*fpp=*/1.1, pool), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG( + BloomFilter64::Create(/*items=*/std::numeric_limits::max(), /*fpp=*/0.1, pool), + "bloom filter size exceeds the supported range"); +} + +TEST(BloomFilter64Test, TestKeepMemoryPoolAlive) { + std::weak_ptr weak_pool; + { + std::shared_ptr pool(GetMemoryPool()); + weak_pool = pool; + ASSERT_OK_AND_ASSIGN(BloomFilter64 bloom_filter, + BloomFilter64::Create(/*items=*/100, /*fpp=*/0.1, pool)); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + } + ASSERT_TRUE(weak_pool.expired()); +} + TEST(BloomFilter64Test, TestCompatibleWithJava) { // data: -10, -5, 0, 13, 100, 200, 500 std::vector se_bytes = {241, 255, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0}; @@ -74,10 +108,10 @@ TEST(BloomFilter64Test, TestCompatibleWithJava) { ASSERT_TRUE(bloom_filter.TestHash(value)); } - BloomFilter64 bloom_filter2(10, 0.01, pool); + ASSERT_OK_AND_ASSIGN(BloomFilter64 bloom_filter2, BloomFilter64::Create(10, 0.01, pool)); ASSERT_EQ(7, bloom_filter2.GetNumHashFunctions()); - ASSERT_EQ(se_bytes.size() * BloomFilter64::BYTE_SIZE, bloom_filter2.num_bits_); - ASSERT_EQ(se_bytes.size(), bloom_filter2.GetBitSet().bytes_->size()); + ASSERT_EQ(se_bytes.size() * 8, bloom_filter2.GetBitSet().BitSize()); + ASSERT_EQ(se_bytes.size(), bloom_filter2.GetBitSet().ByteLength()); } } // namespace paimon::test diff --git a/src/paimon/common/utils/bloom_filter_test.cpp b/src/paimon/common/utils/bloom_filter_test.cpp index fdd6abb34..0832fdf5c 100644 --- a/src/paimon/common/utils/bloom_filter_test.cpp +++ b/src/paimon/common/utils/bloom_filter_test.cpp @@ -35,7 +35,8 @@ namespace paimon::test { TEST(BloomFilterTest, TestOneSegmentBuilder) { int32_t items = 100; auto pool = GetDefaultPool(); - auto bloom_filter = BloomFilter::Create(items, 0.01); + ASSERT_OK_AND_ASSIGN(std::shared_ptr bloom_filter, + BloomFilter::Create(items, 0.01)); auto seg = MemorySegment::AllocateHeapMemory(1024, pool.get()); ASSERT_OK(bloom_filter->SetMemorySegment(seg)); @@ -54,12 +55,32 @@ TEST(BloomFilterTest, TestOneSegmentBuilder) { } TEST(BloomFilterTest, TestEstimatedHashFunctions) { - ASSERT_EQ(7, BloomFilter::Create(1000, 0.01)->GetNumHashFunctions()); - ASSERT_EQ(7, BloomFilter::Create(10000, 0.01)->GetNumHashFunctions()); - ASSERT_EQ(7, BloomFilter::Create(100000, 0.01)->GetNumHashFunctions()); - ASSERT_EQ(4, BloomFilter::Create(100000, 0.05)->GetNumHashFunctions()); - ASSERT_EQ(7, BloomFilter::Create(1000000, 0.01)->GetNumHashFunctions()); - ASSERT_EQ(4, BloomFilter::Create(1000000, 0.05)->GetNumHashFunctions()); + auto get_num_hash_functions = [](int64_t expected_entries, double fpp) { + EXPECT_OK_AND_ASSIGN(std::shared_ptr bloom_filter, + BloomFilter::Create(expected_entries, fpp)); + return bloom_filter->GetNumHashFunctions(); + }; + ASSERT_EQ(7, get_num_hash_functions(1000, 0.01)); + ASSERT_EQ(7, get_num_hash_functions(10000, 0.01)); + ASSERT_EQ(7, get_num_hash_functions(100000, 0.01)); + ASSERT_EQ(4, get_num_hash_functions(100000, 0.05)); + ASSERT_EQ(7, get_num_hash_functions(1000000, 0.01)); + ASSERT_EQ(4, get_num_hash_functions(1000000, 0.05)); +} + +TEST(BloomFilterTest, TestInvalidExpectedEntriesAndFpp) { + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/0, /*fpp=*/0.1), + "expected entries must be greater than 0"); + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/-1, /*fpp=*/0.1), + "expected entries must be greater than 0"); + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/100, /*fpp=*/0.0), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/100, /*fpp=*/-0.1), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/100, /*fpp=*/1.0), + "fpp must be greater than 0 and less than 1"); + ASSERT_NOK_WITH_MSG(BloomFilter::Create(/*expected_entries=*/100, /*fpp=*/1.1), + "fpp must be greater than 0 and less than 1"); } TEST(BloomFilterTest, TestBloomNumBits) { diff --git a/src/paimon/common/utils/byte_range_combiner.cpp b/src/paimon/common/utils/byte_range_combiner.cpp index 428770037..c306c100f 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 599ca59cf..90942719b 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 19d739a41..5a3cc1c6e 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 722467699..d2786e26e 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 50f602375..c70edb092 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/fields_comparator.cpp b/src/paimon/common/utils/fields_comparator.cpp index 224ec131f..2fee9aa15 100644 --- a/src/paimon/common/utils/fields_comparator.cpp +++ b/src/paimon/common/utils/fields_comparator.cpp @@ -85,28 +85,28 @@ Result FieldsComparator::CompareField( arrow::Type::type type = input_type->id(); switch (type) { case arrow::Type::type::BOOL: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { bool lvalue = lhs.GetBoolean(field_idx); bool rvalue = rhs.GetBoolean(field_idx); return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); }); case arrow::Type::type::INT8: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { int8_t lvalue = lhs.GetByte(field_idx); int8_t rvalue = rhs.GetByte(field_idx); return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); }); case arrow::Type::type::INT16: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { int16_t lvalue = lhs.GetShort(field_idx); int16_t rvalue = rhs.GetShort(field_idx); return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); }); case arrow::Type::type::DATE32: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { int32_t lvalue = lhs.GetDate(field_idx); int32_t rvalue = rhs.GetDate(field_idx); @@ -114,39 +114,36 @@ Result FieldsComparator::CompareField( }); case arrow::Type::type::INT32: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { int32_t lvalue = lhs.GetInt(field_idx); int32_t rvalue = rhs.GetInt(field_idx); return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); }); case arrow::Type::type::INT64: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { int64_t lvalue = lhs.GetLong(field_idx); int64_t rvalue = rhs.GetLong(field_idx); return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); }); case arrow::Type::type::FLOAT: - // TODO(xinyu.lxy): - // currently in java KeyComparatorSupplier: -inf < -0.0 == +0.0 < +inf = nan - // paimon-cpp: -inf < -0.0 == +0.0 < +inf and nan cannot be compared - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { float lvalue = lhs.GetFloat(field_idx); float rvalue = rhs.GetFloat(field_idx); - return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); + return CompareFloatingPoint(lvalue, rvalue); }); case arrow::Type::type::DOUBLE: - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { double lvalue = lhs.GetDouble(field_idx); double rvalue = rhs.GetDouble(field_idx); - return lvalue == rvalue ? 0 : (lvalue < rvalue ? -1 : 1); + return CompareFloatingPoint(lvalue, rvalue); }); case arrow::Type::type::STRING: case arrow::Type::type::BINARY: { - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { auto lvalue = lhs.GetStringView(field_idx); auto rvalue = rhs.GetStringView(field_idx); @@ -157,7 +154,7 @@ Result FieldsComparator::CompareField( case arrow::Type::type::TIMESTAMP: { auto timestamp_type = checked_pointer_cast(input_type); int32_t precision = DateTimeUtils::GetPrecisionFromType(timestamp_type); - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx, precision](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { Timestamp lvalue = lhs.GetTimestamp(field_idx, precision); Timestamp rvalue = rhs.GetTimestamp(field_idx, precision); @@ -168,7 +165,7 @@ Result FieldsComparator::CompareField( auto* decimal_type = checked_cast(input_type.get()); auto precision = decimal_type->precision(); auto scale = decimal_type->scale(); - return FieldsComparator::FieldComparatorFunc( + return FieldComparatorFunc( [field_idx, precision, scale](const InternalRow& lhs, const InternalRow& rhs) -> int32_t { Decimal lvalue = lhs.GetDecimal(field_idx, precision, scale); diff --git a/src/paimon/common/utils/fields_comparator.h b/src/paimon/common/utils/fields_comparator.h index 7dbeca15e..50e0260d9 100644 --- a/src/paimon/common/utils/fields_comparator.h +++ b/src/paimon/common/utils/fields_comparator.h @@ -41,6 +41,9 @@ class DataField; /// A `Comparator` that compares the file store key. class FieldsComparator { public: + using FieldComparatorFunc = + std::function; + static Result> Create( const std::vector& input_data_field, bool is_ascending_order); @@ -82,9 +85,6 @@ class FieldsComparator { } private: - using FieldComparatorFunc = - std::function; - FieldsComparator(bool is_ascending_order, const std::vector& sort_fields, std::vector&& comparators) : is_ascending_order_(is_ascending_order), diff --git a/src/paimon/common/utils/fields_comparator_test.cpp b/src/paimon/common/utils/fields_comparator_test.cpp index 2f64f813c..111de404c 100644 --- a/src/paimon/common/utils/fields_comparator_test.cpp +++ b/src/paimon/common/utils/fields_comparator_test.cpp @@ -20,8 +20,10 @@ #include "paimon/common/utils/fields_comparator.h" #include +#include #include #include +#include #include "arrow/api.h" #include "gtest/gtest.h" @@ -80,6 +82,44 @@ class FieldsComparatorTest : public ::testing::Test { } CheckResult(row1, row2, input_types, sort_fields, has_null); } + + template + void CheckFloatingPointOrder(const std::shared_ptr& type) { + std::shared_ptr pool = GetDefaultPool(); + std::vector data_fields = {DataField(0, arrow::field("f0", type))}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr ascending_comparator, + FieldsComparator::Create(data_fields, + /*is_ascending_order=*/true)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr descending_comparator, + FieldsComparator::Create(data_fields, + /*is_ascending_order=*/false)); + + const T nan = std::numeric_limits::quiet_NaN(); + const std::vector values = {-std::numeric_limits::infinity(), + static_cast(-1), + static_cast(-0.0), + static_cast(0.0), + static_cast(1), + std::numeric_limits::infinity(), + nan}; + std::vector rows; + rows.reserve(values.size()); + for (T value : values) { + rows.emplace_back(BinaryRowGenerator::GenerateRow({value}, pool.get())); + } + + for (size_t i = 0; i < rows.size(); ++i) { + for (size_t j = 0; j < rows.size(); ++j) { + int32_t expected = i == j ? 0 : (i < j ? -1 : 1); + ASSERT_EQ(expected, ascending_comparator->CompareTo(rows[i], rows[j])); + ASSERT_EQ(-expected, descending_comparator->CompareTo(rows[i], rows[j])); + } + } + + BinaryRow negative_nan_row = BinaryRowGenerator::GenerateRow({-nan}, pool.get()); + ASSERT_EQ(0, ascending_comparator->CompareTo(rows.back(), negative_nan_row)); + ASSERT_EQ(0, ascending_comparator->CompareTo(negative_nan_row, rows.back())); + } }; TEST_F(FieldsComparatorTest, TestSimple) { @@ -202,6 +242,11 @@ TEST_F(FieldsComparatorTest, TestSimple) { } } +TEST_F(FieldsComparatorTest, TestFloatingPointOrder) { + CheckFloatingPointOrder(arrow::float32()); + CheckFloatingPointOrder(arrow::float64()); +} + TEST_F(FieldsComparatorTest, TestTimestampType) { auto pool = GetDefaultPool(); // test ts with different precision diff --git a/src/paimon/common/utils/http_client.cpp b/src/paimon/common/utils/http_client.cpp index e61c7faa8..1542f58a1 100644 --- a/src/paimon/common/utils/http_client.cpp +++ b/src/paimon/common/utils/http_client.cpp @@ -164,6 +164,9 @@ CurlHttpClient::~CurlHttpClient() = default; Result CurlHttpClient::Execute(const HttpRequest& request, const HttpBodyConsumer& consumer) const { + if (request.request_timeout_ms < 0) { + return Status::Invalid("HTTP request timeout must not be negative"); + } for (int32_t attempt = 0; attempt < kMaxAttempts; ++attempt) { CURL* handle = impl_->Acquire(); if (handle == nullptr) { @@ -184,6 +187,8 @@ Result CurlHttpClient::Execute(const HttpRequest& request, curl_easy_setopt(handle, CURLOPT_HTTPHEADER, headers); curl_easy_setopt(handle, CURLOPT_NOSIGNAL, 1L); curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, 30000L); + curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, + static_cast(request.request_timeout_ms)); // NOLINT(runtime/int) curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, WriteCallback); curl_easy_setopt(handle, CURLOPT_WRITEDATA, &context); curl_easy_setopt(handle, CURLOPT_HEADERFUNCTION, HeaderCallback); diff --git a/src/paimon/common/utils/http_client.h b/src/paimon/common/utils/http_client.h index 0dd3e742c..8b04a7003 100644 --- a/src/paimon/common/utils/http_client.h +++ b/src/paimon/common/utils/http_client.h @@ -46,6 +46,8 @@ struct HttpRequest { HttpMethod method = HttpMethod::GET; std::string url; HttpHeaders headers; + /// Overall request timeout in milliseconds; zero keeps libcurl's no-timeout default. + int32_t request_timeout_ms = 0; }; struct HttpResponse { diff --git a/src/paimon/common/utils/http_client_test.cpp b/src/paimon/common/utils/http_client_test.cpp index 5bf3ce232..f9cf82d96 100644 --- a/src/paimon/common/utils/http_client_test.cpp +++ b/src/paimon/common/utils/http_client_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/math.h b/src/paimon/common/utils/math.h index 54ad6cf73..b9e2e3103 100644 --- a/src/paimon/common/utils/math.h +++ b/src/paimon/common/utils/math.h @@ -28,6 +28,7 @@ #pragma once #include +#include #include #include #include @@ -36,10 +37,56 @@ #include "fmt/format.h" #include "paimon/common/utils/options_utils.h" +#include "paimon/io/byte_order.h" #include "paimon/status.h" namespace paimon { +inline constexpr uint32_t kCanonicalFloatNaNBits = 0x7fc00000; +inline constexpr uint64_t kCanonicalDoubleNaNBits = 0x7ff8000000000000; + +template +inline FloatingPoint FloatingPointFromBits(Bits bits) { + static_assert(std::is_floating_point_v); + static_assert(std::is_integral_v); + static_assert(sizeof(FloatingPoint) == sizeof(Bits)); + FloatingPoint value; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +inline float CanonicalizeFloatingPoint(float value) { + if (std::isnan(value)) { + return FloatingPointFromBits(kCanonicalFloatNaNBits); + } + return value; +} + +inline double CanonicalizeFloatingPoint(double value) { + if (std::isnan(value)) { + return FloatingPointFromBits(kCanonicalDoubleNaNBits); + } + return value; +} + +inline int32_t CanonicalizeFloatToIntBits(float value) { + if (std::isnan(value)) { + return static_cast(kCanonicalFloatNaNBits); + } + int32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +inline int64_t CanonicalizeDoubleToLongBits(double value) { + if (std::isnan(value)) { + return static_cast(kCanonicalDoubleNaNBits); + } + int64_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + template constexpr bool InRange(From value) { static_assert(std::is_integral_v && std::is_integral_v, @@ -136,4 +183,30 @@ inline T EndianSwapValue(T v) { } } +template +inline T ToBigEndian(T value) { + if constexpr (SystemByteOrder() == ByteOrder::PAIMON_LITTLE_ENDIAN) { + return EndianSwapValue(value); + } + return value; +} + +template +inline T ToLittleEndian(T value) { + if constexpr (SystemByteOrder() == ByteOrder::PAIMON_BIG_ENDIAN) { + return EndianSwapValue(value); + } + return value; +} + +template +inline T FromBigEndian(T value) { + return ToBigEndian(value); +} + +template +inline T FromLittleEndian(T value) { + return ToLittleEndian(value); +} + } // namespace paimon diff --git a/src/paimon/common/utils/math_test.cpp b/src/paimon/common/utils/math_test.cpp index 36df1cbab..ce9379f6c 100644 --- a/src/paimon/common/utils/math_test.cpp +++ b/src/paimon/common/utils/math_test.cpp @@ -19,6 +19,8 @@ #include "paimon/common/utils/math.h" +#include +#include #include #include "gtest/gtest.h" @@ -26,6 +28,29 @@ namespace paimon::test { +TEST(MathTest, FloatingPointNaNCanonicalization) { + const auto float_nan = CanonicalizeFloatingPoint(FloatingPointFromBits(0xffc12345U)); + uint32_t float_nan_bits; + std::memcpy(&float_nan_bits, &float_nan, sizeof(float_nan_bits)); + ASSERT_EQ(kCanonicalFloatNaNBits, float_nan_bits); + ASSERT_EQ(static_cast(kCanonicalFloatNaNBits), + CanonicalizeFloatToIntBits(FloatingPointFromBits(0x7fa12345U))); + + const auto double_nan = + CanonicalizeFloatingPoint(FloatingPointFromBits(0xfff8123456789abcULL)); + uint64_t double_nan_bits; + std::memcpy(&double_nan_bits, &double_nan, sizeof(double_nan_bits)); + ASSERT_EQ(kCanonicalDoubleNaNBits, double_nan_bits); + ASSERT_EQ(static_cast(kCanonicalDoubleNaNBits), + CanonicalizeDoubleToLongBits(FloatingPointFromBits(0x7ff123456789abcdULL))); + + const float negative_zero = CanonicalizeFloatingPoint(-0.0f); + uint32_t negative_zero_bits; + std::memcpy(&negative_zero_bits, &negative_zero, sizeof(negative_zero_bits)); + ASSERT_EQ(0x80000000U, negative_zero_bits); + ASSERT_EQ(0x3ff0000000000000, CanonicalizeDoubleToLongBits(1.0)); +} + // Test case: Test EndianSwapValue for different integral types TEST(MathTest, EndianSwapValue) { // Test 16-bit value @@ -44,6 +69,22 @@ TEST(MathTest, EndianSwapValue) { ASSERT_EQ(swapped64, 0xF0DEBC9A78563412); } +TEST(MathTest, ToEndian) { + constexpr uint32_t kValue = 0x12345678; + + const uint32_t big_endian = ToBigEndian(kValue); + std::array big_endian_bytes{}; + std::memcpy(big_endian_bytes.data(), &big_endian, sizeof(big_endian)); + ASSERT_EQ((std::array{0x12, 0x34, 0x56, 0x78}), big_endian_bytes); + ASSERT_EQ(kValue, FromBigEndian(big_endian)); + + const uint32_t little_endian = ToLittleEndian(kValue); + std::array little_endian_bytes{}; + std::memcpy(little_endian_bytes.data(), &little_endian, sizeof(little_endian)); + ASSERT_EQ((std::array{0x78, 0x56, 0x34, 0x12}), little_endian_bytes); + ASSERT_EQ(kValue, FromLittleEndian(little_endian)); +} + TEST(MathTest, InRange) { // signed -> unsigned: negative values out of range, boundary values in range ASSERT_TRUE(InRange(0)); diff --git a/src/paimon/common/utils/options_utils.h b/src/paimon/common/utils/options_utils.h index 90b30b54b..08b09ad2d 100644 --- a/src/paimon/common/utils/options_utils.h +++ b/src/paimon/common/utils/options_utils.h @@ -76,13 +76,42 @@ class OptionsUtils { return value.value(); } + template + static Result> GetOptionalValueFromMap( + const std::map& key_value_map, const std::string& key) { + Result value = GetValueFromMap(key_value_map, key); + if (value.ok()) { + return std::optional(value.value()); + } + if (value.status().IsNotExist()) { + return std::optional(); + } + return value.status(); + } + + static Result GetNonEmptyValueFromMap( + const std::map& key_value_map, const std::string& key) { + Result value = GetValueFromMap(key_value_map, key); + if (!value.ok()) { + return value.status(); + } + if (value.value().empty()) { + return Status::Invalid(fmt::format("value for key {} must not be empty", key)); + } + return value.value(); + } + /// Fetch options with specific prefix and remove prefix for key. + /// @param prefix Prefix used to select options and removed from the returned keys. + /// @param options Options to select from. + /// @return Options whose keys start with and are longer than `prefix`, with the prefix removed + /// from each key. static std::map FetchOptionsWithPrefix( const std::string& prefix, const std::map& options) { std::map options_with_prefix; - int64_t prefix_len = prefix.size(); + const std::string::size_type prefix_len = prefix.size(); for (const auto& [key, value] : options) { - if (StringUtils::StartsWith(key, prefix)) { + if (key.size() > prefix_len && StringUtils::StartsWith(key, prefix)) { options_with_prefix[key.substr(prefix_len)] = value; } } diff --git a/src/paimon/common/utils/options_utils_test.cpp b/src/paimon/common/utils/options_utils_test.cpp index 7a09a9850..61e520874 100644 --- a/src/paimon/common/utils/options_utils_test.cpp +++ b/src/paimon/common/utils/options_utils_test.cpp @@ -63,10 +63,40 @@ TEST(OptionsUtilsTest, TestGetValueFromMap) { ASSERT_EQ(999, empty); } +TEST(OptionsUtilsTest, TestGetOptionalValueFromMap) { + const std::map key_value_map = { + {"key_int", "10"}, {"key_empty", ""}, {"key_invalid", "ab"}}; + + ASSERT_OK_AND_ASSIGN(std::optional optional_value, + OptionsUtils::GetOptionalValueFromMap(key_value_map, "key_int")); + ASSERT_EQ(std::optional(10), optional_value); + ASSERT_OK_AND_ASSIGN( + std::optional optional_missing, + OptionsUtils::GetOptionalValueFromMap(key_value_map, "key_nonexist")); + ASSERT_EQ(std::nullopt, optional_missing); + ASSERT_OK_AND_ASSIGN( + std::optional optional_empty, + OptionsUtils::GetOptionalValueFromMap(key_value_map, "key_empty")); + ASSERT_EQ(std::optional(""), optional_empty); + ASSERT_TRUE(OptionsUtils::GetOptionalValueFromMap(key_value_map, "key_invalid") + .status() + .IsInvalid()); +} + TEST(OptionsUtilsTest, TestFetchOptionsWithPrefix) { - std::map options = {{"key1", "value1"}, {"test.key2", "value2"}}; + std::map options = { + {"key1", "value1"}, {"test.", "empty-key"}, {"test.key2", "value2"}}; auto new_options = OptionsUtils::FetchOptionsWithPrefix("test.", options); std::map expected = {{"key2", "value2"}}; ASSERT_EQ(expected, new_options); } + +TEST(OptionsUtilsTest, TestGetNonEmptyValueFromMap) { + std::map options = {{"present", "value"}, {"empty", ""}}; + ASSERT_OK_AND_ASSIGN(std::string value, + OptionsUtils::GetNonEmptyValueFromMap(options, "present")); + ASSERT_EQ("value", value); + ASSERT_TRUE(OptionsUtils::GetNonEmptyValueFromMap(options, "missing").status().IsNotExist()); + ASSERT_TRUE(OptionsUtils::GetNonEmptyValueFromMap(options, "empty").status().IsInvalid()); +} } // namespace paimon::test diff --git a/src/paimon/common/utils/read_ahead_cache.cpp b/src/paimon/common/utils/read_ahead_cache.cpp index b0001189b..a74b35202 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 196045b7e..f039c22e8 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 676254220..e7900b7b6 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 000000000..cb9c6039f --- /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 000000000..cc6b74cf1 --- /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/sensitive_config_utils.cpp b/src/paimon/common/utils/sensitive_config_utils.cpp index 28740412c..a3cff3c66 100644 --- a/src/paimon/common/utils/sensitive_config_utils.cpp +++ b/src/paimon/common/utils/sensitive_config_utils.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/sensitive_config_utils.h b/src/paimon/common/utils/sensitive_config_utils.h index 717156098..05cc8ac63 100644 --- a/src/paimon/common/utils/sensitive_config_utils.h +++ b/src/paimon/common/utils/sensitive_config_utils.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/sensitive_config_utils_test.cpp b/src/paimon/common/utils/sensitive_config_utils_test.cpp index 5e8b5350b..9e8989f76 100644 --- a/src/paimon/common/utils/sensitive_config_utils_test.cpp +++ b/src/paimon/common/utils/sensitive_config_utils_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/serialization_utils.h b/src/paimon/common/utils/serialization_utils.h index c1e97d033..449766ca6 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 5e612ff2e..56a39a295 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/common/utils/string_utils.cpp b/src/paimon/common/utils/string_utils.cpp index 5b405895d..869e184d0 100644 --- a/src/paimon/common/utils/string_utils.cpp +++ b/src/paimon/common/utils/string_utils.cpp @@ -30,8 +30,28 @@ #include "paimon/status.h" namespace paimon { +namespace { + +bool IsTrimCharacter(unsigned char c) { + // Match the characters removed by Java String::trim for the ASCII strings handled here. + return c <= 0x20; +} + +char ToAsciiLower(unsigned char c) { + return c >= 'A' && c <= 'Z' ? static_cast(c + ('a' - 'A')) : static_cast(c); +} + +char ToAsciiUpper(unsigned char c) { + return c >= 'a' && c <= 'z' ? static_cast(c - ('a' - 'A')) : static_cast(c); +} + +} // namespace + std::string StringUtils::Replace(const std::string& text, const std::string& search_string, const std::string& replacement, int32_t max) { + if (text.empty() || search_string.empty() || max == 0) { + return text; + } std::string str = text; size_t pos = str.find(search_string); int32_t count = 0; @@ -45,6 +65,9 @@ std::string StringUtils::Replace(const std::string& text, const std::string& sea std::string StringUtils::ReplaceLast(const std::string& text, const std::string& old_str, const std::string& new_str) { + if (text.empty() || old_str.empty()) { + return text; + } std::string str = text; size_t pos = str.rfind(old_str); if (pos != std::string::npos) { @@ -54,7 +77,8 @@ std::string StringUtils::ReplaceLast(const std::string& text, const std::string& } bool StringUtils::StartsWith(const std::string& str, const std::string& prefix, size_t start_pos) { - return (str.size() >= prefix.size()) && (str.compare(start_pos, prefix.size(), prefix) == 0); + return start_pos <= str.size() && prefix.size() <= str.size() - start_pos && + str.compare(start_pos, prefix.size(), prefix) == 0; } bool StringUtils::EndsWith(const std::string& str, const std::string& suffix) { size_t s1 = str.size(); @@ -74,26 +98,45 @@ bool StringUtils::IsNullOrWhitespaceOnly(const std::string& str) { } void StringUtils::Trim(std::string* str) { - str->erase(str->find_last_not_of(' ') + 1); - str->erase(0, str->find_first_not_of(' ')); + auto first = std::find_if_not(str->begin(), str->end(), + [](unsigned char c) { return IsTrimCharacter(c); }); + auto last = std::find_if_not(str->rbegin(), str->rend(), [](unsigned char c) { + return IsTrimCharacter(c); + }).base(); + if (first >= last) { + str->clear(); + return; + } + *str = std::string(first, last); } std::string StringUtils::ToLowerCase(const std::string& str) { std::string result; result.reserve(str.length()); - std::transform(str.begin(), str.end(), std::back_inserter(result), - [](unsigned char c) { return std::tolower(c); }); + std::transform(str.begin(), str.end(), std::back_inserter(result), ToAsciiLower); return result; } std::string StringUtils::ToUpperCase(const std::string& str) { std::string result; result.reserve(str.length()); - std::transform(str.begin(), str.end(), std::back_inserter(result), - [](unsigned char c) { return std::toupper(c); }); + std::transform(str.begin(), str.end(), std::back_inserter(result), ToAsciiUpper); return result; } +bool StringUtils::EqualsIgnoreCase(const std::string& left, const std::string& right) { + if (left.size() != right.size()) { + return false; + } + for (size_t i = 0; i < left.size(); ++i) { + if (ToAsciiLower(static_cast(left[i])) != + ToAsciiLower(static_cast(right[i]))) { + return false; + } + } + return true; +} + std::vector StringUtils::Split(const std::string& text, const std::string& sep_str, bool ignore_empty) { std::vector vec; diff --git a/src/paimon/common/utils/string_utils.h b/src/paimon/common/utils/string_utils.h index 3c0906e2e..7681a8d9e 100644 --- a/src/paimon/common/utils/string_utils.h +++ b/src/paimon/common/utils/string_utils.h @@ -50,24 +50,18 @@ class PAIMON_EXPORT StringUtils { public: /// Replaces all occurrences of a string within another string. /// - /// A `null` reference passed to this method is a no-op. - /// /// 
-    /// StringUtils::Replace(null, *, *)        = null
     /// StringUtils::Replace("", *, *)          = ""
-    /// StringUtils::Replace("any", null, *)    = "any"
-    /// StringUtils::Replace("any", *, null)    = "any"
     /// StringUtils::Replace("any", "", *)      = "any"
-    /// StringUtils::Replace("aba", "a", null)  = "aba"
     /// StringUtils::Replace("aba", "a", "")    = "b"
     /// StringUtils::Replace("aba", "a", "z")   = "zbz"
     ///
/// /// @see #replace(string text, string search_string, string replacement, int max) - /// @param text text to search and replace in, may be null - /// @param search_string the String to search for, may be null - /// @param replacement the String to replace it with, may be null - /// @return the text with any replacements processed, `null` if null string input + /// @param text text to search and replace in + /// @param search_string the String to search for + /// @param replacement the String to replace it with + /// @return the text with any replacements processed static std::string Replace(const std::string& text, const std::string& search_string, const std::string& replacement) { return Replace(text, search_string, replacement, -1); @@ -76,16 +70,10 @@ class PAIMON_EXPORT StringUtils { /// Replaces a String with another String inside a larger String, for the first `max` values of /// the search String. /// - /// A `null` reference passed to this method is a no-op. - /// /// 
-    /// StringUtils::Replace(null, *, *, *)         = null
     /// StringUtils::Replace("", *, *, *)           = ""
-    /// StringUtils::Replace("any", null, *, *)     = "any"
-    /// StringUtils::Replace("any", *, null, *)     = "any"
     /// StringUtils::Replace("any", "", *, *)       = "any"
     /// StringUtils::Replace("any", *, *, 0)        = "any"
-    /// StringUtils::Replace("abaa", "a", null, -1) = "abaa"
     /// StringUtils::Replace("abaa", "a", "", -1)   = "b"
     /// StringUtils::Replace("abaa", "a", "z", 0)   = "abaa"
     /// StringUtils::Replace("abaa", "a", "z", 1)   = "zbaa"
@@ -93,11 +81,11 @@ class PAIMON_EXPORT StringUtils {
     /// StringUtils::Replace("abaa", "a", "z", -1)  = "zbzz"
     ///
/// - /// @param text text to search and replace in, may be null - /// @param search_string the String to search for, may be null - /// @param replacement the String to replace it with, may be null + /// @param text text to search and replace in + /// @param search_string the String to search for + /// @param replacement the String to replace it with /// @param max maximum number of values to replace, or `-1` if no maximum - /// @return the text with any replacements processed, `null` if null string input + /// @return the text with any replacements processed static std::string Replace(const std::string& text, const std::string& search_string, const std::string& replacement, int32_t max); @@ -115,6 +103,9 @@ class PAIMON_EXPORT StringUtils { static std::string ToLowerCase(const std::string& str); static std::string ToUpperCase(const std::string& str); + /// Compares two strings using ASCII case folding. + static bool EqualsIgnoreCase(const std::string& left, const std::string& right); + template static std::string VectorToString(const std::vector& vec) { std::vector strs; diff --git a/src/paimon/common/utils/string_utils_test.cpp b/src/paimon/common/utils/string_utils_test.cpp index 11c3e0005..a4f230781 100644 --- a/src/paimon/common/utils/string_utils_test.cpp +++ b/src/paimon/common/utils/string_utils_test.cpp @@ -73,6 +73,8 @@ void StringUtilsTest::CheckOverFlowAndUnderFlow(const std::string& over_flow, } TEST_F(StringUtilsTest, TestReplaceAll) { + ASSERT_EQ("abc", StringUtils::Replace("abc", "", "x")); + ASSERT_EQ("", StringUtils::Replace("", "a", "b")); { std::string origin = "how is is you"; std::string expect = "how are are you"; @@ -118,6 +120,8 @@ TEST_F(StringUtilsTest, TestReplaceAll) { } TEST_F(StringUtilsTest, TestReplaceLast) { + ASSERT_EQ("abc", StringUtils::ReplaceLast("abc", "", "x")); + ASSERT_EQ("", StringUtils::ReplaceLast("", "a", "b")); { std::string origin = "a/b/c//"; std::string expect = "a/b/c/_"; @@ -140,6 +144,7 @@ TEST_F(StringUtilsTest, TestReplaceLast) { } TEST_F(StringUtilsTest, TestReplaceWithMaxCount) { + ASSERT_EQ("abc", StringUtils::Replace("abc", "a", "b", 0)); { std::string origin = "how is is you"; std::string expect = "how are is you"; @@ -236,6 +241,13 @@ TEST_F(StringUtilsTest, TestToUpperCase) { } } +TEST_F(StringUtilsTest, TestEqualsIgnoreCase) { + ASSERT_TRUE(StringUtils::EqualsIgnoreCase("", "")); + ASSERT_TRUE(StringUtils::EqualsIgnoreCase("AbC-123", "aBc-123")); + ASSERT_FALSE(StringUtils::EqualsIgnoreCase("abc", "abcd")); + ASSERT_FALSE(StringUtils::EqualsIgnoreCase("abc", "abx")); +} + TEST_F(StringUtilsTest, TestStartsWith) { { std::string str = "abcde"; @@ -261,6 +273,26 @@ TEST_F(StringUtilsTest, TestStartsWith) { std::string str = ""; ASSERT_TRUE(StringUtils::StartsWith(str, "")); } + { + std::string str = "abc"; + ASSERT_TRUE(StringUtils::StartsWith(str, "", /*start_pos=*/3)); + ASSERT_FALSE(StringUtils::StartsWith(str, "", /*start_pos=*/4)); + ASSERT_FALSE(StringUtils::StartsWith(str, "a", /*start_pos=*/4)); + } +} + +TEST_F(StringUtilsTest, TestTrim) { + std::string value = " \tabc\r\n"; + StringUtils::Trim(&value); + ASSERT_EQ("abc", value); + + value = "\t\r\n"; + StringUtils::Trim(&value); + ASSERT_TRUE(value.empty()); + + value.clear(); + StringUtils::Trim(&value); + ASSERT_TRUE(value.empty()); } TEST_F(StringUtilsTest, TestEndsWith) { { diff --git a/src/paimon/common/utils/url_utils.cpp b/src/paimon/common/utils/url_utils.cpp index 883a4b77f..8d9bfd536 100644 --- a/src/paimon/common/utils/url_utils.cpp +++ b/src/paimon/common/utils/url_utils.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/url_utils.h b/src/paimon/common/utils/url_utils.h index a045dc067..2fedbe73c 100644 --- a/src/paimon/common/utils/url_utils.h +++ b/src/paimon/common/utils/url_utils.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/common/utils/url_utils_test.cpp b/src/paimon/common/utils/url_utils_test.cpp index ebdc2bdc7..56b0620dc 100644 --- a/src/paimon/common/utils/url_utils_test.cpp +++ b/src/paimon/common/utils/url_utils_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/core/append/append_only_writer_test.cpp b/src/paimon/core/append/append_only_writer_test.cpp index 8ef6e135f..114016091 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/append/bucketed_append_compact_manager_test.cpp b/src/paimon/core/append/bucketed_append_compact_manager_test.cpp index ee56825b5..712dd3667 100644 --- a/src/paimon/core/append/bucketed_append_compact_manager_test.cpp +++ b/src/paimon/core/append/bucketed_append_compact_manager_test.cpp @@ -115,7 +115,7 @@ class BucketedAppendCompactManagerTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr NewNamedFile(const std::string& file_name, int64_t file_size, @@ -135,7 +135,7 @@ class BucketedAppendCompactManagerTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateTestDvMaintainer( diff --git a/src/paimon/core/bucket/hive_bucket_function.cpp b/src/paimon/core/bucket/hive_bucket_function.cpp index 913053c1f..e87c292d9 100644 --- a/src/paimon/core/bucket/hive_bucket_function.cpp +++ b/src/paimon/core/bucket/hive_bucket_function.cpp @@ -19,13 +19,12 @@ #include "paimon/core/bucket/hive_bucket_function.h" #include -#include -#include #include #include "fmt/format.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/core/bucket/hive_hasher.h" #include "paimon/status.h" @@ -105,10 +104,8 @@ uint32_t HiveBucketFunction::ComputeHash(const BinaryRow& row, int32_t field_ind uint32_t bits; if (float_value == -0.0f) { bits = 0; - } else if (std::isnan(float_value)) { - bits = 0x7FC00000U; } else { - std::memcpy(&bits, &float_value, sizeof(bits)); + bits = static_cast(CanonicalizeFloatToIntBits(float_value)); } return HiveHasher::HashInt(bits); } @@ -117,10 +114,8 @@ uint32_t HiveBucketFunction::ComputeHash(const BinaryRow& row, int32_t field_ind uint64_t bits; if (double_value == -0.0) { bits = 0; - } else if (std::isnan(double_value)) { - bits = 0x7FF8000000000000ULL; } else { - std::memcpy(&bits, &double_value, sizeof(bits)); + bits = static_cast(CanonicalizeDoubleToLongBits(double_value)); } return HiveHasher::HashLong(bits); } diff --git a/src/paimon/core/bucket/hive_bucket_function_test.cpp b/src/paimon/core/bucket/hive_bucket_function_test.cpp index 21f2a9843..d97a0294b 100644 --- a/src/paimon/core/bucket/hive_bucket_function_test.cpp +++ b/src/paimon/core/bucket/hive_bucket_function_test.cpp @@ -18,12 +18,12 @@ #include "paimon/core/bucket/hive_bucket_function.h" -#include #include #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/data/binary_row_writer.h" +#include "paimon/common/utils/math.h" #include "paimon/core/bucket/hive_hasher.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/binary_row_generator.h" @@ -111,18 +111,6 @@ class HiveBucketFunctionTest : public ::testing::Test { auto pool = GetDefaultPool(); return BinaryRowGenerator::GenerateRow({value}, pool.get()); } - - float FloatFromBits(uint32_t bits) { - float value; - std::memcpy(&value, &bits, sizeof(value)); - return value; - } - - double DoubleFromBits(uint64_t bits) { - double value; - std::memcpy(&value, &bits, sizeof(value)); - return value; - } }; /// Test matching Java: testHiveBucketFunction @@ -235,11 +223,12 @@ TEST_F(HiveBucketFunctionTest, TestFloatNaNCanonicalizationCompatibleWithJava) { ASSERT_OK_AND_ASSIGN(auto func, HiveBucketFunction::Create(field_types)); // Verified with Java HiveBucketFunction: - // Float.NaN, Float.intBitsToFloat(0x7fa12345), and Float.intBitsToFloat(0x7fc00000) - // all hash through Float.floatToIntBits(...) = 0x7fc00000. + // Float.NaN, a payload NaN, and the canonical NaN all hash through + // Float.floatToIntBits(...) to kCanonicalFloatNaNBits. ASSERT_EQ(344, func->Bucket(CreateFloatRow(std::numeric_limits::quiet_NaN()), 1000)); - ASSERT_EQ(344, func->Bucket(CreateFloatRow(FloatFromBits(0x7FA12345U)), 1000)); - ASSERT_EQ(344, func->Bucket(CreateFloatRow(FloatFromBits(0x7FC00000U)), 1000)); + ASSERT_EQ(344, func->Bucket(CreateFloatRow(FloatingPointFromBits(0x7FA12345U)), 1000)); + ASSERT_EQ(344, func->Bucket( + CreateFloatRow(FloatingPointFromBits(kCanonicalFloatNaNBits)), 1000)); } TEST_F(HiveBucketFunctionTest, TestDoubleNaNCanonicalizationCompatibleWithJava) { @@ -248,10 +237,14 @@ TEST_F(HiveBucketFunctionTest, TestDoubleNaNCanonicalizationCompatibleWithJava) // Verified with Java HiveBucketFunction: // Double.NaN, Double.longBitsToDouble(0x7ff123456789abcd), and canonical NaN - // all hash through Double.doubleToLongBits(...) = 0x7ff8000000000000. + // All NaNs hash through Double.doubleToLongBits(...) to kCanonicalDoubleNaNBits. ASSERT_EQ(360, func->Bucket(CreateDoubleRow(std::numeric_limits::quiet_NaN()), 1000)); - ASSERT_EQ(360, func->Bucket(CreateDoubleRow(DoubleFromBits(0x7FF123456789ABCDULL)), 1000)); - ASSERT_EQ(360, func->Bucket(CreateDoubleRow(DoubleFromBits(0x7FF8000000000000ULL)), 1000)); + ASSERT_EQ( + 360, + func->Bucket(CreateDoubleRow(FloatingPointFromBits(0x7FF123456789ABCDULL)), 1000)); + ASSERT_EQ(360, + func->Bucket(CreateDoubleRow(FloatingPointFromBits(kCanonicalDoubleNaNBits)), + 1000)); } TEST_F(HiveBucketFunctionTest, TestTinyintNegativeValuesCompatibleWithJava) { diff --git a/src/paimon/core/catalog/catalog_utils.cpp b/src/paimon/core/catalog/catalog_utils.cpp index 7a6768ab3..eae23ad8e 100644 --- a/src/paimon/core/catalog/catalog_utils.cpp +++ b/src/paimon/core/catalog/catalog_utils.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/core/catalog/catalog_utils.h b/src/paimon/core/catalog/catalog_utils.h index b9a54ac11..e92cd8443 100644 --- a/src/paimon/core/catalog/catalog_utils.h +++ b/src/paimon/core/catalog/catalog_utils.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index 578950fbe..71ba73deb 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -19,6 +19,7 @@ #include "paimon/core/core_options.h" #include +#include #include #include #include @@ -38,6 +39,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 { @@ -50,14 +52,9 @@ class ConfigParser { // Parse basic type configurations template Status Parse(const std::string& key, T* value) const { - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - auto result = StringUtils::StringToValue(iter->second); - if (result) { - *value = result.value(); - return Status::OK(); - } - return Status::Invalid(fmt::format("Invalid Config [{}: {}]", key, iter->second)); + PAIMON_ASSIGN_OR_RAISE(std::optional parsed_value, GetOptionalValue(key)); + if (parsed_value) { + *value = parsed_value.value(); } return Status::OK(); // Return success even if the configuration does not exist } @@ -65,14 +62,9 @@ class ConfigParser { // Parse optional basic type configurations template Status Parse(const std::string& key, std::optional* value) const { - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - auto result = StringUtils::StringToValue(iter->second); - if (result) { - *value = result.value(); - return Status::OK(); - } - return Status::Invalid(fmt::format("Invalid Config [{}: {}]", key, iter->second)); + PAIMON_ASSIGN_OR_RAISE(std::optional parsed_value, GetOptionalValue(key)); + if (parsed_value) { + *value = parsed_value.value(); } return Status::OK(); // Return success even if the configuration does not exist } @@ -81,23 +73,26 @@ class ConfigParser { template Status ParseList(const std::string& key, const std::string& delimiter, std::vector* list, bool need_trim = false) const { - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - auto value_str_vec = StringUtils::Split(iter->second, delimiter, /*ignore_empty=*/true); - for (auto& value_str : value_str_vec) { - if (need_trim) { - StringUtils::Trim(&value_str); - } - if constexpr (std::is_same_v) { - list->emplace_back(value_str); - } else { - auto value = StringUtils::StringToValue(value_str); - if (!value) { - return Status::Invalid( - fmt::format("Invalid Config [{}: {}]", key, iter->second)); - } - list->emplace_back(value.value()); + PAIMON_ASSIGN_OR_RAISE(std::optional config_value, + GetOptionalValue(key)); + if (!config_value) { + return Status::OK(); + } + auto value_str_vec = + StringUtils::Split(config_value.value(), delimiter, /*ignore_empty=*/true); + for (auto& value_str : value_str_vec) { + if (need_trim) { + StringUtils::Trim(&value_str); + } + if constexpr (std::is_same_v) { + list->emplace_back(value_str); + } else { + auto value = StringUtils::StringToValue(value_str); + if (!value) { + return Status::Invalid( + fmt::format("Invalid Config [{}: {}]", key, config_value.value())); } + list->emplace_back(value.value()); } } return Status::OK(); // Return success even if the configuration does not exist @@ -108,9 +103,10 @@ class ConfigParser { Status ParseMemorySize(const std::string& key, T* value) const { static_assert(std::is_same_v || std::is_same_v>, "ParseMemorySize only supports int64_t and std::optional"); - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - PAIMON_ASSIGN_OR_RAISE(*value, MemorySize::ParseBytes(iter->second)); + PAIMON_ASSIGN_OR_RAISE(std::optional config_value, + GetOptionalValue(key)); + if (config_value) { + PAIMON_ASSIGN_OR_RAISE(*value, MemorySize::ParseBytes(config_value.value())); } return Status::OK(); } @@ -120,9 +116,10 @@ class ConfigParser { Status ParseTimeDuration(const std::string& key, T* value) const { static_assert(std::is_same_v || std::is_same_v>, "ParseTimeDuration only supports int64_t and std::optional"); - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - PAIMON_ASSIGN_OR_RAISE(*value, TimeDuration::Parse(iter->second)); + PAIMON_ASSIGN_OR_RAISE(std::optional config_value, + GetOptionalValue(key)); + if (config_value) { + PAIMON_ASSIGN_OR_RAISE(*value, TimeDuration::Parse(config_value.value())); } return Status::OK(); } @@ -131,14 +128,10 @@ class ConfigParser { template Status ParseObject(const std::string& key, const std::string& default_identifier, std::shared_ptr* value) const { - auto iter = config_map_.find(key); - if (iter != config_map_.end()) { - std::string normalized_value = StringUtils::ToLowerCase(iter->second); - PAIMON_ASSIGN_OR_RAISE(*value, Factory::Get(normalized_value, config_map_)); - } else { - PAIMON_ASSIGN_OR_RAISE( - *value, Factory::Get(StringUtils::ToLowerCase(default_identifier), config_map_)); - } + PAIMON_ASSIGN_OR_RAISE(std::string identifier, OptionsUtils::GetValueFromMap( + config_map_, key, default_identifier)); + PAIMON_ASSIGN_OR_RAISE(*value, + Factory::Get(StringUtils::ToLowerCase(identifier), config_map_)); return Status::OK(); } @@ -151,11 +144,10 @@ class ConfigParser { *value = specified_file_system; return Status::OK(); } - std::string default_fs_identifier = "local"; - auto iter = config_map_.find(Options::FILE_SYSTEM); - if (iter != config_map_.end()) { - default_fs_identifier = StringUtils::ToLowerCase(iter->second); - } + PAIMON_ASSIGN_OR_RAISE( + std::string default_fs_identifier, + OptionsUtils::GetValueFromMap(config_map_, Options::FILE_SYSTEM, "local")); + default_fs_identifier = StringUtils::ToLowerCase(default_fs_identifier); *value = std::make_shared(fs_scheme_to_identifier_map, default_fs_identifier, config_map_); return Status::OK(); @@ -163,151 +155,81 @@ class ConfigParser { // Parse SortOrder Status ParseSortOrder(SortOrder* sort_order) const { - auto iter = config_map_.find(Options::SEQUENCE_FIELD_SORT_ORDER); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "ascending") { - *sort_order = SortOrder::ASCENDING; - } else if (str == "descending") { - *sort_order = SortOrder::DESCENDING; - } else { - return Status::Invalid(fmt::format("invalid sort order: {}", str)); - } - } - return Status::OK(); + return ParseEnum( + Options::SEQUENCE_FIELD_SORT_ORDER, + {{"ascending", SortOrder::ASCENDING}, {"descending", SortOrder::DESCENDING}}, + "sort order", sort_order); } // Parse LookupCompactMode Status ParseLookupCompactMode(LookupCompactMode* mode) const { - auto iter = config_map_.find(Options::LOOKUP_COMPACT); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "radical") { - *mode = LookupCompactMode::RADICAL; - } else if (str == "gentle") { - *mode = LookupCompactMode::GENTLE; - } else { - return Status::Invalid(fmt::format("invalid lookup mode: {}", str)); - } - } - return Status::OK(); + return ParseEnum( + Options::LOOKUP_COMPACT, + {{"radical", LookupCompactMode::RADICAL}, {"gentle", LookupCompactMode::GENTLE}}, + "lookup mode", mode); } // Parse SortEngine Status ParseSortEngine(SortEngine* sort_engine) const { - auto iter = config_map_.find(Options::SORT_ENGINE); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "min-heap") { - *sort_engine = SortEngine::MIN_HEAP; - } else if (str == "loser-tree") { - *sort_engine = SortEngine::LOSER_TREE; - } else { - return Status::Invalid(fmt::format("invalid sort engine: {}", str)); - } - } - return Status::OK(); + return ParseEnum( + Options::SORT_ENGINE, + {{"min-heap", SortEngine::MIN_HEAP}, {"loser-tree", SortEngine::LOSER_TREE}}, + "sort engine", sort_engine); } // Parse MergeEngine Status ParseMergeEngine(MergeEngine* merge_engine) const { - auto iter = config_map_.find(Options::MERGE_ENGINE); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "deduplicate") { - *merge_engine = MergeEngine::DEDUPLICATE; - } else if (str == "partial-update") { - *merge_engine = MergeEngine::PARTIAL_UPDATE; - } else if (str == "aggregation") { - *merge_engine = MergeEngine::AGGREGATE; - } else if (str == "first-row") { - *merge_engine = MergeEngine::FIRST_ROW; - } else { - return Status::Invalid(fmt::format("invalid merge engine: {}", str)); - } - } - return Status::OK(); + return ParseEnum(Options::MERGE_ENGINE, + {{"deduplicate", MergeEngine::DEDUPLICATE}, + {"partial-update", MergeEngine::PARTIAL_UPDATE}, + {"aggregation", MergeEngine::AGGREGATE}, + {"first-row", MergeEngine::FIRST_ROW}}, + "merge engine", merge_engine); } // Parse VariantShreddingInferenceMode Status ParseVariantShreddingInferenceMode(VariantShreddingInferenceMode* inference_mode) const { - auto iter = config_map_.find(Options::VARIANT_SHREDDING_INFERENCE_MODE); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "per-file") { - *inference_mode = VariantShreddingInferenceMode::PER_FILE; - } else if (str == "adaptive") { - *inference_mode = VariantShreddingInferenceMode::ADAPTIVE; - } else { - return Status::Invalid( - fmt::format("invalid variant shredding inference mode: {}", str)); - } - } - return Status::OK(); + return ParseEnum(Options::VARIANT_SHREDDING_INFERENCE_MODE, + {{"per-file", VariantShreddingInferenceMode::PER_FILE}, + {"adaptive", VariantShreddingInferenceMode::ADAPTIVE}}, + "variant shredding inference mode", inference_mode); } // Parse ChangelogProducer Status ParseChangelogProducer(ChangelogProducer* changelog_producer) const { - auto iter = config_map_.find(Options::CHANGELOG_PRODUCER); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "none") { - *changelog_producer = ChangelogProducer::NONE; - } else if (str == "input") { - *changelog_producer = ChangelogProducer::INPUT; - } else if (str == "full-compaction") { - *changelog_producer = ChangelogProducer::FULL_COMPACTION; - } else if (str == "lookup") { - *changelog_producer = ChangelogProducer::LOOKUP; - } else { - return Status::Invalid(fmt::format("invalid changelog producer: {}", str)); - } - } - return Status::OK(); + return ParseEnum(Options::CHANGELOG_PRODUCER, + {{"none", ChangelogProducer::NONE}, + {"input", ChangelogProducer::INPUT}, + {"full-compaction", ChangelogProducer::FULL_COMPACTION}, + {"lookup", ChangelogProducer::LOOKUP}}, + "changelog producer", changelog_producer); } // Parse ExternalPathStrategy Status ParseExternalPathStrategy(ExternalPathStrategy* external_path_strategy) const { - auto iter = config_map_.find(Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "none") { - *external_path_strategy = ExternalPathStrategy::NONE; - } else if (str == "specific-fs") { - *external_path_strategy = ExternalPathStrategy::SPECIFIC_FS; - } else if (str == "round-robin") { - *external_path_strategy = ExternalPathStrategy::ROUND_ROBIN; - } else { - return Status::Invalid(fmt::format("invalid external path strategy: {}", str)); - } - } - return Status::OK(); + return ParseEnum(Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY, + {{"none", ExternalPathStrategy::NONE}, + {"specific-fs", ExternalPathStrategy::SPECIFIC_FS}, + {"round-robin", ExternalPathStrategy::ROUND_ROBIN}}, + "external path strategy", external_path_strategy); } // Parse BucketFunctionType Status ParseBucketFunctionType(BucketFunctionType* bucket_function_type) const { - auto iter = config_map_.find(Options::BUCKET_FUNCTION_TYPE); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - if (str == "default") { - *bucket_function_type = BucketFunctionType::DEFAULT; - } else if (str == "mod") { - *bucket_function_type = BucketFunctionType::MOD; - } else if (str == "hive") { - *bucket_function_type = BucketFunctionType::HIVE; - } else { - return Status::Invalid(fmt::format("invalid bucket function type: {}", str)); - } - } - return Status::OK(); + return ParseEnum(Options::BUCKET_FUNCTION_TYPE, + {{"default", BucketFunctionType::DEFAULT}, + {"mod", BucketFunctionType::MOD}, + {"hive", BucketFunctionType::HIVE}}, + "bucket function type", bucket_function_type); } // Parse StartupMode Status ParseStartupMode(StartupMode* startup_mode) const { - auto iter = config_map_.find(Options::SCAN_MODE); - if (iter != config_map_.end()) { - std::string str = StringUtils::ToLowerCase(iter->second); - PAIMON_ASSIGN_OR_RAISE(*startup_mode, StartupMode::FromString(str)); + PAIMON_ASSIGN_OR_RAISE(std::optional value, + GetOptionalValue(Options::SCAN_MODE)); + if (value) { + PAIMON_ASSIGN_OR_RAISE( + *startup_mode, StartupMode::FromString(StringUtils::ToLowerCase(value.value()))); } return Status::OK(); } @@ -371,37 +293,81 @@ class ConfigParser { } private: - const std::map config_map_; + template + Status ParseEnum(const std::string& key, + std::initializer_list> candidates, + const std::string& error_name, T* value) const { + PAIMON_ASSIGN_OR_RAISE(std::optional config_value, + GetOptionalValue(key)); + if (!config_value) { + return Status::OK(); + } + std::string normalized_value = StringUtils::ToLowerCase(config_value.value()); + for (const auto& [candidate, candidate_value] : candidates) { + if (normalized_value == candidate) { + *value = candidate_value; + return Status::OK(); + } + } + return Status::Invalid(fmt::format("invalid {}: {}", error_name, normalized_value)); + } + + template + Result> GetOptionalValue(const std::string& key) const { + Result> result = + OptionsUtils::GetOptionalValueFromMap(config_map_, key); + if (!result.ok()) { + return Status::Invalid( + fmt::format("Invalid Config [{}: {}]", key, config_map_.at(key))); + } + return result.value(); + } + + const std::map& config_map_; }; // Impl is a private implementation of CoreOptions, // storing various configurable fields and their default values. struct CoreOptions::Impl { int64_t page_size = 64 * 1024; - std::optional target_file_size; int64_t target_file_row_num = std::numeric_limits::max(); - std::optional blob_target_file_size; int64_t source_split_target_size = 128 * 1024 * 1024; int64_t source_split_open_file_cost = 4 * 1024 * 1024; 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; int64_t realtime_read_view_ttl_millis = 5 * 60 * 1000; + int64_t write_buffer_spill_max_disk_size = std::numeric_limits::max(); + double variant_shredding_min_field_cardinality_ratio = 0.1; + double variant_shredding_adaptive_retention_ratio = 0.05; + double lookup_cache_bloom_filter_fpp = 0.05; + int64_t cache_page_size = 64 * 1024; // 64KB + int64_t lookup_cache_max_memory = 256 * 1024 * 1024; + double lookup_cache_high_prio_pool_ratio = 0.25; + int64_t lookup_cache_file_retention_ms = 1 * 3600 * 1000; // 1 hour + int64_t lookup_cache_max_disk_size = INT64_MAX; + std::optional target_file_size; + std::optional blob_target_file_size; std::shared_ptr file_format; std::shared_ptr file_system; std::shared_ptr manifest_file_format; std::shared_ptr cache; - std::optional scan_snapshot_id; std::optional scan_timestamp_millis; + std::optional optimized_compaction_interval; + std::optional compaction_total_size_threshold; + std::optional compaction_incremental_size_threshold; + std::shared_ptr changelog_file_format; ExpireConfig expire_config; std::vector sequence_field; std::vector remove_record_on_sequence_group; + std::vector changelog_row_deduplicate_ignore_fields; std::vector blob_fields; std::vector blob_descriptor_fields; std::vector blob_view_fields; @@ -412,17 +378,23 @@ struct CoreOptions::Impl { std::string manifest_compression = "zstd"; std::string branch = BranchManager::DEFAULT_MAIN_BRANCH; std::string data_file_prefix = "data-"; + std::string changelog_file_prefix = "changelog-"; std::string file_system_scheme_to_identifier_map_str; - std::optional field_default_func; std::optional scan_fallback_branch; std::optional data_file_external_paths; std::optional blob_view_upstream_warehouse; - + std::optional changelog_file_compression; + std::optional global_index_external_path; + std::optional scan_tag_name; + CompressOptions lookup_compress_options{"zstd", 1}; + CompressOptions spill_compress_options{"zstd", 1}; std::map raw_options; + std::map> file_format_per_level; + std::map file_compression_per_level; + StatisticsMode realtime_store_statistics_mode = StatisticsMode::NONE; int32_t bucket = -1; - int32_t manifest_merge_min_count = 30; int32_t scan_manifest_entry_cache_max_snapshots = 0; int32_t read_batch_size = 1024; @@ -433,85 +405,68 @@ struct CoreOptions::Impl { int32_t compaction_max_size_amplification_percent = 200; int32_t compaction_size_ratio = 1; int32_t num_sorted_runs_compaction_trigger = 5; - std::optional num_sorted_runs_stop_trigger; - std::optional num_levels; - SortOrder sequence_field_sort_order = SortOrder::ASCENDING; MergeEngine merge_engine = MergeEngine::DEDUPLICATE; SortEngine sort_engine = SortEngine::LOSER_TREE; ChangelogProducer changelog_producer = ChangelogProducer::NONE; ExternalPathStrategy external_path_strategy = ExternalPathStrategy::NONE; LookupCompactMode lookup_compact_mode = LookupCompactMode::RADICAL; - std::optional lookup_compact_max_interval; BucketFunctionType bucket_function_type = BucketFunctionType::DEFAULT; - int32_t file_compression_zstd_level = 1; - int64_t write_buffer_spill_max_disk_size = std::numeric_limits::max(); + CoreOptions::SequenceNumberInitMode write_sequence_number_init_mode = + CoreOptions::SequenceNumberInitMode::SCAN; + VariantShreddingInferenceMode variant_shredding_inference_mode = + VariantShreddingInferenceMode::PER_FILE; + int32_t variant_shredding_max_schema_width = 300; + int32_t variant_shredding_max_schema_depth = 50; + int32_t variant_shredding_max_infer_buffer_row = 4096; + int32_t variant_shredding_adaptive_max_infer_buffer_row = 256; + int32_t compact_off_peak_start_hour = -1; + int32_t compact_off_peak_end_hour = -1; + int32_t compact_off_peak_ratio = 0; + int32_t lookup_remote_level_threshold = INT32_MIN; + std::optional num_sorted_runs_stop_trigger; + std::optional num_levels; + std::optional lookup_compact_max_interval; + std::optional global_index_thread_num; + bool realtime_enabled = false; + bool scan_manifest_entry_lazy_decode_enabled = true; bool ignore_delete = false; bool manifest_delete_file_drop_stats = false; bool write_buffer_spillable = true; bool write_only = false; bool bucket_append_ordered = false; - CoreOptions::SequenceNumberInitMode write_sequence_number_init_mode = - CoreOptions::SequenceNumberInitMode::SCAN; bool deletion_vectors_enabled = false; bool deletion_vectors_bitmap64 = false; bool force_lookup = false; bool lookup_wait = true; + bool changelog_row_deduplicate = false; bool partial_update_remove_record_on_delete = false; bool aggregation_remove_record_on_delete = false; bool table_read_sequence_number_enabled = false; bool key_value_sequence_number_enabled = false; bool file_index_read_enabled = true; bool enable_adaptive_prefetch_strategy = true; + bool prefetch_io_metrics_enabled = false; bool index_file_in_data_file_dir = false; bool row_tracking_enabled = false; bool row_tracking_partition_group_on_commit = true; bool data_evolution_enabled = false; bool variant_infer_shredding_schema = false; - VariantShreddingInferenceMode variant_shredding_inference_mode = - VariantShreddingInferenceMode::PER_FILE; - int32_t variant_shredding_max_schema_width = 300; - int32_t variant_shredding_max_schema_depth = 50; - double variant_shredding_min_field_cardinality_ratio = 0.1; - int32_t variant_shredding_max_infer_buffer_row = 4096; - int32_t variant_shredding_adaptive_max_infer_buffer_row = 256; - double variant_shredding_adaptive_retention_ratio = 0.05; bool blob_view_resolve_enabled = true; bool blob_as_descriptor = false; - std::optional blob_split_by_file_size; bool legacy_partition_name_enabled = true; bool global_index_enabled = true; - std::optional global_index_thread_num; bool commit_force_compact = false; bool commit_discard_duplicate_files = false; bool dynamic_partition_overwrite = true; bool overwrite_upgrade = true; bool compaction_force_rewrite_all_files = false; bool compaction_force_up_level_0 = false; - std::optional global_index_external_path; - - std::optional scan_tag_name; - std::optional optimized_compaction_interval; - std::optional compaction_total_size_threshold; - std::optional compaction_incremental_size_threshold; - int32_t compact_off_peak_start_hour = -1; - int32_t compact_off_peak_end_hour = -1; - int32_t compact_off_peak_ratio = 0; bool lookup_cache_bloom_filter = true; - double lookup_cache_bloom_filter_fpp = 0.05; bool lookup_remote_file_enabled = false; - int32_t lookup_remote_level_threshold = INT32_MIN; - CompressOptions lookup_compress_options{"zstd", 1}; - CompressOptions spill_compress_options{"zstd", 1}; - int64_t cache_page_size = 64 * 1024; // 64KB - std::map> file_format_per_level; - std::map file_compression_per_level; - int64_t lookup_cache_max_memory = 256 * 1024 * 1024; - double lookup_cache_high_prio_pool_ratio = 0.25; - int64_t lookup_cache_file_retention_ms = 1 * 3600 * 1000; // 1 hour - int64_t lookup_cache_max_disk_size = INT64_MAX; + std::optional blob_split_by_file_size; // Parse basic table options: bucket, partition, file sizes, batch sizes, file system, etc. Status ParseBasicOptions( @@ -585,6 +540,8 @@ struct CoreOptions::Impl { PAIMON_RETURN_NOT_OK(parser.ParseExternalPathStrategy(&external_path_strategy)); // Parse data-file.prefix - file name prefix of data files, default "data-" PAIMON_RETURN_NOT_OK(parser.Parse(Options::DATA_FILE_PREFIX, &data_file_prefix)); + // Parse changelog-file.prefix - file name prefix of changelog files, default "changelog-" + PAIMON_RETURN_NOT_OK(parser.Parse(Options::CHANGELOG_FILE_PREFIX, &changelog_file_prefix)); // Parse row-tracking.enabled - whether to enable unique row id for append table PAIMON_RETURN_NOT_OK( parser.Parse(Options::ROW_TRACKING_ENABLED, &row_tracking_enabled)); @@ -633,6 +590,14 @@ struct CoreOptions::Impl { Options::FILE_FORMAT, /*default_identifier=*/"parquet", &file_format)); // Parse file.compression - default file compression, default "zstd" PAIMON_RETURN_NOT_OK(parser.Parse(Options::FILE_COMPRESSION, &file_compression)); + // Parse changelog-file.format - no default value + if (parser.ContainsKey(Options::CHANGELOG_FILE_FORMAT)) { + PAIMON_RETURN_NOT_OK(parser.ParseObject( + Options::CHANGELOG_FILE_FORMAT, file_format->Identifier(), &changelog_file_format)); + } + // Parse changelog-file.compression - no default value + PAIMON_RETURN_NOT_OK( + parser.Parse(Options::CHANGELOG_FILE_COMPRESSION, &changelog_file_compression)); // Parse file.compression.zstd-level - zstd compression level, default 1 PAIMON_RETURN_NOT_OK( parser.Parse(Options::FILE_COMPRESSION_ZSTD_LEVEL, &file_compression_zstd_level)); @@ -763,6 +728,13 @@ struct CoreOptions::Impl { PAIMON_RETURN_NOT_OK(parser.Parse(Options::FIELDS_DEFAULT_AGG_FUNC, &field_default_func)); // Parse changelog-producer - whether to double write to a changelog file, default "none" PAIMON_RETURN_NOT_OK(parser.ParseChangelogProducer(&changelog_producer)); + // Parse changelog-producer.row-deduplicate - skip unchanged row changelogs + PAIMON_RETURN_NOT_OK(parser.Parse(Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, + &changelog_row_deduplicate)); + // Parse changelog-producer.row-deduplicate-ignore-fields - ignored comparison fields + PAIMON_RETURN_NOT_OK(parser.ParseList( + Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, Options::FIELDS_SEPARATOR, + &changelog_row_deduplicate_ignore_fields, /*need_trim=*/true)); // Parse partial-update.remove-record-on-delete - remove whole row on delete PAIMON_RETURN_NOT_OK(parser.Parse(Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_DELETE, &partial_update_remove_record_on_delete)); @@ -803,12 +775,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 +793,10 @@ 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)); + PAIMON_RETURN_NOT_OK( + parser.Parse(Options::PREFETCH_IO_METRICS_ENABLED, &prefetch_io_metrics_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 +806,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 +1047,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 +1160,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; } @@ -1380,6 +1394,10 @@ bool CoreOptions::EnableAdaptivePrefetchStrategy() const { return impl_->enable_adaptive_prefetch_strategy; } +bool CoreOptions::PrefetchIoMetricsEnabled() const { + return impl_->prefetch_io_metrics_enabled; +} + Result> CoreOptions::GetFieldAggFunc( const std::string& field_name) const { ConfigParser parser(impl_->raw_options); @@ -1575,6 +1593,26 @@ ChangelogProducer CoreOptions::GetChangelogProducer() const { return impl_->changelog_producer; } +bool CoreOptions::ChangelogRowDeduplicate() const { + return impl_->changelog_row_deduplicate; +} + +const std::vector& CoreOptions::GetChangelogRowDeduplicateIgnoreFields() const { + return impl_->changelog_row_deduplicate_ignore_fields; +} + +std::string CoreOptions::ChangelogFilePrefix() const { + return impl_->changelog_file_prefix; +} + +std::shared_ptr CoreOptions::GetChangelogFileFormat() const { + return impl_->changelog_file_format; +} + +std::optional CoreOptions::GetChangelogFileCompression() const { + return impl_->changelog_file_compression; +} + LookupStrategy CoreOptions::GetLookupStrategy() const { return LookupStrategy::From( /*is_first_row=*/GetMergeEngine() == MergeEngine::FIRST_ROW, @@ -1654,6 +1692,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 9eb289883..345958e1a 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; @@ -202,11 +207,17 @@ class PAIMON_EXPORT CoreOptions { bool DeletionVectorsBitmap64() const; int64_t DeletionVectorTargetFileSize() const; ChangelogProducer GetChangelogProducer() const; + bool ChangelogRowDeduplicate() const; + const std::vector& GetChangelogRowDeduplicateIgnoreFields() const; + std::string ChangelogFilePrefix() const; + std::shared_ptr GetChangelogFileFormat() const; + std::optional GetChangelogFileCompression() const; LookupStrategy GetLookupStrategy() const; bool NeedLookup() const; bool PrepareCommitWaitCompaction() const; bool FileIndexReadEnabled() const; + int64_t FileIndexInManifestThreshold() const; std::map GetFieldsSequenceGroups() const; bool PartialUpdateRemoveRecordOnDelete() const; @@ -224,6 +235,8 @@ class PAIMON_EXPORT CoreOptions { std::string DataFilePrefix() const; + bool PrefetchIoMetricsEnabled() const; + bool IndexFileInDataFileDir() const; bool RowTrackingEnabled() const; diff --git a/src/paimon/core/core_options_test.cpp b/src/paimon/core/core_options_test.cpp index ebc127edb..c424b9cfe 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" @@ -38,6 +39,7 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({})); ASSERT_EQ(core_options.GetManifestFormat()->Identifier(), "avro"); ASSERT_EQ(core_options.GetFileFormat()->Identifier(), "parquet"); + ASSERT_EQ(nullptr, core_options.GetChangelogFileFormat()); ASSERT_EQ(core_options.GetWriteFileFormat(0)->Identifier(), "parquet"); ASSERT_EQ(core_options.GetWriteFileFormat(3)->Identifier(), "parquet"); ASSERT_TRUE(core_options.GetFileSystem()); @@ -54,7 +56,10 @@ 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(std::nullopt, core_options.GetChangelogFileCompression()); ASSERT_EQ("zstd", core_options.GetWriteFileCompression(0)); ASSERT_EQ("zstd", core_options.GetWriteFileCompression(3)); ASSERT_EQ("zstd", core_options.GetManifestCompression()); @@ -65,6 +70,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()); @@ -121,6 +127,9 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_FALSE(core_options.DeletionVectorsBitmap64()); ASSERT_EQ(2 * 1024 * 1024, core_options.DeletionVectorTargetFileSize()); ASSERT_EQ(ChangelogProducer::NONE, core_options.GetChangelogProducer()); + ASSERT_FALSE(core_options.ChangelogRowDeduplicate()); + ASSERT_TRUE(core_options.GetChangelogRowDeduplicateIgnoreFields().empty()); + ASSERT_EQ("changelog-", core_options.ChangelogFilePrefix()); ASSERT_FALSE(core_options.NeedLookup()); ASSERT_FALSE(core_options.PrepareCommitWaitCompaction()); LookupStrategy expected_lookup_strategy = {/*is_first_row=*/false, @@ -134,9 +143,11 @@ 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()); + ASSERT_FALSE(core_options.PrefetchIoMetricsEnabled()); ASSERT_EQ(core_options.DataFilePrefix(), "data-"); ASSERT_FALSE(core_options.IndexFileInDataFileDir()); ASSERT_FALSE(core_options.RowTrackingEnabled()); @@ -186,6 +197,7 @@ TEST(CoreOptionsTest, TestFromMap) { std::map options = { {Options::FILE_SYSTEM, "Local"}, {Options::FILE_FORMAT, "ORC"}, + {Options::CHANGELOG_FILE_FORMAT, "avro"}, {Options::MANIFEST_FORMAT, "avRo"}, {Options::BUCKET, "3"}, {Options::PAGE_SIZE, "128 kb"}, @@ -217,6 +229,8 @@ 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::PREFETCH_IO_METRICS_ENABLED, "true"}, {Options::SNAPSHOT_NUM_RETAINED_MIN, "15"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "30"}, {Options::SNAPSHOT_EXPIRE_LIMIT, "20"}, @@ -240,6 +254,10 @@ TEST(CoreOptionsTest, TestFromMap) { {Options::DELETION_VECTOR_BITMAP64, "true"}, {Options::DELETION_VECTOR_INDEX_FILE_TARGET_SIZE, "4MB"}, {Options::CHANGELOG_PRODUCER, "full-compaction"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, "true"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, "f0, f2"}, + {Options::CHANGELOG_FILE_PREFIX, "test-changelog-"}, + {Options::CHANGELOG_FILE_COMPRESSION, "lz4"}, {Options::FORCE_LOOKUP, "true"}, {"fields.g_1,g_3.sequence-group", "c,d"}, {Options::AGGREGATION_REMOVE_RECORD_ON_DELETE, "true"}, @@ -248,6 +266,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 +323,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"}, @@ -314,6 +334,7 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_TRUE(fs); ASSERT_EQ(core_options.GetFileFormat()->Identifier(), "orc"); + ASSERT_EQ(core_options.GetChangelogFileFormat()->Identifier(), "avro"); ASSERT_EQ(core_options.GetWriteFileFormat(0)->Identifier(), "avro"); ASSERT_EQ(core_options.GetWriteFileFormat(1)->Identifier(), "orc"); ASSERT_EQ(core_options.GetWriteFileFormat(3)->Identifier(), "parquet"); @@ -353,6 +374,8 @@ 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()); + ASSERT_TRUE(core_options.PrefetchIoMetricsEnabled()); ExpireConfig expire_config = core_options.GetExpireConfig(); ASSERT_EQ(15, expire_config.GetSnapshotRetainMin()); ASSERT_EQ(30, expire_config.GetSnapshotRetainMax()); @@ -381,6 +404,11 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_TRUE(core_options.DeletionVectorsBitmap64()); ASSERT_EQ(4 * 1024 * 1024, core_options.DeletionVectorTargetFileSize()); ASSERT_EQ(ChangelogProducer::FULL_COMPACTION, core_options.GetChangelogProducer()); + ASSERT_TRUE(core_options.ChangelogRowDeduplicate()); + ASSERT_EQ(std::vector({"f0", "f2"}), + core_options.GetChangelogRowDeduplicateIgnoreFields()); + ASSERT_EQ("test-changelog-", core_options.ChangelogFilePrefix()); + ASSERT_EQ(std::optional("lz4"), core_options.GetChangelogFileCompression()); ASSERT_TRUE(core_options.NeedLookup()); ASSERT_TRUE(core_options.PrepareCommitWaitCompaction()); LookupStrategy expected_lookup_strategy = {/*is_first_row=*/false, @@ -398,6 +426,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 +492,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 +522,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 +835,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 1fca5f374..6c0d23be0 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,8 @@ 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(), + /*enable_io_metrics=*/false, pool_)); } else { file_batch_reader = std::make_unique(data, target_type_, batch_size); diff --git a/src/paimon/core/deletionvectors/deletion_vector_test.cpp b/src/paimon/core/deletionvectors/deletion_vector_test.cpp index daf2ecf1c..bff4b46a3 100644 --- a/src/paimon/core/deletionvectors/deletion_vector_test.cpp +++ b/src/paimon/core/deletionvectors/deletion_vector_test.cpp @@ -54,7 +54,7 @@ std::shared_ptr CreateDataFileMeta(const std::string& file_name) { /*min_sequence_number=*/0, /*max_sequence_number=*/0, /*schema_id=*/0, DataFileMeta::DUMMY_LEVEL, std::vector>{}, Timestamp(0, 0), std::nullopt, nullptr, FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, - std::nullopt); + std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } } // namespace diff --git a/src/paimon/core/global_index/indexed_split_test.cpp b/src/paimon/core/global_index/indexed_split_test.cpp index 857f50bfa..03cd2976d 100644 --- a/src/paimon/core/global_index/indexed_split_test.cpp +++ b/src/paimon/core/global_index/indexed_split_test.cpp @@ -17,6 +17,8 @@ * under the License. */ +#include +#include #include #include #include @@ -26,6 +28,7 @@ #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/data/data_define.h" +#include "paimon/common/utils/math.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/fs/local/local_file_system.h" @@ -56,17 +59,20 @@ TEST(IndexedSplitTest, TestSimple) { "file1.orc", 100l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 50l, 249l, 0, 0, std::vector>(), Timestamp(1765535214349l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 50l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 50l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto meta2 = std::make_shared( "file2.orc", 101l, 100l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 250l, 349l, 0, 0, std::vector>(), Timestamp(1765535214349l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 250l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 250l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto meta3 = std::make_shared( "file3.orc", 102l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 1000l, 1199l, 0, 0, std::vector>(), Timestamp(1765535214349l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), @@ -82,8 +88,11 @@ TEST(IndexedSplitTest, TestSimple) { ASSERT_EQ(*result_indexed_split, *expected_indexed_split) << result_indexed_split->ToString(); ASSERT_OK_AND_ASSIGN(std::string serialize_bytes, Split::Serialize(result_indexed_split, pool)); - ASSERT_EQ(serialize_bytes, - std::string(reinterpret_cast(split_bytes.data()), split_bytes.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr roundtrip, + Split::Deserialize(serialize_bytes.data(), serialize_bytes.size(), pool)); + auto roundtrip_indexed_split = std::dynamic_pointer_cast(roundtrip); + ASSERT_EQ(*roundtrip_indexed_split, *expected_indexed_split) + << roundtrip_indexed_split->ToString(); } TEST(IndexedSplitTest, TestIndexedSplitWithScore) { @@ -107,17 +116,20 @@ TEST(IndexedSplitTest, TestIndexedSplitWithScore) { "file1.orc", 100l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 50l, 249l, 0, 0, std::vector>(), Timestamp(1765549435648l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 50l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 50l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto meta2 = std::make_shared( "file2.orc", 101l, 100l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 250l, 349l, 0, 0, std::vector>(), Timestamp(1765549435649l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 250l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 250l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto meta3 = std::make_shared( "file3.orc", 102l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 1000l, 1199l, 0, 0, std::vector>(), Timestamp(1765549435649l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), @@ -139,8 +151,39 @@ TEST(IndexedSplitTest, TestIndexedSplitWithScore) { "rowRanges=[[55, 56],[270, 270],[1001, 1002]], scores=[1.01,2.1,-1.32,4.23,50.74]") != std::string::npos); ASSERT_OK_AND_ASSIGN(std::string serialize_bytes, Split::Serialize(result_indexed_split, pool)); - ASSERT_EQ(serialize_bytes, - std::string(reinterpret_cast(split_bytes.data()), split_bytes.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr roundtrip, + Split::Deserialize(serialize_bytes.data(), serialize_bytes.size(), pool)); + auto roundtrip_indexed_split = std::dynamic_pointer_cast(roundtrip); + ASSERT_EQ(*roundtrip_indexed_split, *expected_indexed_split) + << roundtrip_indexed_split->ToString(); +} + +TEST(IndexedSplitTest, TestSerializeCanonicalizesNaNScore) { + auto pool = GetDefaultPool(); + DataSplitImpl::Builder builder( + /*partition=*/BinaryRow::EmptyRow(), + /*bucket=*/0, /*bucket_path=*/"bucket-0", + /*data_files=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr data_split, builder.Build()); + + const auto payload_nan = FloatingPointFromBits(0xffc12345U); + const auto canonical_nan = FloatingPointFromBits(kCanonicalFloatNaNBits); + auto indexed_split = std::make_shared( + data_split, std::vector{Range(0, 0)}, std::vector{payload_nan}); + auto canonical_indexed_split = std::make_shared( + data_split, std::vector{Range(0, 0)}, std::vector{canonical_nan}); + + ASSERT_OK_AND_ASSIGN(std::string serialized, Split::Serialize(indexed_split, pool)); + ASSERT_OK_AND_ASSIGN(std::string canonical_serialized, + Split::Serialize(canonical_indexed_split, pool)); + ASSERT_EQ(serialized, canonical_serialized); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr roundtrip, + Split::Deserialize(serialized.data(), serialized.size(), pool)); + auto roundtrip_indexed_split = std::dynamic_pointer_cast(roundtrip); + ASSERT_TRUE(roundtrip_indexed_split); + ASSERT_EQ(roundtrip_indexed_split->Scores().size(), 1); + ASSERT_TRUE(std::isnan(roundtrip_indexed_split->Scores()[0])); } TEST(IndexedSplitTest, TestValidate) { @@ -148,7 +191,8 @@ TEST(IndexedSplitTest, TestValidate) { "file.orc", 1l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 1000l, 1199l, 0, 0, std::vector>(), Timestamp(0l, 0), 0, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, 1000l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), diff --git a/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp index d01f18c44..ec65f4ff1 100644 --- a/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp +++ b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp @@ -53,7 +53,7 @@ class PkSortedBucketIndexStateTest : public ::testing::Test { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, file_source, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } /// Builds a payload whose source metadata lists the given sources in the given order. 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 e12374ddb..d0b677a78 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/async_key_value_producer_and_consumer.cpp b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp index 1792b43cd..b353aefc8 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp @@ -25,22 +25,96 @@ #include "arrow/c/abi.h" #include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/reader/batch_reader.h" namespace paimon { class MemoryPool; +namespace { + +class AsyncKeyValueQueueBatchSink : public AsyncKeyValueBatchSink { + public: + AsyncKeyValueQueueBatchSink(std::atomic* consume_finished, + tbb::concurrent_bounded_queue* kv_queue) + : consume_finished_(consume_finished), kv_queue_(kv_queue) {} + + Status Write(AsyncKeyValueBatchType type, std::vector&& rows) override { + if (*consume_finished_) { + return Status::Cancelled("Key value conversion is cancelled"); + } + kv_queue_->push(AsyncKeyValueRowsBatch{type, std::move(rows)}); + return Status::OK(); + } + + bool IsCancelled() const override { + return *consume_finished_; + } + + private: + std::atomic* consume_finished_; + tbb::concurrent_bounded_queue* kv_queue_; +}; + +} // namespace + +std::shared_ptr AsyncKeyValueBatchProducer::GetReaderMetrics() const { + return std::make_shared(); +} + +SortMergeReaderBatchProducer::SortMergeReaderBatchProducer( + std::unique_ptr&& sort_merge_reader, int32_t batch_size) + : sort_merge_reader_(std::move(sort_merge_reader)), + batch_size_(NormalizeProjectionBatchSize(batch_size)) {} + +Status SortMergeReaderBatchProducer::Produce(AsyncKeyValueBatchSink* sink) { + std::vector batch; + batch.reserve(batch_size_); + while (!sink->IsCancelled()) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + sort_merge_reader_->NextBatch()); + if (iterator == nullptr) { + break; + } + while (!sink->IsCancelled()) { + PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); + if (!has_next) { + break; + } + batch.push_back(std::move(iterator->Next())); + if (static_cast(batch.size()) >= batch_size_) { + PAIMON_RETURN_NOT_OK(sink->Write(AsyncKeyValueBatchType::DATA, std::move(batch))); + batch = std::vector(); + batch.reserve(batch_size_); + } + } + } + if (!batch.empty() && !sink->IsCancelled()) { + PAIMON_RETURN_NOT_OK(sink->Write(AsyncKeyValueBatchType::DATA, std::move(batch))); + } + return Status::OK(); +} + +std::shared_ptr SortMergeReaderBatchProducer::GetReaderMetrics() const { + return sort_merge_reader_->GetReaderMetrics(); +} + +void SortMergeReaderBatchProducer::Close() { + if (!closed_) { + sort_merge_reader_->Close(); + closed_ = true; + } +} + template AsyncKeyValueProducerAndConsumer::AsyncKeyValueProducerAndConsumer( - std::unique_ptr&& sort_merge_reader, ConsumerCreator create_consumer, - int32_t batch_size, int32_t consumer_thread_num, const std::shared_ptr& pool) - : batch_size_(std::min(batch_size, MAX_PROJECTION_BATCH_SIZE)), - consumer_thread_num_(consumer_thread_num), - pool_(pool), - sort_merge_reader_(std::move(sort_merge_reader)), - create_consumer_(std::move(create_consumer)) { - kv_queue_.set_capacity(consumer_thread_num * 2); + std::unique_ptr&& producer, ConsumerCreator create_consumer, + int32_t consumer_thread_num) + : consumer_thread_num_(consumer_thread_num), + create_consumer_(std::move(create_consumer)), + producer_(std::move(producer)) { + kv_queue_.set_capacity(consumer_thread_num_ * 2); result_queue_.set_capacity(RESULT_BATCH_COUNT); } @@ -70,6 +144,12 @@ Status AsyncKeyValueProducerAndConsumer::CheckStatusAndCleanUp() { template Result AsyncKeyValueProducerAndConsumer::NextBatch() { + PAIMON_ASSIGN_OR_RAISE(AsyncKeyValueResultBatch result, NextBatchWithType()); + return std::move(result.result); +} + +template +Result> AsyncKeyValueProducerAndConsumer::NextBatchWithType() { if (!producer_future_.valid()) { producer_future_ = std::async(std::launch::async, &AsyncKeyValueProducerAndConsumer::ProduceLoop, this) @@ -78,10 +158,10 @@ Result AsyncKeyValueProducerAndConsumer::NextBatch() { if (consumers_.empty()) { consumers_.reserve(consumer_thread_num_); for (int32_t i = 0; i < consumer_thread_num_; i++) { - Result>> consumer = create_consumer_(); - PAIMON_RETURN_NOT_OK(consumer.status()); + std::unique_ptr> consumer; + PAIMON_ASSIGN_OR_RAISE(consumer, create_consumer_()); auto async_consumer = std::make_unique>( - std::move(consumer).value(), consume_finished_, consumer_finished_count_, kv_queue_, + std::move(consumer), consume_finished_, consumer_finished_count_, kv_queue_, result_queue_); consumers_.push_back(std::move(async_consumer)); } @@ -90,16 +170,16 @@ Result AsyncKeyValueProducerAndConsumer::NextBatch() { if (next_batch_finished_) { // projection reader is eof - return R(); + return AsyncKeyValueResultBatch(); } - R result; + AsyncKeyValueResultBatch result; while (!result_queue_.try_pop(result)) { PAIMON_RETURN_NOT_OK(CheckStatusAndCleanUp()); if (consumer_finished_count_ == consumer_thread_num_ && result_queue_.empty()) { // all consume thread finished next_batch_finished_ = true; - return R(); + return AsyncKeyValueResultBatch(); } usleep(1000); } @@ -109,33 +189,9 @@ Result AsyncKeyValueProducerAndConsumer::NextBatch() { template Status AsyncKeyValueProducerAndConsumer::ProduceLoop() { - std::vector batch; - batch.reserve(batch_size_); - while (!consume_finished_) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, - sort_merge_reader_->NextBatch()); - if (iterator == nullptr) { - break; - } - while (!consume_finished_) { - PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); - if (!has_next) { - break; - } - batch.push_back(std::move(iterator->Next())); - if (static_cast(batch.size()) >= batch_size_) { - kv_queue_.push(std::move(batch)); - batch = std::vector(); - batch.reserve(batch_size_); - } - } - } - // Push remaining rows - if (!batch.empty()) { - kv_queue_.push(std::move(batch)); - } - // Push empty batch as EOF signal - kv_queue_.push(std::vector()); + AsyncKeyValueQueueBatchSink sink(&consume_finished_, &kv_queue_); + PAIMON_RETURN_NOT_OK(producer_->Produce(&sink)); + kv_queue_.push(AsyncKeyValueRowsBatch()); return Status::OK(); } @@ -155,8 +211,9 @@ void AsyncKeyValueProducerAndConsumer::CleanUp() { template void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { - R read_batch; - while (result_queue_.try_pop(read_batch)) { + AsyncKeyValueResultBatch tagged_batch; + while (result_queue_.try_pop(tagged_batch)) { + R& read_batch = tagged_batch.result; if constexpr (std::is_same_v) { if (!BatchReader::IsEofBatch(read_batch)) { ReaderUtils::ReleaseReadBatch(std::move(read_batch)); @@ -168,7 +225,7 @@ void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { } } - std::vector kv_batch; + AsyncKeyValueRowsBatch kv_batch; while (kv_queue_.try_pop(kv_batch)) { } } @@ -178,11 +235,11 @@ template class AsyncKeyValueProducerAndConsumer; template AsyncKeyValueConsumer::AsyncKeyValueConsumer( - std::unique_ptr>&& key_value_consumer, - std::atomic& consume_finished, std::atomic& consumer_finished_count, - tbb::concurrent_bounded_queue>& kv_queue, - tbb::concurrent_bounded_queue& result_queue) - : key_value_consumer_(std::move(key_value_consumer)), + std::unique_ptr>&& consumer, std::atomic& consume_finished, + std::atomic& consumer_finished_count, + tbb::concurrent_bounded_queue& kv_queue, + tbb::concurrent_bounded_queue>& result_queue) + : consumer_(std::move(consumer)), consume_finished_(consume_finished), consumer_finished_count_(consumer_finished_count), kv_queue_(kv_queue), @@ -205,18 +262,18 @@ Status AsyncKeyValueConsumer::GetStatus() const { template Status AsyncKeyValueConsumer::ConsumeLoop() { while (!consume_finished_) { - std::vector key_value_vec; - if (!kv_queue_.try_pop(key_value_vec)) { + AsyncKeyValueRowsBatch rows_batch; + if (!kv_queue_.try_pop(rows_batch)) { usleep(100); continue; } - if (key_value_vec.empty()) { + if (rows_batch.rows.empty()) { // Empty batch is EOF signal; re-push for other consumers - kv_queue_.push(std::move(key_value_vec)); + kv_queue_.push(std::move(rows_batch)); break; } - PAIMON_ASSIGN_OR_RAISE(R result, key_value_consumer_->NextBatch(key_value_vec)); - result_queue_.push(std::move(result)); + PAIMON_ASSIGN_OR_RAISE(R result, consumer_->NextBatch(rows_batch.rows)); + result_queue_.push(AsyncKeyValueResultBatch{rows_batch.type, std::move(result)}); } consumer_finished_count_++; return Status::OK(); @@ -227,7 +284,9 @@ void AsyncKeyValueConsumer::CleanUp() { if (consumer_future_.valid()) { [[maybe_unused]] Status status = consumer_future_.get(); } - key_value_consumer_->CleanUp(); + if (consumer_) { + consumer_->CleanUp(); + } } template class AsyncKeyValueConsumer; diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.h b/src/paimon/core/io/async_key_value_producer_and_consumer.h index af8bbed68..086ff1112 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.h +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -37,21 +38,80 @@ namespace paimon { template class AsyncKeyValueConsumer; -class MemoryPool; class Metrics; -// Asynchronous iterate SortMergeReader (producer) and row-to-array conversion (consumer), support -// multi-threaded conversion, R can be BatchReader::ReadBatch, KeyValueBatch +enum class AsyncKeyValueBatchType { + DATA, + CHANGELOG, +}; + +struct AsyncKeyValueRowsBatch { + AsyncKeyValueBatchType type = AsyncKeyValueBatchType::DATA; + std::vector rows; +}; + +class AsyncKeyValueBatchSink { + public: + virtual ~AsyncKeyValueBatchSink() = default; + + virtual Status Write(AsyncKeyValueBatchType type, std::vector&& rows) = 0; + + virtual bool IsCancelled() const = 0; +}; + +class AsyncKeyValueBatchProducer { + public: + virtual ~AsyncKeyValueBatchProducer() = default; + + virtual Status Produce(AsyncKeyValueBatchSink* sink) = 0; + + virtual std::shared_ptr GetReaderMetrics() const; + + virtual void Close() {} + + protected: + // Limits the number of rows sent to one Arrow projection call. + static int32_t NormalizeProjectionBatchSize(int32_t batch_size) { + return std::min(batch_size, MAX_PROJECTION_BATCH_SIZE); + } + + private: + static constexpr int32_t MAX_PROJECTION_BATCH_SIZE = 100000; +}; + +class SortMergeReaderBatchProducer : public AsyncKeyValueBatchProducer { + public: + SortMergeReaderBatchProducer(std::unique_ptr&& sort_merge_reader, + int32_t batch_size); + + Status Produce(AsyncKeyValueBatchSink* sink) override; + + std::shared_ptr GetReaderMetrics() const override; + + void Close() override; + + private: + std::unique_ptr sort_merge_reader_; + int32_t batch_size_; + bool closed_ = false; +}; + +template +struct AsyncKeyValueResultBatch { + AsyncKeyValueBatchType type = AsyncKeyValueBatchType::DATA; + R result; +}; + +// Asynchronous iterates AsyncKeyValueBatchProducer(producer) and row-to-array conversion +// (consumer), support multi-threaded conversion, R can be BatchReader::ReadBatch or KeyValueBatch. template class AsyncKeyValueProducerAndConsumer { public: using ConsumerCreator = std::function>>()>; - AsyncKeyValueProducerAndConsumer(std::unique_ptr&& sort_merge_reader, - ConsumerCreator create_consumer, int32_t batch_size, - int32_t consumer_thread_num, - const std::shared_ptr& pool); + AsyncKeyValueProducerAndConsumer(std::unique_ptr&& producer, + ConsumerCreator create_consumer, int32_t consumer_thread_num); ~AsyncKeyValueProducerAndConsumer() { CleanUp(); @@ -59,21 +119,20 @@ class AsyncKeyValueProducerAndConsumer { Result NextBatch(); + Result> NextBatchWithType(); + std::shared_ptr GetReaderMetrics() const { - return sort_merge_reader_->GetReaderMetrics(); + return producer_->GetReaderMetrics(); } void Close() { CleanUp(); - sort_merge_reader_->Close(); + producer_->Close(); } private: static constexpr int32_t RESULT_BATCH_COUNT = 3; - // in case write batch size is too large and overflow arrow array - static constexpr int32_t MAX_PROJECTION_BATCH_SIZE = 100000; - void CleanUpQueue(); Status ProduceLoop(); void CleanUp(); @@ -81,11 +140,9 @@ class AsyncKeyValueProducerAndConsumer { Status CheckStatusAndCleanUp(); private: - int32_t batch_size_; int32_t consumer_thread_num_; - std::shared_ptr pool_; - std::unique_ptr sort_merge_reader_; ConsumerCreator create_consumer_; + std::unique_ptr producer_; // produce: merge sort KeyValue and push result KeyValue to kv_queue_, consume: project KeyValue // to arrow array and push result array to result_queue_ @@ -94,18 +151,18 @@ class AsyncKeyValueProducerAndConsumer { std::shared_future producer_future_; std::vector>> consumers_; std::atomic consumer_finished_count_ = 0; - tbb::concurrent_bounded_queue> kv_queue_; - tbb::concurrent_bounded_queue result_queue_; + tbb::concurrent_bounded_queue kv_queue_; + tbb::concurrent_bounded_queue> result_queue_; }; template class AsyncKeyValueConsumer { public: - AsyncKeyValueConsumer(std::unique_ptr>&& key_value_consumer, + AsyncKeyValueConsumer(std::unique_ptr>&& consumer, std::atomic& consume_finished, std::atomic& consumer_finished_count, - tbb::concurrent_bounded_queue>& kv_queue, - tbb::concurrent_bounded_queue& result_queue); + tbb::concurrent_bounded_queue& kv_queue, + tbb::concurrent_bounded_queue>& result_queue); ~AsyncKeyValueConsumer() { CleanUp(); @@ -118,12 +175,12 @@ class AsyncKeyValueConsumer { Status ConsumeLoop(); private: - std::unique_ptr> key_value_consumer_; + std::unique_ptr> consumer_; std::shared_future consumer_future_; std::atomic& consume_finished_; std::atomic& consumer_finished_count_; - tbb::concurrent_bounded_queue>& kv_queue_; - tbb::concurrent_bounded_queue& result_queue_; + tbb::concurrent_bounded_queue& kv_queue_; + tbb::concurrent_bounded_queue>& result_queue_; }; } // namespace paimon diff --git a/src/paimon/core/io/async_key_value_projection_reader.h b/src/paimon/core/io/async_key_value_projection_reader.h index a272a0edf..79c7aa4de 100644 --- a/src/paimon/core/io/async_key_value_projection_reader.h +++ b/src/paimon/core/io/async_key_value_projection_reader.h @@ -38,10 +38,12 @@ class AsyncKeyValueProjectionReader : public BatchReader { -> Result>> { return KeyValueProjectionConsumer::Create(target_schema, target_to_src_mapping, pool); }; + std::unique_ptr producer = + std::make_unique(std::move(sort_merge_reader), + batch_size); producer_and_consumer_ = std::make_unique>( - std::move(sort_merge_reader), create_consumer, batch_size, projection_thread_num, - pool); + std::move(producer), create_consumer, projection_thread_num); } Result NextBatch() override { 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 000000000..5c97bb3da --- /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 000000000..883b719fc --- /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 000000000..1e2c9593f --- /dev/null +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -0,0 +1,280 @@ +/* + * 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, TestBsiAndBloomFilterEmbeddedRoundTrip) { + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bsi.columns", "f0"}, + {"file-index.bloom-filter.columns", "f1"}, + {"file-index.bloom-filter.f1.items", "100"}, + {"file-index.bloom-filter.f1.fpp", "0.01"}, + {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": null, "f1": 30}, + {"f0": 5, "f1": 40}])"))); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish("unused.orc")); + ASSERT_TRUE(result.embedded_index); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(result.embedded_index)); + + ASSERT_OK_AND_ASSIGN(auto bsi_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bsi_readers.size()); + ASSERT_OK_AND_ASSIGN(auto greater_result, bsi_readers[0]->VisitGreaterThan(Literal(1))); + ASSERT_EQ("{3}", greater_result->ToString()); + ASSERT_OK_AND_ASSIGN(auto null_result, bsi_readers[0]->VisitIsNull()); + ASSERT_EQ("{2}", null_result->ToString()); + + ASSERT_OK_AND_ASSIGN(auto bloom_readers, ReadColumn(reader.get(), "f1")); + ASSERT_EQ(1, bloom_readers.size()); + ASSERT_OK_AND_ASSIGN(auto present_result, bloom_readers[0]->VisitEqual(Literal(30))); + ASSERT_TRUE(present_result->IsRemain().value()); +} + +TEST_F(DataFileIndexWriterTest, TestUnavailableWriterFailsCreation) { + ASSERT_NOK_WITH_MSG(CreateWriter({{"file-index.unknown.columns", "f0"}}), + "File index type 'unknown' is not registered"); +} + +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_meta.cpp b/src/paimon/core/io/data_file_meta.cpp index 7ce767f4b..cd8666098 100644 --- a/src/paimon/core/io/data_file_meta.cpp +++ b/src/paimon/core/io/data_file_meta.cpp @@ -72,7 +72,8 @@ Result> DataFileMeta::ForAppend( file_name, file_size, row_count, EmptyMinKey(), EmptyMaxKey(), SimpleStats::EmptyStats(), row_stats, min_sequence_number, max_sequence_number, schema_id, DUMMY_LEVEL, extra_files, Timestamp(/*millisecond=*/local_micro / 1000, /*nano_of_millisecond=*/0), 0ll, - embedded_index, file_source, value_stats_cols, external_path, first_row_id, write_cols); + embedded_index, file_source, value_stats_cols, external_path, first_row_id, write_cols, + /*column_max_sequence_numbers=*/std::nullopt); } Result> DataFileMeta::Upgrade(int32_t new_level) const { @@ -84,7 +85,7 @@ Result> DataFileMeta::Upgrade(int32_t new_level) c file_name, file_size, row_count, min_key, max_key, key_stats, value_stats, min_sequence_number, max_sequence_number, schema_id, new_level, extra_files, creation_time, delete_row_count, embedded_index, file_source, value_stats_cols, external_path, - first_row_id, write_cols); + first_row_id, write_cols, column_max_sequence_numbers); } std::shared_ptr DataFileMeta::CopyWithExtraFiles( @@ -93,7 +94,16 @@ std::shared_ptr DataFileMeta::CopyWithExtraFiles( file_name, file_size, row_count, min_key, max_key, key_stats, value_stats, min_sequence_number, max_sequence_number, schema_id, level, new_extra_files, creation_time, delete_row_count, embedded_index, file_source, value_stats_cols, external_path, - first_row_id, write_cols); + first_row_id, write_cols, column_max_sequence_numbers); +} + +std::shared_ptr DataFileMeta::CopyWithColumnMaxSequenceNumbers( + const std::optional>& new_column_max_sequence_numbers) const { + return std::make_shared( + file_name, file_size, row_count, min_key, max_key, key_stats, value_stats, + min_sequence_number, max_sequence_number, schema_id, level, extra_files, creation_time, + delete_row_count, embedded_index, file_source, value_stats_cols, external_path, + first_row_id, write_cols, new_column_max_sequence_numbers); } std::shared_ptr DataFileMeta::CopyWithoutStats() const { @@ -101,7 +111,7 @@ std::shared_ptr DataFileMeta::CopyWithoutStats() const { file_name, file_size, row_count, min_key, max_key, key_stats, SimpleStats::EmptyStats(), min_sequence_number, max_sequence_number, schema_id, level, extra_files, creation_time, delete_row_count, embedded_index, file_source, std::vector(), external_path, - first_row_id, write_cols); + first_row_id, write_cols, column_max_sequence_numbers); } DataFileMeta::DataFileMeta( @@ -113,7 +123,8 @@ DataFileMeta::DataFileMeta( const std::shared_ptr& _embedded_index, const std::optional& _file_source, const std::optional>& _value_stats_cols, const std::optional& _external_path, const std::optional& _first_row_id, - const std::optional>& _write_cols) + const std::optional>& _write_cols, + const std::optional>& _column_max_sequence_numbers) : file_name(_file_name), file_size(_file_size), row_count(_row_count), @@ -133,7 +144,8 @@ DataFileMeta::DataFileMeta( value_stats_cols(_value_stats_cols), external_path(_external_path), first_row_id(_first_row_id), - write_cols(_write_cols) {} + write_cols(_write_cols), + column_max_sequence_numbers(_column_max_sequence_numbers) {} Result DataFileMeta::FileFormat() const { size_t last_dot_index = file_name.find_last_of("."); @@ -198,7 +210,8 @@ bool DataFileMeta::operator==(const DataFileMeta& other) const { creation_time == other.creation_time && delete_row_count == other.delete_row_count && file_source == other.file_source && value_stats_cols == other.value_stats_cols && external_path == other.external_path && first_row_id == other.first_row_id && - write_cols == other.write_cols; + write_cols == other.write_cols && + column_max_sequence_numbers == other.column_max_sequence_numbers; } bool DataFileMeta::operator!=(const DataFileMeta& other) const { @@ -243,7 +256,8 @@ bool DataFileMeta::TEST_Equal(const DataFileMeta& other) const { level == other.level && delete_row_count == other.delete_row_count && file_source == other.file_source && value_stats_cols == other.value_stats_cols && compare_optional_ignore_name(external_path, other.external_path) && - first_row_id == other.first_row_id && write_cols == other.write_cols; + first_row_id == other.first_row_id && write_cols == other.write_cols && + column_max_sequence_numbers == other.column_max_sequence_numbers; } std::string DataFileMeta::ToString() const { @@ -261,7 +275,8 @@ std::string DataFileMeta::ToString() const { "{}, " "keyStats: {}, valueStats: {}, minSequenceNumber: {}, maxSequenceNumber: {}, schemaId: " "{}, level: {}, extraFiles: {}, creationTime: {}, deleteRowCount: {}, fileSource: {}, " - "valueStatsCols: {}, externalPath: {}, firstRowId: {}, writeCols: {}}}", + "valueStatsCols: {}, externalPath: {}, firstRowId: {}, writeCols: {}, " + "columnMaxSequenceNumbers: {}}}", file_name, file_size, row_count, embedded_index == nullptr ? "null" : std::string(embedded_index->data(), embedded_index->size()), @@ -275,7 +290,10 @@ std::string DataFileMeta::ToString() const { : fmt::format("{}", fmt::join(value_stats_cols.value(), ", ")), external_path == std::nullopt ? "null" : external_path.value(), first_row_id == std::nullopt ? "null" : std::to_string(first_row_id.value()), - write_cols == std::nullopt ? "null" : fmt::format("{}", write_cols.value())); + write_cols == std::nullopt ? "null" : fmt::format("{}", write_cols.value()), + column_max_sequence_numbers == std::nullopt + ? "null" + : fmt::format("{}", column_max_sequence_numbers.value())); } int64_t DataFileMeta::GetMaxSequenceNumber( @@ -315,6 +333,9 @@ const std::shared_ptr& DataFileMeta::DataType() { arrow::field("_FIRST_ROW_ID", arrow::int64(), /*nullable=*/true), arrow::field("_WRITE_COLS", arrow::list(arrow::field("item", arrow::utf8(), /*nullable=*/false)), + /*nullable=*/true), + arrow::field("_WRITE_COLS_SEQUENCES", + arrow::list(arrow::field("item", arrow::int64(), /*nullable=*/false)), /*nullable=*/true)}); return schema; } diff --git a/src/paimon/core/io/data_file_meta.h b/src/paimon/core/io/data_file_meta.h index 98aaa0ee7..436d73785 100644 --- a/src/paimon/core/io/data_file_meta.h +++ b/src/paimon/core/io/data_file_meta.h @@ -58,7 +58,8 @@ struct DataFileMeta { const std::optional>& _value_stats_cols, const std::optional& _external_path, const std::optional& _first_row_id, - const std::optional>& _write_cols); + const std::optional>& _write_cols, + const std::optional>& _column_max_sequence_numbers); static Result> ForAppend( const std::string& file_name, int64_t file_size, int64_t row_count, @@ -83,6 +84,9 @@ struct DataFileMeta { std::shared_ptr CopyWithExtraFiles( const std::vector>& new_extra_files) const; + std::shared_ptr CopyWithColumnMaxSequenceNumbers( + const std::optional>& new_column_max_sequence_numbers) const; + /// Create a copy without value statistics. All other metadata is preserved. /// /// @return A new metadata object with empty value statistics and value-stat columns. @@ -167,5 +171,12 @@ struct DataFileMeta { std::optional first_row_id; std::optional> write_cols; + + /// Maximum sequence number per physical table field after data-evolution compaction. + /// + /// Values follow the table-field order selected by `write_cols` when it is non-null (system + /// fields are ignored), or the file schema field order otherwise. A null value means that only + /// the file-level sequence range is available. + std::optional> column_max_sequence_numbers; }; } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_09_serializer.cpp b/src/paimon/core/io/data_file_meta_09_serializer.cpp index 7eaffc911..7eae41d13 100644 --- a/src/paimon/core/io/data_file_meta_09_serializer.cpp +++ b/src/paimon/core/io/data_file_meta_09_serializer.cpp @@ -114,7 +114,7 @@ Result> DataFileMeta09Serializer::FromRow( embedded_file_index, file_source, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_10_serializer.cpp b/src/paimon/core/io/data_file_meta_10_serializer.cpp index fea60eeb5..bf99f99d8 100644 --- a/src/paimon/core/io/data_file_meta_10_serializer.cpp +++ b/src/paimon/core/io/data_file_meta_10_serializer.cpp @@ -123,7 +123,8 @@ Result> DataFileMeta10Serializer::FromRow( min_sequence_number, max_sequence_number, schema_id, level, InternalRowUtils::FromStringArrayData(extra_files.get()), creation_time, delete_row_count, embedded_file_index, file_source, std::optional>(value_stats_cols), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_12_serializer.cpp b/src/paimon/core/io/data_file_meta_12_serializer.cpp index fe46a85c8..07fa1353a 100644 --- a/src/paimon/core/io/data_file_meta_12_serializer.cpp +++ b/src/paimon/core/io/data_file_meta_12_serializer.cpp @@ -128,7 +128,8 @@ Result> DataFileMeta12Serializer::FromRow( min_sequence_number, max_sequence_number, schema_id, level, InternalRowUtils::FromStringArrayData(extra_files.get()), creation_time, delete_row_count, embedded_file_index, file_source, std::optional>(value_stats_cols), - external_path, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + external_path, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_first_row_id_legacy_serializer.cpp b/src/paimon/core/io/data_file_meta_first_row_id_legacy_serializer.cpp index a834e47f0..3e82290d6 100644 --- a/src/paimon/core/io/data_file_meta_first_row_id_legacy_serializer.cpp +++ b/src/paimon/core/io/data_file_meta_first_row_id_legacy_serializer.cpp @@ -134,7 +134,8 @@ Result> DataFileMetaFirstRowIdLegacySerializer::Fr min_sequence_number, max_sequence_number, schema_id, level, InternalRowUtils::FromStringArrayData(extra_files.get()), creation_time, delete_row_count, embedded_file_index, file_source, std::optional>(value_stats_cols), - external_path, first_row_id, /*write_cols=*/std::nullopt); + external_path, first_row_id, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_serializer.cpp b/src/paimon/core/io/data_file_meta_serializer.cpp index 560299684..9f88ebec0 100644 --- a/src/paimon/core/io/data_file_meta_serializer.cpp +++ b/src/paimon/core/io/data_file_meta_serializer.cpp @@ -24,6 +24,7 @@ #include #include +#include "paimon/common/data/binary_array.h" #include "paimon/common/data/binary_row_writer.h" #include "paimon/common/data/binary_string.h" #include "paimon/common/data/internal_row.h" @@ -41,7 +42,7 @@ class InternalArray; class MemoryPool; Result DataFileMetaSerializer::ToRow(const std::shared_ptr& meta) const { - BinaryRow row(20); + BinaryRow row(NumFields()); BinaryRowWriter writer(&row, 32 * 1024, pool_.get()); writer.WriteString(0, BinaryString::FromString(meta->file_name, pool_.get())); writer.WriteLong(1, meta->file_size); @@ -95,6 +96,12 @@ Result DataFileMetaSerializer::ToRow(const std::shared_ptrwrite_cols.value(), pool_)); } + if (meta->column_max_sequence_numbers == std::nullopt) { + writer.SetNullAt(20); + } else { + writer.WriteArray( + 20, BinaryArray::FromLongArray(meta->column_max_sequence_numbers.value(), pool_.get())); + } writer.Complete(); return row; } @@ -160,6 +167,15 @@ Result> DataFileMetaSerializer::FromRow( } write_cols = InternalRowUtils::FromNotNullStringArrayData(array.get()); } + + std::optional> column_max_sequence_numbers; + if (!row.IsNullAt(20)) { + std::shared_ptr array = row.GetArray(20); + if (array == nullptr) { + return Status::Invalid("invalid column max sequence numbers"); + } + PAIMON_ASSIGN_OR_RAISE(column_max_sequence_numbers, array->ToLongArray()); + } PAIMON_ASSIGN_OR_RAISE(BinaryRow min_values, SerializationUtils::DeserializeBinaryRow(min_key)); PAIMON_ASSIGN_OR_RAISE(BinaryRow max_values, SerializationUtils::DeserializeBinaryRow(max_key)); PAIMON_ASSIGN_OR_RAISE(SimpleStats key_stats, @@ -171,7 +187,7 @@ Result> DataFileMetaSerializer::FromRow( min_sequence_number, max_sequence_number, schema_id, level, InternalRowUtils::FromStringArrayData(extra_files.get()), creation_time, delete_row_count, embedded_file_index, file_source, std::optional>(value_stats_cols), - external_path, first_row_id, write_cols); + external_path, first_row_id, write_cols, column_max_sequence_numbers); } } // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_serializer_test.cpp b/src/paimon/core/io/data_file_meta_serializer_test.cpp index 18ba25788..5fdc452b3 100644 --- a/src/paimon/core/io/data_file_meta_serializer_test.cpp +++ b/src/paimon/core/io/data_file_meta_serializer_test.cpp @@ -21,13 +21,16 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/array/builder_base.h" #include "gtest/gtest.h" #include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/data_file_meta_write_cols_legacy_serializer.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/data/timestamp.h" #include "paimon/io/byte_array_input_stream.h" @@ -49,7 +52,9 @@ class DataFileMetaSerializerTest : public testing::Test { } private: - std::shared_ptr GetDataFileMeta() { + std::shared_ptr GetDataFileMeta( + const std::optional>& column_max_sequence_numbers = + std::vector{16, 32}) { return std::make_shared( "some_file_name", 1024, 8, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), /*min_seq_no=*/16, @@ -58,7 +63,8 @@ class DataFileMetaSerializerTest : public testing::Test { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/3, /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::optional(), - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + column_max_sequence_numbers); } const int32_t TRIES = 100; @@ -68,6 +74,7 @@ class DataFileMetaSerializerTest : public testing::Test { TEST_F(DataFileMetaSerializerTest, TestToFromRow) { DataFileMetaSerializer serializer(memory_pool_); + ASSERT_EQ(serializer.NumFields(), 21); auto expected = GetDataFileMeta(); for (int32_t i = 0; i < TRIES; i++) { ASSERT_OK_AND_ASSIGN(auto row, serializer.ToRow(expected)); @@ -76,6 +83,15 @@ TEST_F(DataFileMetaSerializerTest, TestToFromRow) { } } +TEST_F(DataFileMetaSerializerTest, TestLegacySerializerSchema) { + DataFileMetaWriteColsLegacySerializer serializer(memory_pool_); + ASSERT_EQ(serializer.NumFields(), 20); + auto legacy_type = + checked_pointer_cast(DataFileMetaWriteColsLegacySerializer::DataType()); + ASSERT_EQ(legacy_type->num_fields(), 20); + ASSERT_EQ(legacy_type->field(19)->name(), "_WRITE_COLS"); +} + TEST_F(DataFileMetaSerializerTest, TestSerialize) { DataFileMetaSerializer serializer(memory_pool_); auto expected = GetDataFileMeta(); diff --git a/src/paimon/core/io/data_file_meta_test.cpp b/src/paimon/core/io/data_file_meta_test.cpp index f7283f158..f76bc0ae8 100644 --- a/src/paimon/core/io/data_file_meta_test.cpp +++ b/src/paimon/core/io/data_file_meta_test.cpp @@ -37,7 +37,8 @@ TEST(DataFileMetaTest, TestCopyWithoutStats) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::vector({"f0", "f1"}), /*external_path=*/"file:/tmp/bucket-0/data-0.orc", /*first_row_id=*/100, - /*write_cols=*/std::vector({"f0"})); + /*write_cols=*/std::vector({"f0"}), + /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr result = file_meta->CopyWithoutStats(); @@ -68,7 +69,8 @@ TEST(DataFileMetaTest, TestAddRowCount) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(3, file_meta.AddRowCount().value()); // test null delete row count file_meta.delete_row_count = std::nullopt; @@ -85,7 +87,8 @@ TEST(DataFileMetaTest, TestFileFormat) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_OK_AND_ASSIGN(auto file_format, file_meta.FileFormat()); ASSERT_EQ("orc", file_format); file_meta.file_name = "data-80110e15-97b5-4bcf-ac09-6ca2659a4950-0.parquet"; @@ -107,7 +110,8 @@ TEST(DataFileMetaTest, TestExternalPathDir) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/"file:/tmp/bucket-0/data-80110e15-97b5-4bcf-ac09-6ca2659a4950-0.orc", - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ("file:/tmp/bucket-0", file_meta.ExternalPathDir().value()); file_meta.external_path = std::nullopt; ASSERT_EQ(std::nullopt, file_meta.ExternalPathDir()); @@ -123,7 +127,8 @@ TEST(DataFileMetaTest, TestGetMaxSequenceNumber) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-80110e15-97b5-4bcf-ac09-6ca2659a4950-1.orc", /*file_size=*/645, /*row_count=*/5, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), SimpleStats::EmptyStats(), @@ -133,7 +138,8 @@ TEST(DataFileMetaTest, TestGetMaxSequenceNumber) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(4, DataFileMeta::GetMaxSequenceNumber({file_meta1})); ASSERT_EQ(10, DataFileMeta::GetMaxSequenceNumber({file_meta1, file_meta2})); ASSERT_EQ(-1, DataFileMeta::GetMaxSequenceNumber({})); @@ -152,7 +158,8 @@ TEST(DataFileMetaTest, TestNonNullFirstRowId) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_OK_AND_ASSIGN(int64_t first_row_id, file_meta->NonNullFirstRowId()); ASSERT_EQ(100, first_row_id); } @@ -167,7 +174,7 @@ TEST(DataFileMetaTest, TestNonNullFirstRowId) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ASSERT_NOK_WITH_MSG(file_meta->NonNullFirstRowId(), "First row id of data-1.orc should not be null."); } @@ -184,7 +191,8 @@ TEST(DataFileMetaTest, TestToFileSelection) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); { ASSERT_OK_AND_ASSIGN(std::optional result, @@ -223,7 +231,8 @@ TEST(DataFileMetaTest, TestUpgrade) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); // test normal upgrade ASSERT_OK_AND_ASSIGN(auto new_file_meta, file_meta->Upgrade(10)); ASSERT_EQ(new_file_meta->level, 10); diff --git a/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.cpp b/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.cpp new file mode 100644 index 000000000..ac2d480b9 --- /dev/null +++ b/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.cpp @@ -0,0 +1,156 @@ +/* + * 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_meta_write_cols_legacy_serializer.h" + +#include +#include +#include +#include + +#include "arrow/type.h" +#include "paimon/common/data/internal_row.h" +#include "paimon/common/utils/internal_row_utils.h" +#include "paimon/common/utils/serialization_utils.h" +#include "paimon/core/manifest/file_source.h" +#include "paimon/core/stats/simple_stats.h" +#include "paimon/data/timestamp.h" +#include "paimon/status.h" + +namespace paimon { + +class Bytes; +class InternalArray; + +const std::shared_ptr& DataFileMetaWriteColsLegacySerializer::DataType() { + static std::shared_ptr schema = arrow::struct_( + {arrow::field("_FILE_NAME", arrow::utf8(), /*nullable=*/false), + arrow::field("_FILE_SIZE", arrow::int64(), /*nullable=*/false), + arrow::field("_ROW_COUNT", arrow::int64(), /*nullable=*/false), + arrow::field("_MIN_KEY", arrow::binary(), /*nullable=*/false), + arrow::field("_MAX_KEY", arrow::binary(), /*nullable=*/false), + arrow::field("_KEY_STATS", SimpleStats::DataType(), /*nullable=*/false), + arrow::field("_VALUE_STATS", SimpleStats::DataType(), /*nullable=*/false), + arrow::field("_MIN_SEQUENCE_NUMBER", arrow::int64(), /*nullable=*/false), + arrow::field("_MAX_SEQUENCE_NUMBER", arrow::int64(), /*nullable=*/false), + arrow::field("_SCHEMA_ID", arrow::int64(), /*nullable=*/false), + arrow::field("_LEVEL", arrow::int32(), /*nullable=*/false), + arrow::field("_EXTRA_FILES", + arrow::list(arrow::field("item", arrow::utf8(), /*nullable=*/false)), + /*nullable=*/false), + arrow::field("_CREATION_TIME", arrow::timestamp(arrow::TimeUnit::MILLI), + /*nullable=*/true), + arrow::field("_DELETE_ROW_COUNT", arrow::int64(), /*nullable=*/true), + arrow::field("_EMBEDDED_FILE_INDEX", arrow::binary(), /*nullable=*/true), + arrow::field("_FILE_SOURCE", arrow::int8(), /*nullable=*/true), + arrow::field("_VALUE_STATS_COLS", + arrow::list(arrow::field("item", arrow::utf8(), /*nullable=*/false)), + /*nullable=*/true), + arrow::field("_EXTERNAL_PATH", arrow::utf8(), /*nullable=*/true), + arrow::field("_FIRST_ROW_ID", arrow::int64(), /*nullable=*/true), + arrow::field("_WRITE_COLS", + arrow::list(arrow::field("item", arrow::utf8(), /*nullable=*/false)), + /*nullable=*/true)}); + return schema; +} + +Result DataFileMetaWriteColsLegacySerializer::ToRow( + const std::shared_ptr&) const { + assert(false); + return Status::Invalid("to row for DataFileMetaWriteColsLegacySerializer is invalid"); +} + +Result> DataFileMetaWriteColsLegacySerializer::FromRow( + const InternalRow& row) const { + auto file_name = row.GetString(0); + auto file_size = row.GetLong(1); + auto row_count = row.GetLong(2); + auto min_key = row.GetBinary(3); + auto max_key = row.GetBinary(4); + auto key_stats_row = row.GetRow(5, 3); + auto value_stats_row = row.GetRow(6, 3); + auto min_sequence_number = row.GetLong(7); + auto max_sequence_number = row.GetLong(8); + auto schema_id = row.GetLong(9); + auto level = row.GetInt(10); + std::shared_ptr extra_files = row.GetArray(11); + auto creation_time = row.GetTimestamp(12, 3); + + assert(min_key && max_key && key_stats_row && value_stats_row); + if (extra_files == nullptr) { + return Status::Invalid("extra files is empty"); + } + + std::optional delete_row_count; + if (!row.IsNullAt(13)) { + delete_row_count = row.GetLong(13); + } + std::shared_ptr embedded_file_index; + if (!row.IsNullAt(14)) { + embedded_file_index = row.GetBinary(14); + } + + std::optional file_source; + if (!row.IsNullAt(15)) { + PAIMON_ASSIGN_OR_RAISE(file_source, FileSource::FromByteValue(row.GetByte(15))); + } + + std::optional> value_stats_cols; + if (!row.IsNullAt(16)) { + std::shared_ptr array = row.GetArray(16); + if (array == nullptr) { + return Status::Invalid("invalid value stats cols"); + } + value_stats_cols = InternalRowUtils::FromNotNullStringArrayData(array.get()); + } + + std::optional external_path; + if (!row.IsNullAt(17)) { + external_path = row.GetString(17).ToString(); + } + std::optional first_row_id; + if (!row.IsNullAt(18)) { + first_row_id = row.GetLong(18); + } + + std::optional> write_cols; + if (!row.IsNullAt(19)) { + std::shared_ptr array = row.GetArray(19); + if (array == nullptr) { + return Status::Invalid("invalid write cols"); + } + write_cols = InternalRowUtils::FromNotNullStringArrayData(array.get()); + } + + PAIMON_ASSIGN_OR_RAISE(BinaryRow min_values, SerializationUtils::DeserializeBinaryRow(min_key)); + PAIMON_ASSIGN_OR_RAISE(BinaryRow max_values, SerializationUtils::DeserializeBinaryRow(max_key)); + PAIMON_ASSIGN_OR_RAISE(SimpleStats key_stats, + SimpleStats::FromRow(key_stats_row.get(), pool_.get())); + PAIMON_ASSIGN_OR_RAISE(SimpleStats value_stats, + SimpleStats::FromRow(value_stats_row.get(), pool_.get())); + return std::make_shared( + file_name.ToString(), file_size, row_count, min_values, max_values, key_stats, value_stats, + min_sequence_number, max_sequence_number, schema_id, level, + InternalRowUtils::FromStringArrayData(extra_files.get()), creation_time, delete_row_count, + embedded_file_index, file_source, std::optional>(value_stats_cols), + external_path, first_row_id, write_cols, + /*column_max_sequence_numbers=*/std::nullopt); +} + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.h b/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.h new file mode 100644 index 000000000..0a08d28cc --- /dev/null +++ b/src/paimon/core/io/data_file_meta_write_cols_legacy_serializer.h @@ -0,0 +1,51 @@ +/* + * 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/common/data/binary_row.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/utils/object_serializer.h" +#include "paimon/result.h" + +namespace arrow { +class DataType; +} // namespace arrow + +namespace paimon { +class InternalRow; +class MemoryPool; + +/// Legacy serializer for `DataFileMeta` before column sequence numbers were introduced. +class DataFileMetaWriteColsLegacySerializer + : public ObjectSerializer> { + public: + static const std::shared_ptr& DataType(); + + explicit DataFileMetaWriteColsLegacySerializer(const std::shared_ptr& pool) + : ObjectSerializer>(DataType(), pool) {} + + Result ToRow(const std::shared_ptr& meta) const override; + + Result> FromRow(const InternalRow& row) const override; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/data_file_path_factory.h b/src/paimon/core/io/data_file_path_factory.h index b49154f1b..a9050eeee 100644 --- a/src/paimon/core/io/data_file_path_factory.h +++ b/src/paimon/core/io/data_file_path_factory.h @@ -56,8 +56,9 @@ class DataFilePathFactory : public PathFactory { return NewPath(data_file_prefix_); } - std::string NewChangelogPath() const { - return NewPath(std::string(CHANGELOG_FILE_PREFIX)); + std::string NewChangelogPath(const std::string& changelog_file_prefix, + const std::string& format_identifier) const { + return NewPathFromName(NewFileName(changelog_file_prefix, "." + format_identifier)); } std::string NewBlobPath() const { diff --git a/src/paimon/core/io/data_file_path_factory_test.cpp b/src/paimon/core/io/data_file_path_factory_test.cpp index 12618a2ec..fe4824c94 100644 --- a/src/paimon/core/io/data_file_path_factory_test.cpp +++ b/src/paimon/core/io/data_file_path_factory_test.cpp @@ -59,6 +59,13 @@ TEST_F(DataFilePathFactoryTest, TestNewPath) { ASSERT_EQ(factory_.NewPathFromName("index-file"), "/tmp/index-file"); } +TEST_F(DataFilePathFactoryTest, TestNewChangelogPath) { + std::string path = factory_.NewChangelogPath("changes-", "parquet"); + + ASSERT_TRUE(path.find("/tmp/changes-") != std::string::npos); + ASSERT_TRUE(StringUtils::EndsWith(path, ".parquet")); +} + TEST_F(DataFilePathFactoryTest, TestNewPathWithDataFilePrefixAndExternalPath) { DataFilePathFactory factory; ASSERT_OK_AND_ASSIGN( @@ -97,7 +104,7 @@ TEST_F(DataFilePathFactoryTest, TestToPath) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/"file:/test/bucket-0/example.txt", /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(factory_.ToPath(file_meta), "file:/test/bucket-0/example.txt"); } @@ -118,7 +125,7 @@ TEST_F(DataFilePathFactoryTest, TestToAlignedPath) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/"file:/test/bucket-0/data-0.txt", /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(factory_.ToAlignedPath("index-0", file_meta), "file:/test/bucket-0/index-0"); @@ -135,7 +142,8 @@ TEST_F(DataFilePathFactoryTest, TestCollectFiles) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(factory_.CollectFiles(file_meta), std::vector({"/tmp/data-0.txt"})); file_meta->extra_files = {"data-0.index", "data-1.index"}; diff --git a/src/paimon/core/io/data_file_writer.cpp b/src/paimon/core/io/data_file_writer.cpp index 4ed3e040c..9275fcf25 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 60cc808a9..f56f34956 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 000000000..ccea898a7 --- /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 b929dde83..07195ab7f 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 c727b47d0..cab942f0e 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_evaluator_test.cpp b/src/paimon/core/io/file_index_evaluator_test.cpp index d9730ade1..1824564e6 100644 --- a/src/paimon/core/io/file_index_evaluator_test.cpp +++ b/src/paimon/core/io/file_index_evaluator_test.cpp @@ -319,7 +319,8 @@ TEST_F(FileIndexEvaluatorTest, TestEvaluateEmbeddingIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); CheckResult(/*data_file_path_factory=*/nullptr, data_file_meta); } @@ -339,7 +340,8 @@ TEST_F(FileIndexEvaluatorTest, TestEvaluateNoEmbeddingIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto data_file_path_factory = std::make_shared(); ASSERT_OK(data_file_path_factory->Init(path + "/bucket-0/", /*format_identifier=*/"orc", /*data_file_prefix=*/"data-", nullptr)); @@ -374,7 +376,8 @@ TEST_F(FileIndexEvaluatorTest, TestTimestampType) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto data_file_path_factory = std::make_shared(); ASSERT_OK(data_file_path_factory->Init(path + "/bucket-0/", /*format_identifier=*/"orc", /*data_file_prefix=*/"data-", nullptr)); @@ -408,7 +411,8 @@ TEST_F(FileIndexEvaluatorTest, TestInvalidEvaluate) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto predicate = PredicateBuilder::IsNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_NOK_WITH_MSG( 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 000000000..a9587363a --- /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 000000000..7b7c019bf --- /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 000000000..157203f95 --- /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 9393c7c3c..881712ebe 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,15 @@ 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, + /*column_max_sequence_numbers=*/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 e1e3fd92c..eb7a2efca 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_factories.cpp b/src/paimon/core/io/key_value_data_file_writer_factories.cpp new file mode 100644 index 000000000..5aa5a8caa --- /dev/null +++ b/src/paimon/core/io/key_value_data_file_writer_factories.cpp @@ -0,0 +1,51 @@ +/* + * 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/key_value_data_file_writer_factories.h" + +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" +#include "paimon/core/io/key_value_data_file_writer_factory.h" +#include "paimon/core/io/shredding_key_value_data_file_writer_factory.h" + +namespace paimon { + +Result> +KeyValueDataFileWriterFactories::Create(const CoreOptions& options, int64_t schema_id, + const std::shared_ptr& write_schema, + int32_t level, FileSource file_source, + const std::vector& primary_keys, + const std::shared_ptr& path_factory, + bool create_stats_extractor, bool is_changelog, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan_factory, + ShreddingWritePlanFactories::SelectActive(options, write_schema, pool)); + std::shared_ptr writer_factory; + if (plan_factory != nullptr) { + writer_factory = std::make_shared( + options, schema_id, write_schema, level, file_source, primary_keys, path_factory, + create_stats_extractor, plan_factory, is_changelog, pool); + } else { + writer_factory = std::make_shared( + options, schema_id, write_schema, level, file_source, primary_keys, path_factory, + create_stats_extractor, is_changelog, pool); + } + return writer_factory; +} + +} // namespace paimon diff --git a/src/paimon/core/io/key_value_data_file_writer_factories.h b/src/paimon/core/io/key_value_data_file_writer_factories.h new file mode 100644 index 000000000..335bb94d4 --- /dev/null +++ b/src/paimon/core/io/key_value_data_file_writer_factories.h @@ -0,0 +1,56 @@ +/* + * 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/core/io/data_file_meta.h" +#include "paimon/core/io/single_file_writer_factory.h" +#include "paimon/core/key_value.h" +#include "paimon/core/manifest/file_source.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class CoreOptions; +class DataFilePathFactory; +class MemoryPool; + +/// Creates the appropriate key-value data file writer factory for the configured write schema. +class KeyValueDataFileWriterFactories { + public: + using WriterFactory = SingleFileWriterFactory>; + + static Result> Create( + const CoreOptions& options, int64_t schema_id, + const std::shared_ptr& write_schema, int32_t level, FileSource file_source, + const std::vector& primary_keys, + const std::shared_ptr& path_factory, bool create_stats_extractor, + bool is_changelog, const std::shared_ptr& pool); +}; + +} // 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 07d50b980..e0640ab61 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 @@ -20,10 +20,12 @@ #include "paimon/core/io/key_value_data_file_writer_factory.h" #include +#include #include #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" @@ -36,14 +38,15 @@ KeyValueDataFileWriterFactory::KeyValueDataFileWriterFactory( const std::shared_ptr& write_schema, int32_t level, FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, bool create_stats_extractor, - const std::shared_ptr& pool) + bool is_changelog, const std::shared_ptr& pool) : DataFileWriterFactory(options, schema_id, pool), write_schema_(write_schema), level_(level), file_source_(file_source), primary_keys_(primary_keys), path_factory_(path_factory), - create_stats_extractor_(create_stats_extractor) {} + create_stats_extractor_(create_stats_extractor), + is_changelog_(is_changelog) {} Result>>> KeyValueDataFileWriterFactory::CreateWriter() const { @@ -53,17 +56,50 @@ KeyValueDataFileWriterFactory::CreateWriter() const { return Status::OK(); }; - auto format = options_.GetWriteFileFormat(level_); + std::shared_ptr format = GetFileFormat(); PAIMON_ASSIGN_OR_RAISE(WriterResources resources, CreateWriterResources(*format, write_schema_, create_stats_extractor_)); auto writer = std::make_unique( - options_.GetWriteFileCompression(level_), std::move(converter), schema_id_, level_, - file_source_, primary_keys_, resources.stats_extractor, write_schema_, - path_factory_->IsExternalPath(), pool_); - PAIMON_RETURN_NOT_OK( - writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); + GetFileCompression(), std::move(converter), schema_id_, level_, file_source_, primary_keys_, + resources.stats_extractor, write_schema_, path_factory_->IsExternalPath(), pool_); + if (!is_changelog_) { + 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(), NewFilePath(format->Identifier()), + resources.writer_builder)); return std::unique_ptr>>( std::move(writer)); } +std::shared_ptr KeyValueDataFileWriterFactory::GetFileFormat() const { + if (is_changelog_) { + std::shared_ptr changelog_format = options_.GetChangelogFileFormat(); + if (changelog_format) { + return changelog_format; + } + } + return options_.GetWriteFileFormat(level_); +} + +std::string KeyValueDataFileWriterFactory::GetFileCompression() const { + if (is_changelog_) { + std::optional changelog_compression = options_.GetChangelogFileCompression(); + if (changelog_compression) { + return changelog_compression.value(); + } + } + return options_.GetWriteFileCompression(level_); +} + +std::string KeyValueDataFileWriterFactory::NewFilePath(const std::string& format_identifier) const { + if (is_changelog_) { + return path_factory_->NewChangelogPath(options_.ChangelogFilePrefix(), format_identifier); + } + return path_factory_->NewPath(); +} + } // namespace paimon diff --git a/src/paimon/core/io/key_value_data_file_writer_factory.h b/src/paimon/core/io/key_value_data_file_writer_factory.h index 6ac50aba3..533b84c6e 100644 --- a/src/paimon/core/io/key_value_data_file_writer_factory.h +++ b/src/paimon/core/io/key_value_data_file_writer_factory.h @@ -38,6 +38,7 @@ namespace paimon { class CoreOptions; class DataFilePathFactory; +class FileFormat; class MemoryPool; class KeyValueDataFileWriterFactory @@ -49,19 +50,24 @@ class KeyValueDataFileWriterFactory FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, - bool create_stats_extractor, + bool create_stats_extractor, bool is_changelog, const std::shared_ptr& pool); Result>>> CreateWriter() const override; protected: + std::shared_ptr GetFileFormat() const; + std::string GetFileCompression() const; + std::string NewFilePath(const std::string& format_identifier) const; + std::shared_ptr write_schema_; int32_t level_; FileSource file_source_; std::vector primary_keys_; std::shared_ptr path_factory_; bool create_stats_extractor_; + bool is_changelog_; }; } // namespace paimon 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 1b6b71c69..858e302e8 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 @@ -20,6 +20,7 @@ #include #include +#include #include "arrow/api.h" #include "arrow/array/array_nested.h" diff --git a/src/paimon/core/io/rolling_blob_file_writer_test.cpp b/src/paimon/core/io/rolling_blob_file_writer_test.cpp index 654c91266..b88f4a329 100644 --- a/src/paimon/core/io/rolling_blob_file_writer_test.cpp +++ b/src/paimon/core/io/rolling_blob_file_writer_test.cpp @@ -57,7 +57,8 @@ TEST_F(RollingBlobFileWriterTest, ValidateFileConsistency) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"f0", "f1"})); + /*write_cols=*/std::vector({"f0", "f1"}), + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-xxx.blob", /*file_size=*/764, /*row_count=*/3, @@ -70,7 +71,8 @@ TEST_F(RollingBlobFileWriterTest, ValidateFileConsistency) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"blob"})); + /*write_cols=*/std::vector({"blob"}), + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-xxx.blob", /*file_size=*/3023, /*row_count=*/1, @@ -83,7 +85,8 @@ TEST_F(RollingBlobFileWriterTest, ValidateFileConsistency) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/3, - /*write_cols=*/std::vector({"blob"})); + /*write_cols=*/std::vector({"blob"}), + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_OK(RollingBlobFileWriter::ValidateFileConsistency(file_meta1, {file_meta2, file_meta3})); ASSERT_NOK_WITH_MSG(RollingBlobFileWriter::ValidateFileConsistency(file_meta1, {file_meta2}), "This is a bug: The row count of main file and blob file does not match."); 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 6e4843bbb..0e4e82199 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 8ac583ee0..f56ba8d14 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" @@ -36,10 +37,11 @@ ShreddingKeyValueDataFileWriterFactory::ShreddingKeyValueDataFileWriterFactory( const std::shared_ptr& write_schema, int32_t level, FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, bool create_stats_extractor, - const std::shared_ptr& plan_factory, + const std::shared_ptr& plan_factory, bool is_changelog, const std::shared_ptr& pool) : KeyValueDataFileWriterFactory(options, schema_id, write_schema, level, file_source, - primary_keys, path_factory, create_stats_extractor, pool), + primary_keys, path_factory, create_stats_extractor, + is_changelog, pool), plan_factory_(plan_factory) {} Result>>> @@ -47,7 +49,7 @@ ShreddingKeyValueDataFileWriterFactory::CreateWriter() const { if (!plan_factory_) { return Status::Invalid("Shredding key-value writer requires a write-plan factory."); } - const std::string format_identifier = options_.GetWriteFileFormat(level_)->Identifier(); + const std::string format_identifier = GetFileFormat()->Identifier(); if (plan_factory_->ShouldInferWritePlan()) { auto create_inner = [this](const std::shared_ptr& converter) { return CreateShreddedWriter(converter); @@ -73,7 +75,8 @@ ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter( [factory = plan_factory_, converter]() { return factory->OnFileCompleted(converter); }); return writer; } - auto format = options_.GetWriteFileFormat(level_); + std::shared_ptr format = GetFileFormat(); + std::string compression = GetFileCompression(); std::shared_ptr file_schema = converter->GetPhysicalSchema(); std::function batch_converter = [converter](KeyValueBatch key_value_batch, ::ArrowArray* array) -> Status { @@ -85,13 +88,19 @@ ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter( PAIMON_ASSIGN_OR_RAISE(WriterResources resources, CreateWriterResources(*format, file_schema, create_stats_extractor_)); auto writer = std::make_unique( - options_.GetWriteFileCompression(level_), std::move(batch_converter), schema_id_, level_, - file_source_, primary_keys_, resources.stats_extractor, file_schema, - path_factory_->IsExternalPath(), pool_); - PAIMON_RETURN_NOT_OK( - writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); + compression, std::move(batch_converter), schema_id_, level_, file_source_, primary_keys_, + resources.stats_extractor, file_schema, path_factory_->IsExternalPath(), pool_); + if (!is_changelog_) { + 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(), NewFilePath(format->Identifier()), + resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer = - plan_factory_->CreateMetadataFinalizer(converter, options_.GetWriteFileCompression(level_)); + plan_factory_->CreateMetadataFinalizer(converter, compression); if (finalizer) { writer->SetMetadataFinalizer(std::move(finalizer)); } diff --git a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h index fb967588e..e32dc8a04 100644 --- a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h +++ b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h @@ -41,7 +41,7 @@ class ShreddingKeyValueDataFileWriterFactory : public KeyValueDataFileWriterFact const std::shared_ptr& write_schema, int32_t level, FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, bool create_stats_extractor, - const std::shared_ptr& plan_factory, + const std::shared_ptr& plan_factory, bool is_changelog, const std::shared_ptr& pool); Result>>> diff --git a/src/paimon/core/io/single_file_writer.h b/src/paimon/core/io/single_file_writer.h index 99507b579..6db3a699c 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/vector_file_batch_reader.cpp b/src/paimon/core/io/vector_file_batch_reader.cpp new file mode 100644 index 000000000..f16025c42 --- /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 000000000..b7eab263c --- /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 000000000..e334e62aa --- /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 053405b89..2389cd8d7 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 7438895f5..4a71a1b68 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 2aa2db524..a31f4bf73 100644 --- a/src/paimon/core/manifest/manifest_entry_serializer_test.cpp +++ b/src/paimon/core/manifest/manifest_entry_serializer_test.cpp @@ -43,7 +43,8 @@ class ManifestEntrySerializerTest : public testing::Test { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/3, /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::optional(), - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } }; TEST_F(ManifestEntrySerializerTest, TestToFromRow) { @@ -55,12 +56,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_entry_writer_test.cpp b/src/paimon/core/manifest/manifest_entry_writer_test.cpp index 3c3d00d7e..6c9291b71 100644 --- a/src/paimon/core/manifest/manifest_entry_writer_test.cpp +++ b/src/paimon/core/manifest/manifest_entry_writer_test.cpp @@ -64,7 +64,7 @@ class ManifestEntryWriterTest : public ::testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, first_row_id, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return {FileKind::Add(), BinaryRowGenerator::GenerateRow({10}, pool_.get()), /*bucket=*/0, /*total_buckets=*/-1, meta}; } @@ -126,7 +126,8 @@ TEST_F(ManifestEntryWriterTest, TestSimple) { /*creation_time=*/Timestamp(1743525392885ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto meta2 = std::make_shared( "data-5858a84b-7081-4618-b828-ae3918c5e1f6-0.orc", /*file_size=*/943, /*row_count=*/4, @@ -144,7 +145,8 @@ TEST_F(ManifestEntryWriterTest, TestSimple) { /*creation_time=*/Timestamp(1743525392921ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto entry1 = ManifestEntry(FileKind::Add(), BinaryRowGenerator::GenerateRow({10}, pool_.get()), 0, 2, meta1); diff --git a/src/paimon/core/manifest/manifest_file.cpp b/src/paimon/core/manifest/manifest_file.cpp index 22f2681f6..a556b82ea 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,26 @@ 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(); + ColumnarRow row(fields, pool_, /*row_id=*/0); + for (int64_t i = 0; i < batch->length(); i++) { + row.SetRowId(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 d34764b5f..0211e14d1 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 8f6e0b2ea..df78fc00d 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; } @@ -149,7 +153,7 @@ TEST_F(ManifestFileTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry1 = ManifestEntry(FileKind::Delete(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta1); @@ -167,7 +171,7 @@ TEST_F(ManifestFileTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry2 = ManifestEntry(FileKind::Delete(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta2); @@ -185,7 +189,7 @@ TEST_F(ManifestFileTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry3 = ManifestEntry(FileKind::Delete(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta3); @@ -203,7 +207,7 @@ TEST_F(ManifestFileTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry4 = ManifestEntry(FileKind::Delete(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta4); @@ -221,7 +225,7 @@ TEST_F(ManifestFileTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry5 = ManifestEntry(FileKind::Add(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta5); @@ -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 = @@ -335,7 +437,7 @@ TEST_F(ManifestFileTest, TestWithNullCount) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry1 = ManifestEntry(FileKind::Add(), BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_buckets=*/2, file_meta1); @@ -362,7 +464,7 @@ TEST_F(ManifestFileTest, TestWithNullCount) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry2 = ManifestEntry(FileKind::Add(), BinaryRowGenerator::GenerateRow({20}, pool.get()), /*bucket=*/0, /*total_buckets=*/2, file_meta2); @@ -399,13 +501,16 @@ TEST_F(ManifestFileTest, TestManifestFileCompatibleWithJavaPaimon09) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry = ManifestEntry(FileKind::Add(), /*partition=*/BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/-1, file_meta); 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) { @@ -435,13 +540,16 @@ TEST_F(ManifestFileTest, TestManifestFileCompatibleWithJavaPaimon11) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto manifest_entry = ManifestEntry(FileKind::Add(), /*partition=*/BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/-1, file_meta); 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/manifest/manifest_list.h b/src/paimon/core/manifest/manifest_list.h index 31959a5c9..ff67e8080 100644 --- a/src/paimon/core/manifest/manifest_list.h +++ b/src/paimon/core/manifest/manifest_list.h @@ -113,8 +113,7 @@ class ManifestList : public ObjectsFile { const std::optional& changelog_manifest_list = snapshot.ChangelogManifestList(); if (changelog_manifest_list) { - return Status::NotImplemented("do not support read changelog manifest list"); - // return Read(changelog_manifest_list.value(), /*filter=*/nullptr, manifests); + return Read(changelog_manifest_list.value(), /*filter=*/nullptr, manifests); } else { return Status::OK(); } diff --git a/src/paimon/core/manifest/manifest_list_test.cpp b/src/paimon/core/manifest/manifest_list_test.cpp index 889f18f2b..23949bfdd 100644 --- a/src/paimon/core/manifest/manifest_list_test.cpp +++ b/src/paimon/core/manifest/manifest_list_test.cpp @@ -25,6 +25,7 @@ #include "gtest/gtest.h" #include "paimon/core/core_options.h" #include "paimon/core/manifest/manifest_file_meta.h" +#include "paimon/core/snapshot.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/format/file_format.h" @@ -141,6 +142,33 @@ TEST_F(ManifestListTest, TestEmptyManifestList) { ASSERT_EQ(manifest_file_metas.size(), 0); } +TEST_F(ManifestListTest, TestReadChangelogManifests) { + auto pool = GetDefaultPool(); + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto manifest_list = CreateManifestList("orc", dir->Str(), pool); + ManifestFileMeta expected_meta( + "changelog-manifest", /*file_size=*/100, /*num_added_files=*/1, + /*num_deleted_files=*/0, SimpleStats::EmptyStats(), /*schema_id=*/0, + /*min_bucket=*/0, /*max_bucket=*/0, /*min_level=*/0, /*max_level=*/0, + /*min_row_id=*/std::nullopt, /*max_row_id=*/std::nullopt); + ASSERT_OK_AND_ASSIGN(auto changelog_manifest_list, manifest_list->Write({expected_meta})); + Snapshot snapshot( + /*id=*/1, /*schema_id=*/0, /*base_manifest_list=*/"", + /*base_manifest_list_size=*/std::nullopt, /*delta_manifest_list=*/"", + /*delta_manifest_list_size=*/std::nullopt, + /*changelog_manifest_list=*/changelog_manifest_list.first, + /*changelog_manifest_list_size=*/changelog_manifest_list.second, + /*index_manifest=*/std::nullopt, /*commit_user=*/"user", /*commit_identifier=*/1, + Snapshot::CommitKind::Append(), /*time_millis=*/0, /*total_record_count=*/1, + /*delta_record_count=*/1, /*changelog_record_count=*/1, /*watermark=*/std::nullopt, + /*statistics=*/std::nullopt, /*properties=*/std::nullopt, /*next_row_id=*/std::nullopt); + + std::vector actual_metas; + ASSERT_OK(manifest_list->ReadChangelogManifests(snapshot, &actual_metas)); + ASSERT_EQ(std::vector({expected_meta}), actual_metas); +} + TEST_F(ManifestListTest, TestManifestListCompatibleWithJavaPaimon09) { auto pool = GetDefaultPool(); auto manifest_file_metas = ReadManifestFileMeta("avro", paimon::test::GetDataDir() + "/avro", diff --git a/src/paimon/core/manifest/partition_entry_test.cpp b/src/paimon/core/manifest/partition_entry_test.cpp index 8b7dd6269..213e0dbaa 100644 --- a/src/paimon/core/manifest/partition_entry_test.cpp +++ b/src/paimon/core/manifest/partition_entry_test.cpp @@ -50,7 +50,7 @@ class PartitionEntryTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } }; diff --git a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp index 3efd40f64..4f632425b 100644 --- a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.cpp @@ -17,7 +17,142 @@ */ #include "paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h" + +#include + +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/io/key_value_meta_projection_consumer.h" +#include "paimon/core/io/row_to_arrow_array_converter.h" +#include "paimon/format/file_format.h" namespace paimon { + +namespace { + +using CancellationChecker = std::function; + +class ChangelogCompactionBatchProducer : public AsyncKeyValueBatchProducer { + public: + ChangelogCompactionBatchProducer( + std::unique_ptr&& sort_merge_reader, int32_t write_batch_size, + std::shared_ptr>&& merge_function_wrapper, + const FieldsComparator::FieldComparatorFunc& key_comparator, + const CancellationChecker& cancellation_checker, bool drop_delete, bool produce_data, + bool produce_changelog) + : sort_merge_reader_(std::move(sort_merge_reader)), + write_batch_size_(NormalizeProjectionBatchSize(write_batch_size)), + merge_function_wrapper_(std::move(merge_function_wrapper)), + key_comparator_(key_comparator), + cancellation_checker_(cancellation_checker), + drop_delete_(drop_delete), + produce_data_(produce_data), + produce_changelog_(produce_changelog) {} + + Status Produce(AsyncKeyValueBatchSink* sink) override { + std::vector compact_buffer; + std::vector changelog_buffer; + compact_buffer.reserve(write_batch_size_); + changelog_buffer.reserve(write_batch_size_); + + auto flush = [&](AsyncKeyValueBatchType type, std::vector* buffer) -> Status { + if (buffer->empty()) { + return Status::OK(); + } + std::vector rows = std::move(*buffer); + buffer->clear(); + buffer->reserve(write_batch_size_); + return sink->Write(type, std::move(rows)); + }; + + auto emit_result = [&](ChangelogResult&& result) -> Status { + if (produce_data_ && result.result && + (!drop_delete_ || result.result->value_kind->IsAdd())) { + compact_buffer.emplace_back(std::move(result.result).value()); + if (static_cast(compact_buffer.size()) >= write_batch_size_) { + PAIMON_RETURN_NOT_OK(flush(AsyncKeyValueBatchType::DATA, &compact_buffer)); + } + } + if (produce_changelog_) { + for (auto& changelog : result.changelogs) { + changelog_buffer.emplace_back(std::move(changelog)); + if (static_cast(changelog_buffer.size()) >= write_batch_size_) { + PAIMON_RETURN_NOT_OK( + flush(AsyncKeyValueBatchType::CHANGELOG, &changelog_buffer)); + } + } + } + return Status::OK(); + }; + + std::shared_ptr current_key; + auto finish_group = [&]() -> Status { + if (!current_key) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE(std::optional result, + merge_function_wrapper_->GetResult()); + current_key.reset(); + if (result) { + PAIMON_RETURN_NOT_OK(emit_result(std::move(result).value())); + } + return Status::OK(); + }; + + while (true) { + if (cancellation_checker_()) { + return Status::Cancelled("Compaction is cancelled"); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + sort_merge_reader_->NextBatch()); + if (!iterator) { + break; + } + while (true) { + PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); + if (!has_next) { + break; + } + KeyValue key_value = iterator->Next(); + if (current_key && key_comparator_(*current_key, *key_value.key) != 0) { + PAIMON_RETURN_NOT_OK(finish_group()); + } + if (!current_key) { + merge_function_wrapper_->Reset(); + current_key = key_value.key; + } + PAIMON_RETURN_NOT_OK(merge_function_wrapper_->Add(std::move(key_value))); + } + } + PAIMON_RETURN_NOT_OK(finish_group()); + sort_merge_reader_->Close(); + closed_ = true; + PAIMON_RETURN_NOT_OK(flush(AsyncKeyValueBatchType::DATA, &compact_buffer)); + PAIMON_RETURN_NOT_OK(flush(AsyncKeyValueBatchType::CHANGELOG, &changelog_buffer)); + return Status::OK(); + } + + void Close() override { + if (closed_) { + return; + } + sort_merge_reader_->Close(); + closed_ = true; + } + + private: + std::unique_ptr sort_merge_reader_; + int32_t write_batch_size_; + std::shared_ptr> merge_function_wrapper_; + FieldsComparator::FieldComparatorFunc key_comparator_; + CancellationChecker cancellation_checker_; + bool drop_delete_; + bool produce_data_; + bool produce_changelog_; + bool closed_ = false; +}; + +} // namespace + ChangelogMergeTreeRewriter::ChangelogMergeTreeRewriter( int32_t max_level, bool force_drop_delete, const BinaryRow& partition, int32_t bucket, int64_t schema_id, const std::vector& trimmed_primary_keys, @@ -26,14 +161,18 @@ ChangelogMergeTreeRewriter::ChangelogMergeTreeRewriter( const std::shared_ptr& path_factory_cache, std::unique_ptr&& merge_file_split_read, MergeFunctionWrapperFactory merge_function_wrapper_factory, - const std::shared_ptr& cancellation_controller, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, + bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& pool) : MergeTreeCompactRewriter( partition, bucket, schema_id, trimmed_primary_keys, options, data_schema, write_schema, std::move(dv_factory), path_factory_cache, std::move(merge_file_split_read), std::move(merge_function_wrapper_factory), cancellation_controller, pool), max_level_(max_level), - force_drop_delete_(force_drop_delete) {} + force_drop_delete_(force_drop_delete), + changelog_merge_function_wrapper_factory_( + std::move(changelog_merge_function_wrapper_factory)), + produce_changelog_(produce_changelog) {} Result ChangelogMergeTreeRewriter::Rewrite( int32_t output_level, bool drop_delete, const std::vector>& sections) { @@ -78,31 +217,78 @@ Result ChangelogMergeTreeRewriter::RewriteOrProduceChangelog( bool rewrite_compact_file) { PAIMON_ASSIGN_OR_RAISE(MergeTreeCompactRewriter::KeyValueConsumerCreator create_consumer, GenerateKeyValueConsumer()); - std::vector> reader_holders; - auto before = ExtractFilesFromSections(sections); std::unique_ptr compact_file_writer; if (rewrite_compact_file) { PAIMON_ASSIGN_OR_RAISE(compact_file_writer, CreateRollingRowWriter(output_level)); } - // TODO(xinyu.lxy): produce changelog + std::unique_ptr changelog_file_writer; + if (produce_changelog_) { + PAIMON_ASSIGN_OR_RAISE(changelog_file_writer, CreateRollingChangelogWriter(output_level)); + } + + std::vector> reader_holders; ScopeGuard write_guard([&]() -> void { if (compact_file_writer) { compact_file_writer->Abort(); + compact_file_writer.reset(); + } + if (changelog_file_writer) { + changelog_file_writer->Abort(); + changelog_file_writer.reset(); } - merge_file_split_read_.reset(); for (const auto& reader : reader_holders) { reader->Close(); } + merge_file_split_read_.reset(); }); + bool produce_data = compact_file_writer != nullptr; + bool produce_changelog = changelog_file_writer != nullptr; + FieldsComparator::FieldComparatorFunc key_comparator = [this](const InternalRow& lhs, + const InternalRow& rhs) { + return merge_file_split_read_->GetKeyComparator()->CompareTo(lhs, rhs); + }; + CancellationChecker cancellation_checker = [this]() { return IsCancelled(); }; + for (const auto& section : sections) { - PAIMON_RETURN_NOT_OK(MergeReadAndWrite(output_level, drop_delete, section, create_consumer, - compact_file_writer.get(), &reader_holders)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + CreateRawSortMergeReaderForSection(section)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> merge_function_wrapper, + changelog_merge_function_wrapper_factory_(output_level)); + std::unique_ptr producer = + std::make_unique( + std::move(sort_merge_reader), options_.GetWriteBatchSize(), + std::move(merge_function_wrapper), key_comparator, cancellation_checker, + drop_delete, produce_data, produce_changelog); + auto producer_and_consumer = + std::make_shared>( + std::move(producer), create_consumer, /*consumer_thread_num=*/1); + reader_holders.emplace_back(producer_and_consumer); + + while (true) { + if (IsCancelled()) { + return Status::Cancelled("Compaction is cancelled"); + } + PAIMON_ASSIGN_OR_RAISE(AsyncKeyValueResultBatch output, + producer_and_consumer->NextBatchWithType()); + if (output.result.batch == nullptr) { + break; + } + if (output.type == AsyncKeyValueBatchType::DATA) { + PAIMON_RETURN_NOT_OK(compact_file_writer->Write(std::move(output.result))); + } else { + PAIMON_RETURN_NOT_OK(changelog_file_writer->Write(std::move(output.result))); + } + } } if (compact_file_writer) { PAIMON_RETURN_NOT_OK(compact_file_writer->Close()); } + if (changelog_file_writer) { + PAIMON_RETURN_NOT_OK(changelog_file_writer->Close()); + } std::vector> after; if (compact_file_writer) { PAIMON_ASSIGN_OR_RAISE(after, compact_file_writer->GetResult()); @@ -118,8 +304,12 @@ Result ChangelogMergeTreeRewriter::RewriteOrProduceChangelog( NotifyRewriteCompactBefore(before); } PAIMON_ASSIGN_OR_RAISE(after, NotifyRewriteCompactAfter(after)); + std::vector> changelog_files; + if (changelog_file_writer) { + PAIMON_ASSIGN_OR_RAISE(changelog_files, changelog_file_writer->GetResult()); + } write_guard.Release(); - return CompactResult(before, after); + return CompactResult(before, after, changelog_files); } } // namespace paimon diff --git a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h index c5d8e8914..c080b30a7 100644 --- a/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h +++ b/src/paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h @@ -22,11 +22,15 @@ #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/changelog_result.h" #include "paimon/core/mergetree/compact/merge_tree_compact_rewriter.h" namespace paimon { /// A `MergeTreeCompactRewriter` which produces changelog files while performing compaction. class ChangelogMergeTreeRewriter : public MergeTreeCompactRewriter { public: + using ChangelogMergeFunctionWrapperFactory = + std::function>>(int32_t)>; + Result Rewrite(int32_t output_level, bool drop_delete, const std::vector>& sections) override; @@ -42,6 +46,8 @@ class ChangelogMergeTreeRewriter : public MergeTreeCompactRewriter { const std::shared_ptr& path_factory_cache, std::unique_ptr&& merge_file_split_read, MergeFunctionWrapperFactory merge_function_wrapper_factory, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, + bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& pool); @@ -85,5 +91,8 @@ class ChangelogMergeTreeRewriter : public MergeTreeCompactRewriter { Result RewriteOrProduceChangelog( int32_t output_level, const std::vector>& sections, bool drop_delete, bool rewrite_compact_file); + + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory_; + bool produce_changelog_; }; } // namespace paimon diff --git a/src/paimon/core/mergetree/compact/changelog_result.h b/src/paimon/core/mergetree/compact/changelog_result.h new file mode 100644 index 000000000..778183ed7 --- /dev/null +++ b/src/paimon/core/mergetree/compact/changelog_result.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 +#include + +#include "paimon/core/key_value.h" + +namespace paimon { + +/// The result of merging all records with the same primary key while producing changelog. +struct ChangelogResult { + std::optional result; + std::vector changelogs; +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/compact_strategy_test.cpp b/src/paimon/core/mergetree/compact/compact_strategy_test.cpp index a0abebe33..e8508a794 100644 --- a/src/paimon/core/mergetree/compact/compact_strategy_test.cpp +++ b/src/paimon/core/mergetree/compact/compact_strategy_test.cpp @@ -38,7 +38,7 @@ class CompactStrategyTest : public testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return {level, SortedRun::FromSingle(file_meta)}; } diff --git a/src/paimon/core/mergetree/compact/early_full_compaction_test.cpp b/src/paimon/core/mergetree/compact/early_full_compaction_test.cpp index a4cc9c663..37cc061f4 100644 --- a/src/paimon/core/mergetree/compact/early_full_compaction_test.cpp +++ b/src/paimon/core/mergetree/compact/early_full_compaction_test.cpp @@ -55,7 +55,7 @@ class EarlyFullCompactionTest : public testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return {level, SortedRun::FromSingle(file_meta)}; } diff --git a/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper.h b/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper.h index 319ed5ab6..e0d7e6d7f 100644 --- a/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper.h +++ b/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper.h @@ -24,7 +24,9 @@ #include #include +#include "paimon/common/data/serializer/row_compacted_serializer.h" #include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/changelog_result.h" #include "paimon/core/mergetree/compact/first_row_merge_function.h" #include "paimon/core/mergetree/compact/merge_function_wrapper.h" #include "paimon/result.h" @@ -32,12 +34,15 @@ namespace paimon { /// Wrapper for `MergeFunction`s to produce changelog by lookup for first row. -class FirstRowMergeFunctionWrapper : public MergeFunctionWrapper { +class FirstRowMergeFunctionWrapper : public MergeFunctionWrapper { public: FirstRowMergeFunctionWrapper( std::unique_ptr&& merge_function, - std::function(const std::shared_ptr&)> contains) - : merge_function_(std::move(merge_function)), contains_(std::move(contains)) {} + std::function(const std::shared_ptr&)> contains, + std::unique_ptr&& value_serializer) + : merge_function_(std::move(merge_function)), + contains_(std::move(contains)), + value_serializer_(std::move(value_serializer)) {} void Reset() override { merge_function_->Reset(); @@ -47,11 +52,13 @@ class FirstRowMergeFunctionWrapper : public MergeFunctionWrapper { return merge_function_->Add(std::move(kv)); } - Result> GetResult() override { + Result> GetResult() override { PAIMON_ASSIGN_OR_RAISE(std::optional result, merge_function_->GetResult()); + ChangelogResult changelog_result; if (merge_function_->ContainsHighLevel()) { + changelog_result.result = std::move(result); Reset(); - return result; + return std::optional(std::move(changelog_result)); } if (!result) { Reset(); @@ -62,17 +69,27 @@ class FirstRowMergeFunctionWrapper : public MergeFunctionWrapper { if (contains) { // empty Reset(); - return std::optional(); + return std::optional(std::move(changelog_result)); } - // new record, output changelog - // TODO(xinyu.lxy) support changelog + // TODO(lisizhuo.lsz): avoid serialize & deserialize here. + if (value_serializer_) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, + value_serializer_->SerializeToBytes(*result->value)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr changelog_value, + value_serializer_->Deserialize(bytes)); + changelog_result.changelogs.emplace_back(result->value_kind, result->sequence_number, + result->level, result->key, + std::move(changelog_value)); + } + changelog_result.result = std::move(result); Reset(); - return result; + return std::optional(std::move(changelog_result)); } private: std::unique_ptr merge_function_; std::function(const std::shared_ptr&)> contains_; + std::unique_ptr value_serializer_; }; } // namespace paimon diff --git a/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper_test.cpp b/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper_test.cpp index 36a2ef01e..47a2d38ed 100644 --- a/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper_test.cpp +++ b/src/paimon/core/mergetree/compact/first_row_merge_function_wrapper_test.cpp @@ -29,6 +29,15 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { +std::unique_ptr CreateValueSerializer( + const std::shared_ptr& pool) { + return RowCompactedSerializer::Create(arrow::schema({arrow::field("value", arrow::int32())}), + pool) + .value(); +} +} // namespace + TEST(FirstRowMergeFunctionWrapperTest, TestSimple) { auto pool = GetDefaultPool(); KeyValue kv1(RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, /*key=*/ @@ -45,14 +54,17 @@ TEST(FirstRowMergeFunctionWrapperTest, TestSimple) { auto contains = [](const std::shared_ptr& row) { return true; }; - FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains)); + FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains), + CreateValueSerializer(pool)); wrapper.Reset(); ASSERT_OK(wrapper.Add(std::move(kv1))); ASSERT_OK(wrapper.Add(std::move(kv2))); ASSERT_OK(wrapper.Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper.GetResult()); ASSERT_TRUE(result); - ASSERT_EQ(result.value().sequence_number, 0); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 0); + ASSERT_TRUE(result->changelogs.empty()); } TEST(FirstRowMergeFunctionWrapperTest, TestAllLevel0WithContain) { @@ -71,13 +83,16 @@ TEST(FirstRowMergeFunctionWrapperTest, TestAllLevel0WithContain) { auto contains = [](const std::shared_ptr& row) { return true; }; - FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains)); + FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains), + CreateValueSerializer(pool)); wrapper.Reset(); ASSERT_OK(wrapper.Add(std::move(kv1))); ASSERT_OK(wrapper.Add(std::move(kv2))); ASSERT_OK(wrapper.Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper.GetResult()); - ASSERT_FALSE(result); + ASSERT_TRUE(result); + ASSERT_FALSE(result->result); + ASSERT_TRUE(result->changelogs.empty()); } TEST(FirstRowMergeFunctionWrapperTest, TestAllLevel0WithoutContain) { @@ -96,14 +111,18 @@ TEST(FirstRowMergeFunctionWrapperTest, TestAllLevel0WithoutContain) { auto contains = [](const std::shared_ptr& row) { return false; }; - FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains)); + FirstRowMergeFunctionWrapper wrapper(std::move(mfunc), std::move(contains), + CreateValueSerializer(pool)); wrapper.Reset(); ASSERT_OK(wrapper.Add(std::move(kv1))); ASSERT_OK(wrapper.Add(std::move(kv2))); ASSERT_OK(wrapper.Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper.GetResult()); ASSERT_TRUE(result); - ASSERT_EQ(result.value().sequence_number, 0); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 0); + ASSERT_EQ(result->changelogs.size(), 1); + ASSERT_EQ(result->changelogs[0].sequence_number, 0); } } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/force_up_level0_compaction_test.cpp b/src/paimon/core/mergetree/compact/force_up_level0_compaction_test.cpp index 18a11229c..c1d26a569 100644 --- a/src/paimon/core/mergetree/compact/force_up_level0_compaction_test.cpp +++ b/src/paimon/core/mergetree/compact/force_up_level0_compaction_test.cpp @@ -39,7 +39,7 @@ class ForceUpLevel0CompactionTest : public testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return {level, SortedRun::FromSingle(file_meta)}; } diff --git a/src/paimon/core/mergetree/compact/internal_row_equalizer.h b/src/paimon/core/mergetree/compact/internal_row_equalizer.h new file mode 100644 index 000000000..22f15d8bd --- /dev/null +++ b/src/paimon/core/mergetree/compact/internal_row_equalizer.h @@ -0,0 +1,250 @@ +/* + * 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 + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/data_getters.h" +#include "paimon/common/data/internal_array.h" +#include "paimon/common/data/internal_map.h" +#include "paimon/common/data/internal_row.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon { + +/// Creates equality functions for internal rows, including nested values. +/// Java's RecordEqualiser also compares the RowKind embedded in InternalRow. This comparator +/// currently compares field values only. This does not affect the current lookup changelog results +/// because changelog kinds are tracked separately by KeyValue::value_kind. +class InternalRowEqualizer { + public: + static Result Create( + const std::shared_ptr& schema, + const std::vector& ignore_fields) { + std::set ignored(ignore_fields.begin(), ignore_fields.end()); + std::vector> equalizers; + for (int32_t i = 0; i < schema->num_fields(); ++i) { + if (ignored.find(schema->field(i)->name()) != ignored.end()) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(ValueEqualizer equalizer, + CreateValueEqualizer(schema->field(i)->type())); + equalizers.emplace_back(i, std::move(equalizer)); + } + return FieldsComparator::FieldComparatorFunc( + [equalizers = std::move(equalizers)](const InternalRow& lhs, const InternalRow& rhs) { + for (const auto& [field_idx, equalizer] : equalizers) { + if (!EqualAt(lhs, field_idx, rhs, field_idx, equalizer)) { + return 1; + } + } + return 0; + }); + } + + private: + using ValueEqualizer = + std::function; + + static bool EqualAt(const DataGetters& lhs, int32_t lhs_pos, const DataGetters& rhs, + int32_t rhs_pos, const ValueEqualizer& equalizer) { + bool lhs_null = lhs.IsNullAt(lhs_pos); + bool rhs_null = rhs.IsNullAt(rhs_pos); + if (lhs_null || rhs_null) { + return lhs_null == rhs_null; + } + return equalizer(lhs, lhs_pos, rhs, rhs_pos); + } + + static Result CreateValueEqualizer( + const std::shared_ptr& type) { + switch (type->id()) { + case arrow::Type::BOOL: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetBoolean(pos); }); + case arrow::Type::INT8: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetByte(pos); }); + case arrow::Type::INT16: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetShort(pos); }); + case arrow::Type::INT32: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetInt(pos); }); + case arrow::Type::DATE32: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetDate(pos); }); + case arrow::Type::INT64: + return PrimitiveEqualizer( + [](const DataGetters& row, int32_t pos) { return row.GetLong(pos); }); + case arrow::Type::FLOAT: + return ValueEqualizer([](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return FieldsComparator::CompareFloatingPoint(lhs.GetFloat(lhs_pos), + rhs.GetFloat(rhs_pos)) == 0; + }); + case arrow::Type::DOUBLE: + return ValueEqualizer([](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return FieldsComparator::CompareFloatingPoint(lhs.GetDouble(lhs_pos), + rhs.GetDouble(rhs_pos)) == 0; + }); + case arrow::Type::STRING: + case arrow::Type::BINARY: + return ValueEqualizer([](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return lhs.GetStringView(lhs_pos) == rhs.GetStringView(rhs_pos); + }); + case arrow::Type::TIMESTAMP: { + std::shared_ptr timestamp_type = + checked_pointer_cast(type); + int32_t precision = DateTimeUtils::GetPrecisionFromType(timestamp_type); + return ValueEqualizer([precision](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return lhs.GetTimestamp(lhs_pos, precision) == + rhs.GetTimestamp(rhs_pos, precision); + }); + } + case arrow::Type::DECIMAL128: { + std::shared_ptr decimal_type = + checked_pointer_cast(type); + int32_t precision = decimal_type->precision(); + int32_t scale = decimal_type->scale(); + return ValueEqualizer([precision, scale](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return lhs.GetDecimal(lhs_pos, precision, scale) + .CompareTo(rhs.GetDecimal(rhs_pos, precision, scale)) == 0; + }); + } + case arrow::Type::LIST: { + std::shared_ptr list_type = + checked_pointer_cast(type); + PAIMON_ASSIGN_OR_RAISE(ValueEqualizer element_equalizer, + CreateValueEqualizer(list_type->value_type())); + return ValueEqualizer([element_equalizer = std::move(element_equalizer)]( + const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + std::shared_ptr lhs_array = lhs.GetArray(lhs_pos); + std::shared_ptr rhs_array = rhs.GetArray(rhs_pos); + if (lhs_array->Size() != rhs_array->Size()) { + return false; + } + for (int32_t i = 0; i < lhs_array->Size(); ++i) { + if (!EqualAt(*lhs_array, i, *rhs_array, i, element_equalizer)) { + return false; + } + } + return true; + }); + } + case arrow::Type::MAP: { + std::shared_ptr map_type = + checked_pointer_cast(type); + PAIMON_ASSIGN_OR_RAISE(ValueEqualizer key_equalizer, + CreateValueEqualizer(map_type->key_type())); + PAIMON_ASSIGN_OR_RAISE(ValueEqualizer item_equalizer, + CreateValueEqualizer(map_type->item_type())); + return ValueEqualizer([key_equalizer = std::move(key_equalizer), + item_equalizer = std::move(item_equalizer)]( + const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + std::shared_ptr lhs_map = lhs.GetMap(lhs_pos); + std::shared_ptr rhs_map = rhs.GetMap(rhs_pos); + if (lhs_map->Size() != rhs_map->Size()) { + return false; + } + std::shared_ptr lhs_keys = lhs_map->KeyArray(); + std::shared_ptr rhs_keys = rhs_map->KeyArray(); + std::shared_ptr lhs_values = lhs_map->ValueArray(); + std::shared_ptr rhs_values = rhs_map->ValueArray(); + std::vector matched(rhs_map->Size(), false); + for (int32_t lhs_index = 0; lhs_index < lhs_map->Size(); ++lhs_index) { + bool found = false; + for (int32_t rhs_index = 0; rhs_index < rhs_map->Size(); ++rhs_index) { + if (matched[rhs_index]) { + continue; + } + if (EqualAt(*lhs_keys, lhs_index, *rhs_keys, rhs_index, + key_equalizer) && + EqualAt(*lhs_values, lhs_index, *rhs_values, rhs_index, + item_equalizer)) { + matched[rhs_index] = true; + found = true; + break; + } + } + if (!found) { + return false; + } + } + return true; + }); + } + case arrow::Type::STRUCT: { + std::shared_ptr struct_type = + checked_pointer_cast(type); + std::vector field_equalizers; + field_equalizers.reserve(struct_type->num_fields()); + for (const auto& field : struct_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(ValueEqualizer field_equalizer, + CreateValueEqualizer(field->type())); + field_equalizers.emplace_back(std::move(field_equalizer)); + } + int32_t field_count = struct_type->num_fields(); + return ValueEqualizer([field_equalizers = std::move(field_equalizers), field_count]( + const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + std::shared_ptr lhs_row = lhs.GetRow(lhs_pos, field_count); + std::shared_ptr rhs_row = rhs.GetRow(rhs_pos, field_count); + for (int32_t i = 0; i < field_count; ++i) { + if (!EqualAt(*lhs_row, i, *rhs_row, i, field_equalizers[i])) { + return false; + } + } + return true; + }); + } + default: + return Status::NotImplemented( + fmt::format("Do not support equality for type {}", type->ToString())); + } + } + + template + static ValueEqualizer PrimitiveEqualizer(Getter getter) { + return [getter = std::move(getter)](const DataGetters& lhs, int32_t lhs_pos, + const DataGetters& rhs, int32_t rhs_pos) { + return static_cast(getter(lhs, lhs_pos)) == static_cast(getter(rhs, rhs_pos)); + }; + } +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/internal_row_equalizer_test.cpp b/src/paimon/core/mergetree/compact/internal_row_equalizer_test.cpp new file mode 100644 index 000000000..80bc37b55 --- /dev/null +++ b/src/paimon/core/mergetree/compact/internal_row_equalizer_test.cpp @@ -0,0 +1,204 @@ +/* + * 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/mergetree/compact/internal_row_equalizer.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/binary_array.h" +#include "paimon/common/data/binary_array_writer.h" +#include "paimon/common/data/binary_map.h" +#include "paimon/common/data/generic_row.h" +#include "paimon/common/data/internal_array.h" +#include "paimon/common/data/internal_map.h" +#include "paimon/common/data/internal_row.h" +#include "paimon/common/utils/decimal_utils.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr CreateIntArray(const std::vector& values, + MemoryPool* pool) { + return std::make_shared(BinaryArray::FromIntArray(values, pool)); +} + +std::shared_ptr CreateNullableIntArray(const std::vector& values, + int32_t null_pos, MemoryPool* pool) { + auto array = std::make_shared(); + BinaryArrayWriter writer(array.get(), static_cast(values.size()), sizeof(int32_t), + pool); + for (int32_t i = 0; i < static_cast(values.size()); ++i) { + if (i == null_pos) { + writer.SetNullValue(i); + } else { + writer.WriteInt(i, values[i]); + } + } + writer.Complete(); + return array; +} + +std::shared_ptr CreateIntMap(const std::vector& keys, + const std::vector& values, MemoryPool* pool) { + BinaryArray key_array = BinaryArray::FromIntArray(keys, pool); + BinaryArray value_array = BinaryArray::FromIntArray(values, pool); + return BinaryMap::ValueOf(key_array, value_array, pool); +} + +std::shared_ptr CreateNestedRow(int32_t value, double floating_point) { + return GenericRow::Of({value, floating_point}); +} + +} // namespace + +TEST(InternalRowEqualizerTest, PrimitiveTypesAndIgnoreFields) { + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr schema = arrow::schema( + {arrow::field("boolean", arrow::boolean()), arrow::field("tinyint", arrow::int8()), + arrow::field("smallint", arrow::int16()), arrow::field("int", arrow::int32()), + arrow::field("date", arrow::date32()), arrow::field("bigint", arrow::int64()), + arrow::field("float", arrow::float32()), arrow::field("double", arrow::float64()), + arrow::field("string", arrow::utf8()), arrow::field("binary", arrow::binary()), + arrow::field("timestamp", arrow::timestamp(arrow::TimeUnit::MICRO)), + arrow::field("decimal", arrow::decimal128(10, 2)), + arrow::field("ignored", arrow::int32())}); + + const float float_nan = std::numeric_limits::quiet_NaN(); + const double double_nan = std::numeric_limits::quiet_NaN(); + auto binary = std::make_shared("binary", pool.get()); + Decimal decimal(10, 2, DecimalUtils::StrToInt128("12345").value()); + std::vector left_values = {true, + static_cast(1), + static_cast(2), + static_cast(3), + int32_t{4}, + int64_t{5}, + float_nan, + double_nan, + std::string_view("string"), + binary, + Timestamp(1234, 567000), + decimal, + int32_t{10}}; + std::vector right_values = left_values; + right_values.back() = int32_t{20}; + + std::unique_ptr left = GenericRow::Of(left_values); + std::unique_ptr right = GenericRow::Of(right_values); + ASSERT_OK_AND_ASSIGN(FieldsComparator::FieldComparatorFunc equalizer, + InternalRowEqualizer::Create(schema, {"ignored"})); + ASSERT_EQ(0, equalizer(*left, *right)); + + right->SetField(/*pos=*/3, int32_t{30}); + ASSERT_NE(0, equalizer(*left, *right)); +} + +TEST(InternalRowEqualizerTest, NullAndFloatingPointSemantics) { + std::shared_ptr schema = arrow::schema( + {arrow::field("value", arrow::float64()), arrow::field("nullable", arrow::int32())}); + ASSERT_OK_AND_ASSIGN(FieldsComparator::FieldComparatorFunc equalizer, + InternalRowEqualizer::Create(schema, {})); + + std::unique_ptr negative_zero = + GenericRow::Of({static_cast(-0.0), NullType()}); + std::unique_ptr positive_zero = + GenericRow::Of({static_cast(0.0), NullType()}); + ASSERT_NE(0, equalizer(*negative_zero, *positive_zero)); + + double nan1 = std::numeric_limits::quiet_NaN(); + double nan2 = -std::numeric_limits::quiet_NaN(); + std::unique_ptr left_nan = GenericRow::Of({nan1, NullType()}); + std::unique_ptr right_nan = GenericRow::Of({nan2, NullType()}); + ASSERT_EQ(0, equalizer(*left_nan, *right_nan)); + + right_nan->SetField(/*pos=*/1, int32_t{1}); + ASSERT_NE(0, equalizer(*left_nan, *right_nan)); +} + +TEST(InternalRowEqualizerTest, NestedTypes) { + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr schema = arrow::schema( + {arrow::field("array", arrow::list(arrow::int32())), + arrow::field("map", arrow::map(arrow::int32(), arrow::int32())), + arrow::field("row", arrow::struct_({arrow::field("value", arrow::int32()), + arrow::field("floating", arrow::float64())}))}); + ASSERT_OK_AND_ASSIGN(FieldsComparator::FieldComparatorFunc equalizer, + InternalRowEqualizer::Create(schema, {})); + + std::unique_ptr left = + GenericRow::Of({CreateNullableIntArray({1, 0, 3}, /*null_pos=*/1, pool.get()), + CreateIntMap({1, 2}, {10, 20}, pool.get()), + CreateNestedRow(100, std::numeric_limits::quiet_NaN())}); + std::unique_ptr right = + GenericRow::Of({CreateNullableIntArray({1, 9, 3}, /*null_pos=*/1, pool.get()), + CreateIntMap({1, 2}, {10, 20}, pool.get()), + CreateNestedRow(100, -std::numeric_limits::quiet_NaN())}); + ASSERT_EQ(0, equalizer(*left, *right)); + + right->SetField(/*pos=*/0, CreateIntArray({1, 2}, pool.get())); + ASSERT_NE(0, equalizer(*left, *right)); + + right->SetField(/*pos=*/0, CreateNullableIntArray({1, 0, 3}, /*null_pos=*/1, pool.get())); + right->SetField(/*pos=*/0, CreateNullableIntArray({1, 0, 4}, /*null_pos=*/1, pool.get())); + ASSERT_NE(0, equalizer(*left, *right)); + + right->SetField(/*pos=*/0, CreateNullableIntArray({1, 0, 3}, /*null_pos=*/1, pool.get())); + right->SetField(/*pos=*/1, CreateIntMap({1, 3}, {10, 20}, pool.get())); + ASSERT_NE(0, equalizer(*left, *right)); + + right->SetField(/*pos=*/1, CreateIntMap({1, 2}, {10, 20}, pool.get())); + right->SetField(/*pos=*/2, CreateNestedRow(101, 1.0)); + ASSERT_NE(0, equalizer(*left, *right)); +} + +TEST(InternalRowEqualizerTest, MapEqualityDoesNotDependOnEntryOrder) { + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr schema = + arrow::schema({arrow::field("map", arrow::map(arrow::int32(), arrow::int32()))}); + ASSERT_OK_AND_ASSIGN(FieldsComparator::FieldComparatorFunc equalizer, + InternalRowEqualizer::Create(schema, {})); + + std::unique_ptr left = + GenericRow::Of({CreateIntMap({1, 2, 3}, {10, 20, 30}, pool.get())}); + std::unique_ptr reordered = + GenericRow::Of({CreateIntMap({3, 1, 2}, {30, 10, 20}, pool.get())}); + ASSERT_EQ(0, equalizer(*left, *reordered)); + + reordered->SetField(/*pos=*/0, CreateIntMap({3, 1, 2}, {30, 10, 21}, pool.get())); + ASSERT_NE(0, equalizer(*left, *reordered)); +} + +TEST(InternalRowEqualizerTest, UnsupportedType) { + ASSERT_NOK_WITH_MSG(InternalRowEqualizer::Create( + arrow::schema({arrow::field("unsupported", arrow::null())}), {}), + "Do not support equality for type null"); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/interval_partition_test.cpp b/src/paimon/core/mergetree/compact/interval_partition_test.cpp index 861a35c4c..9c68000ca 100644 --- a/src/paimon/core/mergetree/compact/interval_partition_test.cpp +++ b/src/paimon/core/mergetree/compact/interval_partition_test.cpp @@ -125,7 +125,7 @@ class IntervalPartitionTest : public testing::Test { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } private: diff --git a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h index d2b21b2c0..34f7cf970 100644 --- a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h +++ b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h @@ -25,9 +25,11 @@ #include #include +#include "paimon/common/data/serializer/row_compacted_serializer.h" #include "paimon/common/utils/fields_comparator.h" #include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" #include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/changelog_result.h" #include "paimon/core/mergetree/compact/lookup_merge_function.h" #include "paimon/core/mergetree/compact/merge_function_wrapper.h" #include "paimon/core/mergetree/lookup/file_position.h" @@ -47,23 +49,28 @@ namespace paimon { /// should be AFTER. /// With level-0 record, without level-x record, need to lookup the history value of the upper /// level as BEFORE. -/// TODO(xinyu.lxy) : add changelog template -class LookupChangelogMergeFunctionWrapper : public MergeFunctionWrapper { +class LookupChangelogMergeFunctionWrapper : public MergeFunctionWrapper { public: static Result> Create( std::unique_ptr&& merge_function, std::function>(const std::shared_ptr&)> lookup, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& deletion_vectors_maintainer, - const std::shared_ptr& comparator) { + const std::shared_ptr& comparator, + std::unique_ptr&& value_serializer, + FieldsComparator::FieldComparatorFunc value_equalizer) { if (lookup_strategy.deletion_vector && !deletion_vectors_maintainer) { return Status::Invalid("deletionVectorsMaintainer should not be null, there is a bug."); } + if (should_produce_changelog && !value_serializer) { + return Status::Invalid("valueSerializer is required when producing changelog."); + } return std::unique_ptr( - new LookupChangelogMergeFunctionWrapper(std::move(merge_function), std::move(lookup), - lookup_strategy, deletion_vectors_maintainer, - comparator)); + new LookupChangelogMergeFunctionWrapper( + std::move(merge_function), std::move(lookup), lookup_strategy, + should_produce_changelog, deletion_vectors_maintainer, comparator, + std::move(value_serializer), std::move(value_equalizer))); } void Reset() override { merge_function_->Reset(); @@ -73,12 +80,17 @@ class LookupChangelogMergeFunctionWrapper : public MergeFunctionWrapperAdd(std::move(kv)); } - Result> GetResult() override { + Result> GetResult() override { // 1. Find the latest high level record and compute containLevel0 - std::optional high_level_idx = merge_function_->PickHighLevelIdx(); + const KeyValue* high_level = merge_function_->PickHighLevel(); + bool contain_level0 = merge_function_->ContainLevel0(); + std::optional before; + if (contain_level0 && should_produce_changelog_ && high_level != nullptr) { + PAIMON_ASSIGN_OR_RAISE(before, CloneKeyValue(*high_level, high_level->value_kind)); + } // 2. Lookup if latest high level record is absent - if (high_level_idx == std::nullopt) { + if (high_level == nullptr) { std::optional lookup_high_level; PAIMON_ASSIGN_OR_RAISE(std::optional lookup_result, lookup_(merge_function_->GetKey())); @@ -105,30 +117,83 @@ class LookupChangelogMergeFunctionWrapper : public MergeFunctionWrapperInsertInto(std::move(lookup_high_level), comparator_); } } // 3. Calculate result PAIMON_ASSIGN_OR_RAISE(std::optional result, merge_function_->GetResult()); - Reset(); + // 4. Set changelog when there's level-0 records - // TODO(liancheng.lsz): setChangelog - return result; + ChangelogResult changelog_result; + if (contain_level0 && should_produce_changelog_) { + PAIMON_RETURN_NOT_OK( + SetChangelog(std::move(before), result, &changelog_result.changelogs)); + } + changelog_result.result = std::move(result); + Reset(); + return std::optional(std::move(changelog_result)); } private: LookupChangelogMergeFunctionWrapper( std::unique_ptr&& merge_function, std::function>(const std::shared_ptr&)> lookup, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& deletion_vectors_maintainer, - const std::shared_ptr& user_defined_seq_comparator) + const std::shared_ptr& user_defined_seq_comparator, + std::unique_ptr&& value_serializer, + FieldsComparator::FieldComparatorFunc value_equalizer) : merge_function_(std::move(merge_function)), lookup_(std::move(lookup)), lookup_strategy_(lookup_strategy), + should_produce_changelog_(should_produce_changelog), deletion_vectors_maintainer_(deletion_vectors_maintainer), - comparator_(CreateSequenceComparator(user_defined_seq_comparator)) {} + comparator_(CreateSequenceComparator(user_defined_seq_comparator)), + value_serializer_(std::move(value_serializer)), + value_equalizer_(std::move(value_equalizer)) {} + + Result CloneKeyValue(const KeyValue& from, const RowKind* value_kind) { + // TODO(lisizhuo.lsz): avoid serialize & deserialize here. + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, + value_serializer_->SerializeToBytes(*from.value)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr value, + value_serializer_->Deserialize(bytes)); + return KeyValue(value_kind, from.sequence_number, KeyValue::UNKNOWN_LEVEL, from.key, + std::move(value)); + } + + Status SetChangelog(std::optional&& before, const std::optional& after, + std::vector* changelogs) { + if (!before || !before->value_kind->IsAdd()) { + if (after && after->value_kind->IsAdd()) { + PAIMON_ASSIGN_OR_RAISE(KeyValue insert, + CloneKeyValue(after.value(), RowKind::Insert())); + changelogs->emplace_back(std::move(insert)); + } + return Status::OK(); + } + + if (!after || !after->value_kind->IsAdd()) { + before->value_kind = RowKind::Delete(); + changelogs->emplace_back(std::move(before.value())); + return Status::OK(); + } + + if (!value_equalizer_ || value_equalizer_(*before->value, *after->value) != 0) { + before->value_kind = RowKind::UpdateBefore(); + PAIMON_ASSIGN_OR_RAISE(KeyValue update_after, + CloneKeyValue(after.value(), RowKind::UpdateAfter())); + changelogs->emplace_back(std::move(before.value())); + changelogs->emplace_back(std::move(update_after)); + } + return Status::OK(); + } static std::function CreateSequenceComparator( const std::shared_ptr& user_defined_seq_comparator) { @@ -150,8 +215,11 @@ class LookupChangelogMergeFunctionWrapper : public MergeFunctionWrapper merge_function_; std::function>(const std::shared_ptr&)> lookup_; LookupStrategy lookup_strategy_; + bool should_produce_changelog_; std::shared_ptr deletion_vectors_maintainer_; std::function comparator_; + std::unique_ptr value_serializer_; + FieldsComparator::FieldComparatorFunc value_equalizer_; }; } // namespace paimon diff --git a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp index 446beb091..c4d255d9a 100644 --- a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp @@ -26,12 +26,22 @@ #include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" #include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/internal_row_equalizer.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/key_value_checker.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { +std::unique_ptr CreateValueSerializer( + const std::shared_ptr& pool) { + return RowCompactedSerializer::Create(arrow::schema({arrow::field("value", arrow::int32())}), + pool) + .value(); +} +} // namespace + TEST(LookupChangelogMergeFunctionWrapperTest, TestCreateInvalid) { auto pool = GetDefaultPool(); auto mfunc = std::make_unique(/*ignore_delete=*/true); @@ -45,8 +55,10 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestCreateInvalid) { /*deletion_vector=*/true, /*force_lookup=*/true); ASSERT_NOK_WITH_MSG(LookupChangelogMergeFunctionWrapper::Create( std::move(lookup_mfunc), lookup, lookup_strategy, + /*should_produce_changelog=*/false, /*deletion_vectors_maintainer=*/nullptr, - /*comparator=*/nullptr), + /*comparator=*/nullptr, CreateValueSerializer(pool), + /*value_equalizer=*/{}), "deletionVectorsMaintainer should not be null, there is a bug."); } @@ -69,12 +81,15 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestSimple) { KeyValue(RowKind::Insert(), /*sequence_number=*/1000, /*level=*/3, key, BinaryRowGenerator::GenerateRowPtr({1001}, pool.get()))); }; - LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/false, + LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/true, /*deletion_vector=*/false, /*force_lookup=*/true); ASSERT_OK_AND_ASSIGN(auto wrapper, LookupChangelogMergeFunctionWrapper::Create( std::move(lookup_mfunc), lookup, lookup_strategy, + /*should_produce_changelog=*/false, /*deletion_vectors_maintainer=*/nullptr, - /*comparator=*/nullptr)); + /*comparator=*/nullptr, + /*value_serializer=*/nullptr, + /*value_equalizer=*/{})); wrapper->Reset(); ASSERT_OK(wrapper->Add(std::move(kv1))); @@ -82,7 +97,9 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestSimple) { ASSERT_OK(wrapper->Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper->GetResult()); ASSERT_TRUE(result); - ASSERT_EQ(result.value().sequence_number, 2); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 2); + ASSERT_TRUE(result->changelogs.empty()); } TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookup) { @@ -104,12 +121,14 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookup) { KeyValue(RowKind::Insert(), /*sequence_number=*/0, /*level=*/3, key, BinaryRowGenerator::GenerateRowPtr({1001}, pool.get()))); }; - LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/false, + LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/true, /*deletion_vector=*/false, /*force_lookup=*/true); ASSERT_OK_AND_ASSIGN(auto wrapper, LookupChangelogMergeFunctionWrapper::Create( std::move(lookup_mfunc), lookup, lookup_strategy, + /*should_produce_changelog=*/true, /*deletion_vectors_maintainer=*/nullptr, - /*comparator=*/nullptr)); + /*comparator=*/nullptr, CreateValueSerializer(pool), + /*value_equalizer=*/{})); wrapper->Reset(); ASSERT_OK(wrapper->Add(std::move(kv1))); @@ -117,7 +136,98 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookup) { ASSERT_OK(wrapper->Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper->GetResult()); ASSERT_TRUE(result); - ASSERT_EQ(result.value().sequence_number, 3); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 3); + ASSERT_EQ(result->changelogs.size(), 2); + ASSERT_EQ(result->changelogs[0].value_kind, RowKind::UpdateBefore()); + ASSERT_EQ(result->changelogs[0].level, KeyValue::UNKNOWN_LEVEL); + ASSERT_EQ(result->changelogs[0].value->GetInt(0), 1001); + ASSERT_EQ(result->changelogs[1].value_kind, RowKind::UpdateAfter()); + ASSERT_EQ(result->changelogs[1].level, KeyValue::UNKNOWN_LEVEL); + ASSERT_EQ(result->changelogs[1].value->GetInt(0), 300); +} + +TEST(LookupChangelogMergeFunctionWrapperTest, TestRowDeduplicate) { + auto pool = GetDefaultPool(); + KeyValue kv(RowKind::Insert(), /*sequence_number=*/1, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), + BinaryRowGenerator::GenerateRowPtr({300}, pool.get())); + + auto merge_function = std::make_unique(/*ignore_delete=*/true); + auto lookup_merge_function = std::make_unique(std::move(merge_function)); + auto lookup = [&](const std::shared_ptr& key) -> Result> { + return std::optional( + KeyValue(RowKind::Insert(), /*sequence_number=*/0, /*level=*/3, key, + BinaryRowGenerator::GenerateRowPtr({300}, pool.get()))); + }; + LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/true, + /*deletion_vector=*/false, /*force_lookup=*/true); + auto value_equalizer = [](const InternalRow& lhs, const InternalRow& rhs) { + return lhs.GetInt(0) == rhs.GetInt(0) ? 0 : 1; + }; + ASSERT_OK_AND_ASSIGN(auto wrapper, + LookupChangelogMergeFunctionWrapper::Create( + std::move(lookup_merge_function), lookup, lookup_strategy, + /*should_produce_changelog=*/true, + /*deletion_vectors_maintainer=*/nullptr, + /*comparator=*/nullptr, CreateValueSerializer(pool), value_equalizer)); + + wrapper->Reset(); + ASSERT_OK(wrapper->Add(std::move(kv))); + ASSERT_OK_AND_ASSIGN(auto result, wrapper->GetResult()); + ASSERT_TRUE(result); + ASSERT_TRUE(result->result); + ASSERT_TRUE(result->changelogs.empty()); +} + +TEST(LookupChangelogMergeFunctionWrapperTest, TestRowDeduplicateWithIgnoreFields) { + auto pool = GetDefaultPool(); + auto value_schema = arrow::schema( + {arrow::field("value", arrow::int32()), arrow::field("ignored", arrow::int32())}); + auto merge_function = std::make_unique(/*ignore_delete=*/true); + auto lookup_merge_function = std::make_unique(std::move(merge_function)); + auto lookup = [&](const std::shared_ptr& key) -> Result> { + return std::optional( + KeyValue(RowKind::Insert(), /*sequence_number=*/0, /*level=*/3, key, + BinaryRowGenerator::GenerateRowPtr({300, 1}, pool.get()))); + }; + LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/true, + /*deletion_vector=*/false, /*force_lookup=*/true); + ASSERT_OK_AND_ASSIGN(auto value_equalizer, + InternalRowEqualizer::Create(value_schema, {"ignored"})); + ASSERT_OK_AND_ASSIGN(auto value_serializer, RowCompactedSerializer::Create(value_schema, pool)); + ASSERT_OK_AND_ASSIGN(auto wrapper, + LookupChangelogMergeFunctionWrapper::Create( + std::move(lookup_merge_function), lookup, lookup_strategy, + /*should_produce_changelog=*/true, + /*deletion_vectors_maintainer=*/nullptr, + /*comparator=*/nullptr, std::move(value_serializer), value_equalizer)); + + // A change to an ignored field does not produce an update changelog. + wrapper->Reset(); + ASSERT_OK(wrapper->Add(KeyValue(RowKind::Insert(), /*sequence_number=*/1, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), + BinaryRowGenerator::GenerateRowPtr({300, 2}, pool.get())))); + ASSERT_OK_AND_ASSIGN(auto ignored_field_result, wrapper->GetResult()); + ASSERT_TRUE(ignored_field_result); + ASSERT_TRUE(ignored_field_result->result); + ASSERT_TRUE(ignored_field_result->changelogs.empty()); + + // A change to a non-ignored field still produces update-before and update-after. + wrapper->Reset(); + ASSERT_OK(wrapper->Add(KeyValue(RowKind::Insert(), /*sequence_number=*/2, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({20}, pool.get()), + BinaryRowGenerator::GenerateRowPtr({301, 2}, pool.get())))); + ASSERT_OK_AND_ASSIGN(auto value_field_result, wrapper->GetResult()); + ASSERT_TRUE(value_field_result); + ASSERT_TRUE(value_field_result->result); + ASSERT_EQ(value_field_result->changelogs.size(), 2); + ASSERT_EQ(value_field_result->changelogs[0].value_kind, RowKind::UpdateBefore()); + ASSERT_EQ(value_field_result->changelogs[0].value->GetInt(0), 300); + ASSERT_EQ(value_field_result->changelogs[0].value->GetInt(1), 1); + ASSERT_EQ(value_field_result->changelogs[1].value_kind, RowKind::UpdateAfter()); + ASSERT_EQ(value_field_result->changelogs[1].value->GetInt(0), 301); + ASSERT_EQ(value_field_result->changelogs[1].value->GetInt(1), 2); } TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookupWithDv) { @@ -154,8 +264,66 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookupWithDv) { /*deletion_vector=*/true, /*force_lookup=*/false); ASSERT_OK_AND_ASSIGN(auto wrapper, LookupChangelogMergeFunctionWrapper::Create( - std::move(lookup_mfunc), lookup, lookup_strategy, dv_maintainer, - /*comparator=*/nullptr)); + std::move(lookup_mfunc), lookup, lookup_strategy, + /*should_produce_changelog=*/false, dv_maintainer, + /*comparator=*/nullptr, + /*value_serializer=*/nullptr, + /*value_equalizer=*/{})); + + wrapper->Reset(); + ASSERT_OK(wrapper->Add(std::move(kv1))); + ASSERT_OK(wrapper->Add(std::move(kv2))); + ASSERT_OK(wrapper->Add(std::move(kv3))); + ASSERT_OK_AND_ASSIGN(auto result, wrapper->GetResult()); + ASSERT_TRUE(result); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 3); + ASSERT_EQ(result->result->value->GetInt(0), 100 + 200 + 300 + 1001); + + auto dv = dv_maintainer->DeletionVectorOf("data.file"); + ASSERT_TRUE(dv); + ASSERT_FALSE(dv.value()->IsDeleted(0).value()); + ASSERT_TRUE(dv.value()->IsDeleted(10).value()); +} + +TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookupWithDvAndChangelog) { + auto pool = GetDefaultPool(); + KeyValue kv1(RowKind::Insert(), /*sequence_number=*/1, /*level=*/0, /*key=*/ + BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), + /*value=*/BinaryRowGenerator::GenerateRowPtr({100}, pool.get())); + KeyValue kv2(RowKind::Insert(), /*sequence_number=*/2, /*level=*/0, + /*key=*/BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), + /*value=*/BinaryRowGenerator::GenerateRowPtr({200}, pool.get())); + KeyValue kv3(RowKind::Insert(), /*sequence_number=*/3, /*level=*/0, + /*key=*/BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), + /*value=*/BinaryRowGenerator::GenerateRowPtr({300}, pool.get())); + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, + CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); + ASSERT_OK_AND_ASSIGN(auto mfunc, AggregateMergeFunction::Create( + arrow::schema({arrow::field("value", arrow::int32())}), + {"key"}, core_options, pool)); + auto lookup_mfunc = std::make_unique(std::move(mfunc)); + auto lookup = + [&](const std::shared_ptr& key) -> Result> { + return std::optional( + {KeyValue(RowKind::Insert(), /*sequence_number=*/0, /*level=*/3, key, + BinaryRowGenerator::GenerateRowPtr({1001}, pool.get())), + "data.file", /*row_position=*/10}); + }; + auto dv_index_file = + std::make_shared(/*fs=*/nullptr, /*path_factory=*/nullptr, + /*bitmap64=*/false, pool); + std::map> deletion_vectors; + auto dv_maintainer = std::make_shared(dv_index_file, deletion_vectors); + + LookupStrategy lookup_strategy(/*is_first_row=*/false, /*produce_changelog=*/true, + /*deletion_vector=*/true, /*force_lookup=*/false); + ASSERT_OK_AND_ASSIGN(auto wrapper, + LookupChangelogMergeFunctionWrapper::Create( + std::move(lookup_mfunc), lookup, lookup_strategy, + /*should_produce_changelog=*/true, dv_maintainer, + /*comparator=*/nullptr, CreateValueSerializer(pool), + /*value_equalizer=*/{})); wrapper->Reset(); ASSERT_OK(wrapper->Add(std::move(kv1))); @@ -163,8 +331,16 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookupWithDv) { ASSERT_OK(wrapper->Add(std::move(kv3))); ASSERT_OK_AND_ASSIGN(auto result, wrapper->GetResult()); ASSERT_TRUE(result); - ASSERT_EQ(result.value().sequence_number, 3); - ASSERT_EQ(result.value().value->GetInt(0), 100 + 200 + 300 + 1001); + ASSERT_TRUE(result->result); + ASSERT_EQ(result->result->sequence_number, 3); + ASSERT_EQ(result->result->value->GetInt(0), 100 + 200 + 300 + 1001); + ASSERT_EQ(result->changelogs.size(), 2); + ASSERT_EQ(result->changelogs[0].value_kind, RowKind::UpdateBefore()); + ASSERT_EQ(result->changelogs[0].sequence_number, 0); + ASSERT_EQ(result->changelogs[0].value->GetInt(0), 1001); + ASSERT_EQ(result->changelogs[1].value_kind, RowKind::UpdateAfter()); + ASSERT_EQ(result->changelogs[1].sequence_number, 3); + ASSERT_EQ(result->changelogs[1].value->GetInt(0), 100 + 200 + 300 + 1001); auto dv = dv_maintainer->DeletionVectorOf("data.file"); ASSERT_TRUE(dv); diff --git a/src/paimon/core/mergetree/compact/lookup_merge_function.h b/src/paimon/core/mergetree/compact/lookup_merge_function.h index e50d85eb2..7d8ee1a02 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_function.h +++ b/src/paimon/core/mergetree/compact/lookup_merge_function.h @@ -100,6 +100,11 @@ class LookupMergeFunction : public MergeFunction { return high_level_idx; } + const KeyValue* PickHighLevel() const { + std::optional high_level_idx = PickHighLevelIdx(); + return high_level_idx ? &candidates_[high_level_idx.value()] : nullptr; + } + private: std::unique_ptr merge_function_; std::vector candidates_; diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp index 994b0e3fc..1a818dd01 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.cpp @@ -23,8 +23,10 @@ #include "paimon/common/table/special_fields.h" #include "paimon/core/mergetree/compact/first_row_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/mergetree/lookup/file_position.h" #include "paimon/core/mergetree/lookup/positioned_key_value.h" +#include "paimon/core/utils/primary_key_table_utils.h" namespace paimon { template @@ -38,7 +40,8 @@ LookupMergeTreeCompactRewriter::LookupMergeTreeCompactRewriter( const std::shared_ptr& path_factory_cache, std::unique_ptr&& merge_file_split_read, MergeFunctionWrapperFactory merge_function_wrapper_factory, - const std::shared_ptr& cancellation_controller, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, + bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager, const std::shared_ptr& pool) : ChangelogMergeTreeRewriter( @@ -46,6 +49,7 @@ LookupMergeTreeCompactRewriter::LookupMergeTreeCompactRewriter( trimmed_primary_keys, options, data_schema, write_schema, DeletionVector::CreateFactory(dv_maintainer), path_factory_cache, std::move(merge_file_split_read), std::move(merge_function_wrapper_factory), + std::move(changelog_merge_function_wrapper_factory), produce_changelog, cancellation_controller, pool), lookup_levels_(std::move(lookup_levels)), dv_maintainer_(dv_maintainer), @@ -56,10 +60,10 @@ Result>> LookupMergeTreeCompactRewriter::Create( int32_t max_level, std::unique_ptr>&& lookup_levels, const std::shared_ptr& dv_maintainer, - MergeFunctionWrapperFactory merge_function_wrapper_factory, int32_t bucket, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, int32_t bucket, const BinaryRow& partition, const std::shared_ptr& table_schema, const std::shared_ptr& path_factory_cache, - const CoreOptions& options, + const CoreOptions& options, bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager, const std::shared_ptr& pool) { @@ -78,14 +82,9 @@ LookupMergeTreeCompactRewriter::Create( .WithMemoryPool(pool); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, read_context_builder.Finish()); - // TODO(xinyu.lxy): temporarily disabled pre-buffer for parquet, which may cause high memory - // usage during compaction. Will fix via parquet format refactor. - auto new_options = options.ToMap(); - if (new_options.find("parquet.read.enable-pre-buffer") == new_options.end()) { - new_options["parquet.read.enable-pre-buffer"] = "false"; - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, new_options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr internal_context, + InternalReadContext::Create(read_context, table_schema, options.ToMap())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); @@ -93,44 +92,60 @@ LookupMergeTreeCompactRewriter::Create( std::unique_ptr merge_file_split_read, MergeFileSplitRead::Create(path_factory, internal_context, pool, CreateDefaultExecutor())); + MergeFunctionWrapperFactory merge_function_wrapper_factory = + [data_schema, options, trimmed_primary_keys, pool]( + int32_t /*output_level*/) -> Result>> { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + data_schema, trimmed_primary_keys, options, pool)); + if (options.NeedLookup() && options.GetMergeEngine() != MergeEngine::FIRST_ROW) { + merge_function = std::make_unique(std::move(merge_function)); + } + return std::make_shared(std::move(merge_function)); + }; + return std::unique_ptr(new LookupMergeTreeCompactRewriter( std::move(lookup_levels), dv_maintainer, max_level, partition, bucket, table_schema->Id(), trimmed_primary_keys, options, data_schema, write_schema, path_factory_cache, std::move(merge_file_split_read), std::move(merge_function_wrapper_factory), + std::move(changelog_merge_function_wrapper_factory), produce_changelog, cancellation_controller, remote_lookup_file_manager, pool)); } template -std::shared_ptr> +std::shared_ptr> LookupMergeTreeCompactRewriter::CreateFirstRowMergeFunctionWrapper( std::unique_ptr&& merge_func, int32_t output_level, - LookupLevels* lookup_levels) { + std::unique_ptr&& value_serializer, LookupLevels* lookup_levels) { auto contains = [output_level, lookup_levels](const std::shared_ptr& key) -> Result { PAIMON_ASSIGN_OR_RAISE(std::optional contain, lookup_levels->Lookup(key, output_level + 1)); return contain != std::nullopt; }; - return std::make_shared(std::move(merge_func), - std::move(contains)); + return std::make_shared( + std::move(merge_func), std::move(contains), std::move(value_serializer)); } template -Result>> +Result>> LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( std::unique_ptr&& merge_func, int32_t output_level, const std::shared_ptr& deletion_vectors_maintainer, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& user_defined_seq_comparator, - LookupLevels* lookup_levels) { + std::unique_ptr&& value_serializer, + FieldsComparator::FieldComparatorFunc value_equalizer, LookupLevels* lookup_levels) { auto lookup = [output_level, lookup_levels]( const std::shared_ptr& key) -> Result> { return lookup_levels->Lookup(key, output_level + 1); }; - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr> wrapper, - LookupChangelogMergeFunctionWrapper::Create( - std::move(merge_func), std::move(lookup), lookup_strategy, - deletion_vectors_maintainer, user_defined_seq_comparator)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> wrapper, + LookupChangelogMergeFunctionWrapper::Create( + std::move(merge_func), std::move(lookup), lookup_strategy, should_produce_changelog, + deletion_vectors_maintainer, user_defined_seq_comparator, std::move(value_serializer), + std::move(value_equalizer))); return wrapper; } diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.h b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.h index 06c5f47f4..a461177ba 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.h +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.h @@ -19,10 +19,12 @@ #pragma once #include "arrow/api.h" +#include "paimon/common/data/serializer/row_compacted_serializer.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h" #include "paimon/core/mergetree/compact/first_row_merge_function.h" +#include "paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper.h" #include "paimon/core/mergetree/compact/lookup_merge_function.h" #include "paimon/core/mergetree/lookup/remote_lookup_file_manager.h" #include "paimon/core/mergetree/lookup_levels.h" @@ -38,10 +40,11 @@ class LookupMergeTreeCompactRewriter : public ChangelogMergeTreeRewriter { static Result> Create( int32_t max_level, std::unique_ptr>&& lookup_levels, const std::shared_ptr& dv_maintainer, - MergeFunctionWrapperFactory merge_function_wrapper_factory, int32_t bucket, - const BinaryRow& partition, const std::shared_ptr& table_schema, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, + int32_t bucket, const BinaryRow& partition, + const std::shared_ptr& table_schema, const std::shared_ptr& path_factory_cache, - const CoreOptions& options, + const CoreOptions& options, bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager, const std::shared_ptr& pool); @@ -50,16 +53,20 @@ class LookupMergeTreeCompactRewriter : public ChangelogMergeTreeRewriter { return lookup_levels_->Close(); } - static std::shared_ptr> CreateFirstRowMergeFunctionWrapper( - std::unique_ptr&& merge_func, int32_t output_level, - LookupLevels* lookup_levels); + static std::shared_ptr> + CreateFirstRowMergeFunctionWrapper(std::unique_ptr&& merge_func, + int32_t output_level, + std::unique_ptr&& value_serializer, + LookupLevels* lookup_levels); - static Result>> CreateLookupMergeFunctionWrapper( + static Result>> + CreateLookupMergeFunctionWrapper( std::unique_ptr&& merge_func, int32_t output_level, const std::shared_ptr& deletion_vectors_maintainer, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& user_defined_seq_comparator, - LookupLevels* lookup_levels); + std::unique_ptr&& value_serializer, + FieldsComparator::FieldComparatorFunc value_equalizer, LookupLevels* lookup_levels); private: LookupMergeTreeCompactRewriter( @@ -72,6 +79,8 @@ class LookupMergeTreeCompactRewriter : public ChangelogMergeTreeRewriter { const std::shared_ptr& path_factory_cache, std::unique_ptr&& merge_file_split_read, MergeFunctionWrapperFactory merge_function_wrapper_factory, + ChangelogMergeFunctionWrapperFactory changelog_merge_function_wrapper_factory, + bool produce_changelog, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager, const std::shared_ptr& pool); diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp index bbc1608c4..c930c5ca7 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp @@ -24,6 +24,7 @@ #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/catalog/catalog.h" +#include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/factories/io_hook.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/fields_comparator.h" @@ -35,6 +36,7 @@ #include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" #include "paimon/core/mergetree/compact/changelog_merge_tree_rewriter.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/internal_row_equalizer.h" #include "paimon/core/mergetree/compact/interval_partition.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/mergetree/lookup/default_lookup_serializer_factory.h" @@ -148,16 +150,25 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam>> CreateCompactRewriterForFirstRow( const std::string& table_path, const std::shared_ptr& table_schema, - const CoreOptions& options, std::unique_ptr>&& lookup_levels) const { + const CoreOptions& options, std::unique_ptr>&& lookup_levels, + std::optional produce_changelog = std::nullopt) const { auto path_factory_cache = std::make_shared(table_path, table_schema, options, pool_); + bool should_produce_changelog = + produce_changelog.value_or(options.GetLookupStrategy().produce_changelog); auto merge_function_wrapper_factory = - [lookup_levels_ptr = lookup_levels.get()]( - int32_t output_level) -> Result>> { - std::shared_ptr> merge_function_wrapper = + [lookup_levels_ptr = lookup_levels.get(), data_schema = arrow_schema_, + should_produce_changelog, pool = pool_](int32_t output_level) + -> Result>> { + std::unique_ptr value_serializer; + if (should_produce_changelog) { + PAIMON_ASSIGN_OR_RAISE(value_serializer, + RowCompactedSerializer::Create(data_schema, pool)); + } + std::shared_ptr> merge_function_wrapper = LookupMergeTreeCompactRewriter::CreateFirstRowMergeFunctionWrapper( std::make_unique(/*ignore_delete=*/true), output_level, - lookup_levels_ptr); + std::move(value_serializer), lookup_levels_ptr); return merge_function_wrapper; }; auto cancellation_controller = std::make_shared(); @@ -166,7 +177,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam>> @@ -177,18 +189,29 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam(table_path, table_schema, options, pool_); auto merge_function_wrapper_factory = [this, table_schema, options, lookup_levels_ptr = lookup_levels.get(), - lookup_strategy = options.GetLookupStrategy()]( - int32_t output_level) -> Result>> { - PAIMON_ASSIGN_OR_RAISE(auto merge_func, - AggregateMergeFunction::Create( - arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), - options, GetDefaultPool())); + lookup_strategy = options.GetLookupStrategy()](int32_t output_level) + -> Result>> { PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, + auto merge_func, + AggregateMergeFunction::Create( + arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), options, pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr value_serializer, + RowCompactedSerializer::Create(arrow_schema_, pool_)); + FieldsComparator::FieldComparatorFunc value_equalizer; + if (options.ChangelogRowDeduplicate()) { + PAIMON_ASSIGN_OR_RAISE( + value_equalizer, + InternalRowEqualizer::Create(arrow_schema_, + options.GetChangelogRowDeduplicateIgnoreFields())); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( std::make_unique(std::move(merge_func)), output_level, /*deletion_vectors_maintainer=*/nullptr, lookup_strategy, - /*user_defined_seq_comparator=*/nullptr, lookup_levels_ptr)); + /*should_produce_changelog=*/lookup_strategy.produce_changelog, + /*user_defined_seq_comparator=*/nullptr, std::move(value_serializer), + std::move(value_equalizer), lookup_levels_ptr)); return merge_function_wrapper; }; auto cancellation_controller = std::make_shared(); @@ -196,7 +219,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam>> @@ -213,17 +237,22 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam(dv_index_file, deletion_vectors); - auto merge_function_wrapper_factory = - [lookup_levels_ptr = lookup_levels.get(), lookup_strategy = options.GetLookupStrategy(), - dv_maintainer_ptr = dv_maintainer]( - int32_t output_level) -> Result>> { + auto merge_function_wrapper_factory = [this, lookup_levels_ptr = lookup_levels.get(), + lookup_strategy = options.GetLookupStrategy(), + dv_maintainer_ptr = + dv_maintainer](int32_t output_level) + -> Result>> { auto merge_func = std::make_unique(false); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr value_serializer, + RowCompactedSerializer::Create(arrow_schema_, pool_)); PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, + std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( std::make_unique(std::move(merge_func)), output_level, dv_maintainer_ptr, lookup_strategy, - /*user_defined_seq_comparator=*/nullptr, lookup_levels_ptr)); + /*should_produce_changelog=*/lookup_strategy.produce_changelog, + /*user_defined_seq_comparator=*/nullptr, std::move(value_serializer), + /*value_equalizer=*/{}, lookup_levels_ptr)); return merge_function_wrapper; }; auto cancellation_controller = std::make_shared(); @@ -231,7 +260,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam>> @@ -253,19 +283,31 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam Result>> { - PAIMON_ASSIGN_OR_RAISE(auto merge_func, - AggregateMergeFunction::Create( - arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), - options, GetDefaultPool())); + lookup_strategy = options.GetLookupStrategy(), + dv_maintainer_ptr = dv_maintainer](int32_t output_level) + -> Result>> { + PAIMON_ASSIGN_OR_RAISE( + auto merge_func, + AggregateMergeFunction::Create( + arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), options, pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr value_serializer, + RowCompactedSerializer::Create(arrow_schema_, pool_)); + FieldsComparator::FieldComparatorFunc value_equalizer; + if (options.ChangelogRowDeduplicate()) { + PAIMON_ASSIGN_OR_RAISE( + value_equalizer, + InternalRowEqualizer::Create(arrow_schema_, + options.GetChangelogRowDeduplicateIgnoreFields())); + } PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, + std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter:: CreateLookupMergeFunctionWrapper( std::make_unique(std::move(merge_func)), output_level, dv_maintainer_ptr, lookup_strategy, - /*user_defined_seq_comparator=*/nullptr, lookup_levels_ptr)); + /*should_produce_changelog=*/lookup_strategy.produce_changelog, + /*user_defined_seq_comparator=*/nullptr, std::move(value_serializer), + std::move(value_equalizer), lookup_levels_ptr)); return merge_function_wrapper; }; auto cancellation_controller = std::make_shared(); @@ -273,7 +315,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParamEquals(*result_array)) << result_array->ToString(); } + void CheckShreddingFileSchema(const std::string& file_name, + const std::shared_ptr& table_schema, + const std::string& file_format_name, + const std::shared_ptr& expected_physical_schema, + int32_t field_index, + const MapSharedShreddingFieldMeta& expected_meta) const { + ASSERT_OK_AND_ASSIGN(auto file_format, + FileFormatFactory::Get(file_format_name, table_schema->Options())); + ASSERT_OK_AND_ASSIGN(auto reader_builder, + file_format->CreateReaderBuilder(/*batch_size=*/10)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_name)); + ASSERT_OK_AND_ASSIGN(auto file_batch_reader, reader_builder->Build(input_stream)); + ASSERT_OK_AND_ASSIGN(auto c_file_schema, file_batch_reader->GetFileSchema()); + std::shared_ptr file_schema = + arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); + + ASSERT_TRUE(file_schema->Equals(*expected_physical_schema, /*check_metadata=*/false)) + << "Expected schema:\n" + << expected_physical_schema->ToString() << "\nActual schema:\n" + << file_schema->ToString(); + std::shared_ptr metadata = + file_schema->field(field_index)->metadata(); + ASSERT_NE(nullptr, metadata); + ASSERT_OK_AND_ASSIGN(MapSharedShreddingFieldMeta actual_meta, + MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy())); + ASSERT_EQ(expected_meta, actual_meta); + } + Result> CreateFileStorePathFactory( const std::string& table_path, const CoreOptions& options) const { PAIMON_ASSIGN_OR_RAISE(std::vector external_paths, @@ -422,7 +493,8 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam options = {{Options::MERGE_ENGINE, "first-row"}, - {Options::FILE_FORMAT, "orc"}}; + {Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "lookup"}}; ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options)); auto schema_manager = std::make_shared(fs_, table_path); @@ -447,6 +519,7 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowRewrite) { /*output_level=*/5, /*drop_delete=*/true, runs)); ASSERT_EQ(2, compact_result.Before().size()); ASSERT_EQ(1, compact_result.After().size()); + ASSERT_EQ(1, compact_result.Changelog().size()); // check compact result const auto& compact_file_meta = compact_result.After()[0]; @@ -460,7 +533,8 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowRewrite) { BinaryRowGenerator::GenerateStats({1, 5}, {5, 33}, {0, 0}, pool_.get()), /*min_sequence_number=*/0l, /*max_sequence_number=*/3l, /*schema_id=*/0, /*level=*/5, std::vector>(), Timestamp(0l, 0), /*delete_row_count=*/0, - nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_TRUE(expected_file_meta->TEST_Equal(*compact_file_meta)); // check compact file exist @@ -482,6 +556,148 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowRewrite) { &expected_array); ASSERT_TRUE(array_status.ok()); CheckResult(compact_file_name, table_schema, "orc", expected_array); + std::string changelog_file_name = + table_path + "/bucket-0/" + compact_result.Changelog()[0]->file_name; + CheckResult(changelog_file_name, table_schema, "orc", expected_array); +} + +TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowLooksUpExistingKeys) { + std::map options = {{Options::MERGE_ENGINE, "first-row"}, + {Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "lookup"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); + ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options)); + auto schema_manager = std::make_shared(fs_, table_path); + ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager->ReadSchema(0)); + + ASSERT_OK_AND_ASSIGN(auto level0_file, + NewFiles(/*level=*/0, /*last_sequence_number=*/0, table_path, core_options, + "[[1, 111], [2, 22]]")); + ASSERT_OK_AND_ASSIGN(auto high_level_file, NewFiles(/*level=*/2, /*last_sequence_number=*/-1, + table_path, core_options, "[[1, 11]]")); + auto processor_factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto lookup_levels, + CreateLookupLevels(table_path, table_schema, processor_factory, + std::vector>{ + level0_file, high_level_file})); + ASSERT_OK_AND_ASSIGN(auto rewriter, + CreateCompactRewriterForFirstRow(table_path, table_schema, core_options, + std::move(lookup_levels))); + ASSERT_OK_AND_ASSIGN( + auto runs, GenerateSortedRuns(std::vector>{level0_file})); + ASSERT_OK_AND_ASSIGN(auto compact_result, rewriter->Rewrite( + /*output_level=*/1, /*drop_delete=*/true, runs)); + + ASSERT_EQ(1, compact_result.After().size()); + ASSERT_EQ(1, compact_result.After()[0]->row_count); + ASSERT_EQ(1, compact_result.Changelog().size()); + ASSERT_EQ(1, compact_result.Changelog()[0]->row_count); + + auto type_with_special_fields = + arrow::struct_(SpecialFields::CompleteSequenceAndValueKindField(arrow_schema_)->fields()); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(type_with_special_fields, + {"[[2, 0, 2, 22]]"}, &expected) + .ok()); + CheckResult(table_path + "/bucket-0/" + compact_result.After()[0]->file_name, table_schema, + "orc", expected); + CheckResult(table_path + "/bucket-0/" + compact_result.Changelog()[0]->file_name, table_schema, + "orc", expected); +} + +TEST_F(LookupMergeTreeCompactRewriterTest, TestLookupChangelogCanBeDisabled) { + std::map options = {{Options::MERGE_ENGINE, "first-row"}, + {Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "lookup"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); + ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options)); + auto schema_manager = std::make_shared(fs_, table_path); + ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager->ReadSchema(0)); + ASSERT_OK_AND_ASSIGN(auto file, NewFiles(/*level=*/0, /*last_sequence_number=*/-1, table_path, + core_options, "[[1, 11], [2, 22]]")); + auto processor_factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto lookup_levels, CreateLookupLevels(table_path, table_schema, + processor_factory, {file})); + ASSERT_OK_AND_ASSIGN(auto rewriter, + CreateCompactRewriterForFirstRow(table_path, table_schema, core_options, + std::move(lookup_levels), + /*produce_changelog=*/false)); + ASSERT_OK_AND_ASSIGN(auto runs, GenerateSortedRuns({file})); + ASSERT_OK_AND_ASSIGN(auto compact_result, rewriter->Rewrite( + /*output_level=*/1, /*drop_delete=*/true, runs)); + + ASSERT_EQ(1, compact_result.After().size()); + ASSERT_TRUE(compact_result.Changelog().empty()); +} + +TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowRewriteWithSharedShreddingChangelog) { + arrow::FieldVector fields = { + arrow::field("key", arrow::int32()), + arrow::field("tags", arrow::map(arrow::utf8(), arrow::int64())), + }; + arrow_schema_ = arrow::schema(fields); + key_schema_ = arrow::schema({fields[0]}); + std::map options = { + {Options::MERGE_ENGINE, "first-row"}, + {Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::CHANGELOG_FILE_FORMAT, "parquet"}, + {"fields.tags.map.storage-layout", "shared-shredding"}, + {"fields.tags.map.shared-shredding.max-columns", "3"}, + {"fields.tags.map.shared-shredding.column-placement-policy", "plain"}, + }; + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); + ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options)); + auto schema_manager = std::make_shared(fs_, table_path); + ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager->ReadSchema(0)); + + ASSERT_OK_AND_ASSIGN( + auto file0, NewFiles(/*level=*/0, /*last_sequence_number=*/-1, table_path, core_options, + R"([[1, [["a", 11]]], [3, [["c", 33]]], [5, [["e", 55]]]])")); + ASSERT_OK_AND_ASSIGN(auto file1, + NewFiles(/*level=*/0, /*last_sequence_number=*/2, table_path, core_options, + R"([[2, [["b", 22]]], [5, [["f", 555]]]])")); + std::vector> files = {file0, file1}; + auto processor_factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto lookup_levels, CreateLookupLevels(table_path, table_schema, + processor_factory, files)); + ASSERT_OK_AND_ASSIGN(auto rewriter, + CreateCompactRewriterForFirstRow(table_path, table_schema, core_options, + std::move(lookup_levels))); + ASSERT_OK_AND_ASSIGN(auto runs, GenerateSortedRuns(files)); + ASSERT_OK_AND_ASSIGN(auto compact_result, rewriter->Rewrite( + /*output_level=*/5, /*drop_delete=*/true, runs)); + + ASSERT_EQ(1, compact_result.Changelog().size()); + const std::shared_ptr& changelog_file = compact_result.Changelog()[0]; + ASSERT_EQ(4, changelog_file->row_count); + ASSERT_EQ(FileSource::Append(), changelog_file->file_source); + ASSERT_TRUE(StringUtils::EndsWith(changelog_file->file_name, ".parquet")); + std::string changelog_file_name = table_path + "/bucket-0/" + changelog_file->file_name; + + std::shared_ptr write_schema = SpecialFields::CompleteSequenceAndValueKindField( + DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields())); + ASSERT_OK_AND_ASSIGN(auto physical_schema, MapSharedShreddingUtils::LogicalToPhysicalSchema( + write_schema, {{"tags", 3}})); + std::shared_ptr expected_changelog; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_(physical_schema->fields()), {R"([ + [0, 0, 1, [[0, -1, -1], 11, null, null, null]], + [3, 0, 2, [[1, -1, -1], 22, null, null, null]], + [1, 0, 3, [[2, -1, -1], 33, null, null, null]], + [2, 0, 5, [[3, -1, -1], 55, null, null, null]] + ])"}, + &expected_changelog) + .ok()); + CheckResult(changelog_file_name, table_schema, "parquet", expected_changelog); + + MapSharedShreddingFieldMeta expected_meta; + expected_meta.name_to_id = {{"a", 0}, {"b", 1}, {"c", 2}, {"e", 3}}; + expected_meta.field_to_columns = {{0, {0}}, {1, {0}}, {2, {0}}, {3, {0}}}; + expected_meta.num_columns = 3; + expected_meta.max_row_width = 1; + CheckShreddingFileSchema(changelog_file_name, table_schema, "parquet", physical_schema, + /*field_index=*/3, expected_meta); } TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowUpgrade) { @@ -525,7 +741,8 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestFirstRowUpgrade) { BinaryRowGenerator::GenerateStats({1, 5}, {5, 33}, {0, 0}, pool_.get()), /*min_sequence_number=*/0l, /*max_sequence_number=*/2l, /*schema_id=*/0, /*level=*/5, std::vector>(), Timestamp(0l, 0), /*delete_row_count=*/0, - nullptr, FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + nullptr, FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_TRUE(expected_file_meta->TEST_Equal(*compact_file_meta)) << compact_file_meta->ToString(); } @@ -639,7 +856,7 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestRewriteWithAllHighLevel) { TEST_F(LookupMergeTreeCompactRewriterTest, TestRewriteWithForceLookupAndSumAgg) { std::map options = {{Options::MERGE_ENGINE, "aggregation"}, {Options::FILE_FORMAT, "orc"}, - {Options::FORCE_LOOKUP, "true"}, + {Options::CHANGELOG_PRODUCER, "lookup"}, {Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}}; ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap(options)); ASSERT_OK_AND_ASSIGN(auto table_path, CreateTable(options)); @@ -668,6 +885,7 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestRewriteWithForceLookupAndSumAgg) /*output_level=*/4, /*drop_delete=*/true, runs)); ASSERT_EQ(2, compact_result.Before().size()); ASSERT_EQ(1, compact_result.After().size()); + ASSERT_EQ(1, compact_result.Changelog().size()); const auto& compact_file_meta = compact_result.After()[0]; // check compact file exist @@ -688,6 +906,20 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestRewriteWithForceLookupAndSumAgg) &expected_array); ASSERT_TRUE(array_status.ok()); CheckResult(compact_file_name, table_schema, "orc", expected_array); + + const auto& changelog_file_meta = compact_result.Changelog()[0]; + std::string changelog_file_name = table_path + "/bucket-0/" + changelog_file_meta->file_name; + std::shared_ptr expected_changelog; + auto changelog_status = + arrow::ipc::internal::json::ChunkedArrayFromJSON(type_with_special_fields, {R"([ +[6, 0, 2, 244], +[4, 0, 4, 44], +[2, 1, 5, 55], +[7, 2, 5, 615] +])"}, + &expected_changelog); + ASSERT_TRUE(changelog_status.ok()); + CheckResult(changelog_file_name, table_schema, "orc", expected_changelog); } TEST_F(LookupMergeTreeCompactRewriterTest, TestRewriteWithDvAndDeduplicate) { @@ -1069,7 +1301,8 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { BinaryRowGenerator::GenerateStats({1, 5}, {5, 33}, {0, 0}, pool_.get()), /*min_sequence_number=*/0l, /*max_sequence_number=*/3l, /*schema_id=*/0, level, std::vector>(), Timestamp(0l, 0), delete_row_count, nullptr, - FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); }; { std::map options = {}; @@ -1080,7 +1313,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/1, /*delete_row_count=*/std::nullopt); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::NoChangelogNoRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/2, file)); @@ -1095,7 +1330,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/0, /*delete_row_count=*/std::nullopt); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::ChangelogWithRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/5, file)); @@ -1114,7 +1351,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/0, /*delete_row_count=*/1); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::ChangelogWithRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/2, file)); @@ -1129,7 +1368,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/0, /*delete_row_count=*/std::nullopt); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::ChangelogNoRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/5, file)); @@ -1144,7 +1385,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/0, /*delete_row_count=*/std::nullopt); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::ChangelogNoRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/2, file)); @@ -1160,7 +1403,9 @@ TEST_F(LookupMergeTreeCompactRewriterTest, TestGenerateUpgradeStrategy) { /*trimmed_primary_keys=*/{"key"}, core_options, /*data_schema=*/nullptr, /*write_schema=*/nullptr, /*path_factory_cache=*/nullptr, /*merge_file_split_read=*/nullptr, /*merge_function_wrapper_factory=*/nullptr, - cancellation_controller, /*remote_lookup_file_manager=*/nullptr, pool_); + /*changelog_merge_function_wrapper_factory=*/nullptr, + /*produce_changelog=*/false, cancellation_controller, + /*remote_lookup_file_manager=*/nullptr, pool_); auto file = create_meta(/*level=*/0, /*delete_row_count=*/std::nullopt); ASSERT_EQ(ChangelogMergeTreeRewriter::UpgradeStrategy::ChangelogWithRewrite(), rewriter.GenerateUpgradeStrategy(/*output_level=*/2, file)); diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp index e5cb3abf6..258336698 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp @@ -25,6 +25,7 @@ #include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" #include "paimon/core/mergetree/compact/early_full_compaction.h" #include "paimon/core/mergetree/compact/force_up_level0_compaction.h" +#include "paimon/core/mergetree/compact/internal_row_equalizer.h" #include "paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter.h" #include "paimon/core/mergetree/compact/merge_tree_compact_manager.h" #include "paimon/core/mergetree/compact/merge_tree_compact_rewriter.h" @@ -50,6 +51,24 @@ namespace paimon { namespace { +Result CreateChangelogValueEqualizer( + const std::shared_ptr& schema, const CoreOptions& options, + bool produce_changelog) { + if (!produce_changelog || !options.ChangelogRowDeduplicate()) { + return FieldsComparator::FieldComparatorFunc(); + } + return InternalRowEqualizer::Create(schema, options.GetChangelogRowDeduplicateIgnoreFields()); +} + +Result> CreateChangelogValueSerializer( + const std::shared_ptr& schema, bool produce_changelog, + const std::shared_ptr& pool) { + if (!produce_changelog) { + return std::unique_ptr(); + } + return RowCompactedSerializer::Create(schema, pool); +} + template Result>> CreateLookupLevelsInternal( const CoreOptions& options, const std::shared_ptr& schema_manager, @@ -174,6 +193,8 @@ Result> MergeTreeCompactManagerFactory::CreateL const LookupStrategy& lookup_strategy, const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller) const { + const bool should_produce_changelog = + lookup_strategy.produce_changelog && !ignore_previous_files_; PAIMON_ASSIGN_OR_RAISE(std::shared_ptr remote_lookup_file_manager, CreateRemoteLookupFileManager(partition, bucket)); if (lookup_strategy.is_first_row) { @@ -190,157 +211,133 @@ Result> MergeTreeCompactManagerFactory::CreateL options_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, table_schema_, partition, bucket, levels, processor_factory, dv_maintainer, lookup_file_cache_, remote_lookup_file_manager, pool_)); - auto merge_function_wrapper_factory = - [lookup_levels_ptr = lookup_levels.get(), ignore_delete = options_.IgnoreDelete()]( - int32_t output_level) -> Result>> { - std::shared_ptr> merge_function_wrapper = + auto merge_function_wrapper_factory = [lookup_levels_ptr = lookup_levels.get(), + data_schema = schema_, should_produce_changelog, + ignore_delete = options_.IgnoreDelete(), + pool = pool_](int32_t output_level) + -> Result>> { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr value_serializer, + CreateChangelogValueSerializer(data_schema, should_produce_changelog, pool)); + std::shared_ptr> merge_function_wrapper = LookupMergeTreeCompactRewriter::CreateFirstRowMergeFunctionWrapper( std::make_unique(ignore_delete), output_level, - lookup_levels_ptr); + std::move(value_serializer), lookup_levels_ptr); return merge_function_wrapper; }; - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr> rewriter, - LookupMergeTreeCompactRewriter::Create( - max_level, std::move(lookup_levels), dv_maintainer, - std::move(merge_function_wrapper_factory), bucket, partition, - table_schema_, path_factory_cache, options_, - cancellation_controller, remote_lookup_file_manager, pool_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr> rewriter, + LookupMergeTreeCompactRewriter::Create( + max_level, std::move(lookup_levels), dv_maintainer, + std::move(merge_function_wrapper_factory), bucket, partition, table_schema_, + path_factory_cache, options_, should_produce_changelog, cancellation_controller, + remote_lookup_file_manager, pool_)); return std::shared_ptr(std::move(rewriter)); } if (lookup_strategy.deletion_vector) { return CreateLookupRewriterWithDeletionVector( partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, - path_factory_cache, cancellation_controller, remote_lookup_file_manager); + should_produce_changelog, path_factory_cache, cancellation_controller, + remote_lookup_file_manager); } return CreateLookupRewriterWithoutDeletionVector( - partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, path_factory_cache, - cancellation_controller, remote_lookup_file_manager); + partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, + should_produce_changelog, path_factory_cache, cancellation_controller, + remote_lookup_file_manager); } +template Result> -MergeTreeCompactManagerFactory::CreateLookupRewriterWithDeletionVector( +MergeTreeCompactManagerFactory::CreateLookupRewriterInternal( const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, const std::shared_ptr& dv_maintainer, int32_t max_level, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, + const std::shared_ptr::Factory>& processor_factory, const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager) const { - auto merge_engine = options_.GetMergeEngine(); PAIMON_ASSIGN_OR_RAISE(std::vector trimmed_primary_keys, table_schema_->TrimmedPrimaryKeys()); - if (lookup_strategy.produce_changelog || merge_engine != MergeEngine::DEDUPLICATE || - !options_.GetSequenceField().empty()) { - auto processor_factory = std::make_shared(schema_); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr> lookup_levels, - CreateLookupLevelsInternal( - options_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, - table_schema_, partition, bucket, levels, processor_factory, dv_maintainer, - lookup_file_cache_, remote_lookup_file_manager, pool_)); - auto merge_function_wrapper_factory = - [data_schema = schema_, options = options_, trimmed_primary_keys, - lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, pool = pool_, - user_defined_seq_comparator = user_defined_seq_comparator_]( - int32_t output_level) -> Result>> { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, - PrimaryKeyTableUtils::CreateMergeFunction( - data_schema, trimmed_primary_keys, options, pool)); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, - LookupMergeTreeCompactRewriter:: - CreateLookupMergeFunctionWrapper( - std::make_unique(std::move(merge_func)), output_level, - dv_maintainer_ptr, lookup_strategy, user_defined_seq_comparator, - lookup_levels_ptr)); - return merge_function_wrapper; - }; - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr> rewriter, - LookupMergeTreeCompactRewriter::Create( - max_level, std::move(lookup_levels), dv_maintainer, - std::move(merge_function_wrapper_factory), bucket, partition, table_schema_, - path_factory_cache, options_, cancellation_controller, remote_lookup_file_manager, - pool_)); - return std::shared_ptr(std::move(rewriter)); - } - auto processor_factory = std::make_shared(); PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr> lookup_levels, - CreateLookupLevelsInternal( - options_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, - table_schema_, partition, bucket, levels, processor_factory, dv_maintainer, - lookup_file_cache_, remote_lookup_file_manager, pool_)); + std::unique_ptr> lookup_levels, + CreateLookupLevelsInternal(options_, schema_manager_, io_manager_, cache_manager_, + file_store_path_factory_, table_schema_, partition, bucket, + levels, processor_factory, dv_maintainer, lookup_file_cache_, + remote_lookup_file_manager, pool_)); auto merge_function_wrapper_factory = [data_schema = schema_, options = options_, trimmed_primary_keys, - lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, pool = pool_, - user_defined_seq_comparator = user_defined_seq_comparator_]( - int32_t output_level) -> Result>> { + lookup_levels_ptr = lookup_levels.get(), lookup_strategy, should_produce_changelog, + dv_maintainer_ptr = dv_maintainer, + user_defined_seq_comparator = user_defined_seq_comparator_, + pool = pool_](int32_t output_level) + -> Result>> { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, PrimaryKeyTableUtils::CreateMergeFunction( data_schema, trimmed_primary_keys, options, pool)); PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, - LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( - std::make_unique(std::move(merge_func)), output_level, - dv_maintainer_ptr, lookup_strategy, user_defined_seq_comparator, - lookup_levels_ptr)); - return merge_function_wrapper; + std::unique_ptr value_serializer, + CreateChangelogValueSerializer(data_schema, should_produce_changelog, pool)); + PAIMON_ASSIGN_OR_RAISE( + FieldsComparator::FieldComparatorFunc value_equalizer, + CreateChangelogValueEqualizer(data_schema, options, should_produce_changelog)); + return LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( + std::make_unique(std::move(merge_func)), output_level, + dv_maintainer_ptr, lookup_strategy, should_produce_changelog, + user_defined_seq_comparator, std::move(value_serializer), std::move(value_equalizer), + lookup_levels_ptr); }; - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr> rewriter, - LookupMergeTreeCompactRewriter::Create( - max_level, std::move(lookup_levels), dv_maintainer, - std::move(merge_function_wrapper_factory), bucket, partition, - table_schema_, path_factory_cache, options_, cancellation_controller, - remote_lookup_file_manager, pool_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr> rewriter, + LookupMergeTreeCompactRewriter::Create( + max_level, std::move(lookup_levels), dv_maintainer, + std::move(merge_function_wrapper_factory), bucket, partition, table_schema_, + path_factory_cache, options_, should_produce_changelog, cancellation_controller, + remote_lookup_file_manager, pool_)); return std::shared_ptr(std::move(rewriter)); } Result> -MergeTreeCompactManagerFactory::CreateLookupRewriterWithoutDeletionVector( +MergeTreeCompactManagerFactory::CreateLookupRewriterWithDeletionVector( const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, const std::shared_ptr& dv_maintainer, int32_t max_level, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager) const { - PAIMON_ASSIGN_OR_RAISE(std::vector trimmed_primary_keys, - table_schema_->TrimmedPrimaryKeys()); - auto processor_factory = std::make_shared(schema_); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr> lookup_levels, - CreateLookupLevelsInternal( - options_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, - table_schema_, partition, bucket, levels, processor_factory, dv_maintainer, - lookup_file_cache_, remote_lookup_file_manager, pool_)); - auto merge_function_wrapper_factory = - [data_schema = schema_, options = options_, trimmed_primary_keys, - lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, pool = pool_, - user_defined_seq_comparator = user_defined_seq_comparator_]( - int32_t output_level) -> Result>> { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, - PrimaryKeyTableUtils::CreateMergeFunction( - data_schema, trimmed_primary_keys, options, pool)); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, - LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( - std::make_unique(std::move(merge_func)), output_level, - dv_maintainer_ptr, lookup_strategy, user_defined_seq_comparator, - lookup_levels_ptr)); - return merge_function_wrapper; - }; + auto merge_engine = options_.GetMergeEngine(); + if (lookup_strategy.produce_changelog || merge_engine != MergeEngine::DEDUPLICATE || + !options_.GetSequenceField().empty()) { + std::shared_ptr::Factory> processor_factory = + std::make_shared(schema_); + return CreateLookupRewriterInternal( + partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, + should_produce_changelog, processor_factory, path_factory_cache, + cancellation_controller, remote_lookup_file_manager); + } + std::shared_ptr::Factory> processor_factory = + std::make_shared(); + return CreateLookupRewriterInternal( + partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, + should_produce_changelog, processor_factory, path_factory_cache, cancellation_controller, + remote_lookup_file_manager); +} - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr> rewriter, - LookupMergeTreeCompactRewriter::Create( - max_level, std::move(lookup_levels), dv_maintainer, - std::move(merge_function_wrapper_factory), bucket, partition, - table_schema_, path_factory_cache, options_, cancellation_controller, - remote_lookup_file_manager, pool_)); - return std::shared_ptr(std::move(rewriter)); +Result> +MergeTreeCompactManagerFactory::CreateLookupRewriterWithoutDeletionVector( + const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, + const std::shared_ptr& dv_maintainer, int32_t max_level, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, + const std::shared_ptr& path_factory_cache, + const std::shared_ptr& cancellation_controller, + const std::shared_ptr& remote_lookup_file_manager) const { + std::shared_ptr::Factory> processor_factory = + std::make_shared(schema_); + return CreateLookupRewriterInternal( + partition, bucket, levels, dv_maintainer, max_level, lookup_strategy, + should_produce_changelog, processor_factory, path_factory_cache, cancellation_controller, + remote_lookup_file_manager); } Result> diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.h b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.h index 13a059620..2eaa6ffd4 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.h +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.h @@ -71,7 +71,8 @@ class MergeTreeCompactManagerFactory { const std::shared_ptr& io_manager, const std::shared_ptr& cache_manager, const std::shared_ptr& file_store_path_factory, - const std::string& root_path, const std::shared_ptr& pool) + const std::string& root_path, bool ignore_previous_files, + const std::shared_ptr& pool) : options_(options), pool_(pool), key_comparator_(key_comparator), @@ -83,7 +84,8 @@ class MergeTreeCompactManagerFactory { io_manager_(io_manager), cache_manager_(cache_manager), file_store_path_factory_(file_store_path_factory), - root_path_(root_path) {} + root_path_(root_path), + ignore_previous_files_(ignore_previous_files) {} std::shared_ptr CreateCompactStrategy() const; @@ -115,10 +117,20 @@ class MergeTreeCompactManagerFactory { const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller) const; + template + Result> CreateLookupRewriterInternal( + const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, + const std::shared_ptr& dv_maintainer, int32_t max_level, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, + const std::shared_ptr::Factory>& processor_factory, + const std::shared_ptr& path_factory_cache, + const std::shared_ptr& cancellation_controller, + const std::shared_ptr& remote_lookup_file_manager) const; + Result> CreateLookupRewriterWithDeletionVector( const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, const std::shared_ptr& dv_maintainer, int32_t max_level, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager) const; @@ -126,7 +138,7 @@ class MergeTreeCompactManagerFactory { Result> CreateLookupRewriterWithoutDeletionVector( const BinaryRow& partition, int32_t bucket, const std::shared_ptr& levels, const std::shared_ptr& dv_maintainer, int32_t max_level, - const LookupStrategy& lookup_strategy, + const LookupStrategy& lookup_strategy, bool should_produce_changelog, const std::shared_ptr& path_factory_cache, const std::shared_ptr& cancellation_controller, const std::shared_ptr& remote_lookup_file_manager) const; @@ -146,6 +158,7 @@ class MergeTreeCompactManagerFactory { std::shared_ptr cache_manager_; std::shared_ptr file_store_path_factory_; std::string root_path_; + bool ignore_previous_files_; std::shared_ptr lookup_file_cache_; }; diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory_test.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory_test.cpp index 7473b7c0b..06a4ec197 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory_test.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory_test.cpp @@ -60,7 +60,8 @@ class MergeTreeCompactManagerFactoryStrategyTest : public ::testing::Test { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); return {level, SortedRun::FromSingle(file_meta)}; } @@ -87,6 +88,7 @@ class MergeTreeCompactManagerFactoryStrategyTest : public ::testing::Test { /*cache_manager=*/nullptr, /*file_store_path_factory=*/nullptr, /*root_path=*/"", + /*ignore_previous_files=*/false, /*pool=*/nullptr); } }; @@ -332,7 +334,8 @@ TEST_F(MergeTreeCompactManagerFactoryWriteTest, ASSERT_NOK_WITH_MSG(CreateSingleStringFileStoreWrite( {{"bucket", "1"}, {Options::CHANGELOG_PRODUCER, "full-compaction"}}, /*with_io_manager=*/false), - "C++ Paimon does not support changelog-producer yet"); + "C++ Paimon only supports 'none', 'input' and 'lookup' " + "changelog-producer now"); } TEST_F(MergeTreeCompactManagerFactoryWriteTest, @@ -377,13 +380,11 @@ TEST_F(MergeTreeCompactManagerFactoryWriteTest, } TEST_F(MergeTreeCompactManagerFactoryWriteTest, - TestCreateFileStoreWriteShouldFailWhenLookupChangelogConfigured) { - ASSERT_NOK_WITH_MSG( - CreateSingleStringFileStoreWrite({{"bucket", "1"}, - {Options::DELETION_VECTORS_ENABLED, "true"}, - {Options::CHANGELOG_PRODUCER, "lookup"}}, - /*with_io_manager=*/true), - "C++ Paimon does not support changelog-producer yet"); + TestCreateFileStoreWriteShouldSucceedWhenLookupChangelogConfigured) { + ASSERT_OK(CreateSingleStringFileStoreWrite({{"bucket", "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {Options::CHANGELOG_PRODUCER, "lookup"}}, + /*with_io_manager=*/true)); } } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp index 19b8bc363..9186eb1c4 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_test.cpp @@ -142,7 +142,7 @@ class TestRewriter final : public CompactRewriter { /*max_sequence_number=*/max_sequence, /*schema_id=*/0, output_level, std::vector>(), Timestamp(1, 0), std::nullopt, nullptr, FileSource::Append(), std::nullopt, - std::nullopt, std::nullopt, std::nullopt); + std::nullopt, std::nullopt, std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return CompactResult(before, {after}); } @@ -208,7 +208,7 @@ class MergeTreeCompactManagerTest : public testing::Test { /*max_sequence_number=*/max_sequence, /*schema_id=*/0, minmax.level, std::vector>(), Timestamp(1, 0), std::nullopt, nullptr, FileSource::Append(), std::nullopt, - std::nullopt, std::nullopt, std::nullopt); + std::nullopt, std::nullopt, std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } StrategyFn TestStrategy() { diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp index 1b64be2da..37c1cfa04 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp @@ -23,14 +23,12 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" -#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/scope_guard.h" -#include "paimon/core/io/key_value_data_file_writer_factory.h" +#include "paimon/core/io/key_value_data_file_writer_factories.h" #include "paimon/core/io/key_value_meta_projection_consumer.h" #include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/io/row_to_arrow_array_converter.h" -#include "paimon/core/io/shredding_key_value_data_file_writer_factory.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/format/file_format.h" @@ -84,14 +82,9 @@ Result> MergeTreeCompactRewriter::Crea .WithMemoryPool(pool); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, read_context_builder.Finish()); - // TODO(xinyu.lxy): temporarily disabled pre-buffer for parquet, which may cause high memory - // usage during compaction. Will fix via parquet format refactor. - auto new_options = options.ToMap(); - if (new_options.find("parquet.read.enable-pre-buffer") == new_options.end()) { - new_options["parquet.read.enable-pre-buffer"] = "false"; - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_context, - InternalReadContext::Create(read_context, table_schema, new_options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr internal_context, + InternalReadContext::Create(read_context, table_schema, options.ToMap())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr path_factory, path_factory_cache->GetOrCreatePathFactory(options.GetFileFormat()->Identifier())); @@ -138,20 +131,31 @@ MergeTreeCompactRewriter::CreateRollingRowWriter(int32_t level) { auto format = options_.GetWriteFileFormat(level); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, CreateDataFilePathFactory(format->Identifier())); - std::shared_ptr>> factory; PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr plan_factory, - ShreddingWritePlanFactories::SelectActive(options_, write_schema_, pool_)); - if (plan_factory != nullptr) { - factory = std::make_shared( + std::shared_ptr factory, + KeyValueDataFileWriterFactories::Create( options_, schema_id_, write_schema_, level, FileSource::Compact(), trimmed_primary_keys_, data_file_path_factory, /*create_stats_extractor=*/true, - plan_factory, pool_); - } else { - factory = std::make_shared( - options_, schema_id_, write_schema_, level, FileSource::Compact(), - trimmed_primary_keys_, data_file_path_factory, /*create_stats_extractor=*/true, pool_); + /*is_changelog=*/false, pool_)); + return std::make_unique( + options_.GetTargetFileSize(/*has_primary_key=*/true), + /*target_file_row_num=*/std::numeric_limits::max(), factory); +} + +Result> +MergeTreeCompactRewriter::CreateRollingChangelogWriter(int32_t level) { + std::shared_ptr format = options_.GetChangelogFileFormat(); + if (!format) { + format = options_.GetWriteFileFormat(level); } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + CreateDataFilePathFactory(format->Identifier())); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr factory, + KeyValueDataFileWriterFactories::Create( + options_, schema_id_, write_schema_, level, FileSource::Append(), trimmed_primary_keys_, + data_file_path_factory, /*create_stats_extractor=*/true, + /*is_changelog=*/true, pool_)); return std::make_unique( options_.GetTargetFileSize(/*has_primary_key=*/true), /*target_file_row_num=*/std::numeric_limits::max(), factory); @@ -183,6 +187,19 @@ Result> MergeTreeCompactRewriter::CreateDat return path_factory->CreateDataFilePathFactory(partition_, bucket_); } +Result> +MergeTreeCompactRewriter::CreateRawSortMergeReaderForSection( + const std::vector& section) { + if (!merge_file_split_read_) { + return Status::Invalid( + "merge_file_split_read in MergeTreeCompactRewriter cannot be nullptr"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + CreateDataFilePathFactory(options_.GetFileFormat()->Identifier())); + return merge_file_split_read_->CreateRawSortMergeReaderForSection( + section, partition_, dv_factory_, /*predicate=*/nullptr, data_file_path_factory); +} + Status MergeTreeCompactRewriter::MergeReadAndWrite( int32_t output_level, bool drop_delete, const std::vector& section, const MergeTreeCompactRewriter::KeyValueConsumerCreator& create_consumer, @@ -227,11 +244,12 @@ Status MergeTreeCompactRewriter::MergeReadAndWrite( return Status::OK(); } - // consumer batch size is WriteBatchSize + std::unique_ptr producer = + std::make_unique(std::move(sort_merge_reader), + options_.GetWriteBatchSize()); auto async_key_value_producer_consumer = std::make_shared>( - std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), - /*projection_thread_num=*/1, pool_); + std::move(producer), create_consumer, /*projection_thread_num=*/1); reader_holders.push_back(async_key_value_producer_consumer); // read KeyValueBatch from SortMergeReader and write to RollingWriter while (true) { diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h index c9c624704..c15e16fc5 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.h @@ -92,6 +92,8 @@ class MergeTreeCompactRewriter : public CompactRewriter { Result> CreateRollingRowWriter(int32_t level); + Result> CreateRollingChangelogWriter(int32_t level); + Result GenerateKeyValueConsumer() const; Status MergeReadAndWrite(int32_t output_level, bool drop_delete, @@ -100,6 +102,13 @@ class MergeTreeCompactRewriter : public CompactRewriter { KeyValueRollingFileWriter* rolling_writer, std::vector>* reader_holders_ptr); + Result> CreateRawSortMergeReaderForSection( + const std::vector& section); + + bool IsCancelled() const { + return cancellation_controller_->IsCancelled(); + } + protected: CoreOptions options_; std::unique_ptr merge_file_split_read_; @@ -108,7 +117,6 @@ class MergeTreeCompactRewriter : public CompactRewriter { Result> CreateDataFilePathFactory( const std::string& format); - private: std::shared_ptr pool_; BinaryRow partition_; int32_t bucket_; diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp index cb26540aa..dc19c4726 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter_test.cpp @@ -148,7 +148,8 @@ TEST_F(MergeTreeCompactRewriterTest, TestSimple) { pool_.get()), /*min_sequence_number=*/0l, /*max_sequence_number=*/10l, /*schema_id=*/0, /*level=*/5, std::vector>(), Timestamp(0l, 0), /*delete_row_count=*/0, - nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_TRUE(expected_file_meta->TEST_Equal(*compact_file_meta)); // check compact file exist std::string compact_file_name = @@ -245,7 +246,8 @@ TEST_F(MergeTreeCompactRewriterTest, TestNotDropDelete) { pool_.get()), /*min_sequence_number=*/0l, /*max_sequence_number=*/11l, /*schema_id=*/0, /*level=*/5, std::vector>(), Timestamp(0l, 0), /*delete_row_count=*/2, - nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + nullptr, FileSource::Compact(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_TRUE(expected_file_meta->TEST_Equal(*compact_file_meta)); std::string compact_file_name = diff --git a/src/paimon/core/mergetree/compact/universal_compaction_test.cpp b/src/paimon/core/mergetree/compact/universal_compaction_test.cpp index 74abacaef..9b2e6cd1f 100644 --- a/src/paimon/core/mergetree/compact/universal_compaction_test.cpp +++ b/src/paimon/core/mergetree/compact/universal_compaction_test.cpp @@ -56,7 +56,7 @@ class UniversalCompactionTest : public testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return {level, SortedRun::FromSingle(file_meta)}; } diff --git a/src/paimon/core/mergetree/external_sort_buffer.cpp b/src/paimon/core/mergetree/external_sort_buffer.cpp index 9bfeec899..718a2de62 100644 --- a/src/paimon/core/mergetree/external_sort_buffer.cpp +++ b/src/paimon/core/mergetree/external_sort_buffer.cpp @@ -222,10 +222,11 @@ Result ExternalSortBuffer::SpillToDisk( -> Result>> { return KeyValueMetaProjectionConsumer::Create(target_schema, pool); }; + std::unique_ptr producer = + std::make_unique(std::move(sorted_reader), write_batch_size); auto async_key_value_producer_consumer = std::make_unique>( - std::move(sorted_reader), create_consumer, write_batch_size, - /*projection_thread_num=*/1, pool_); + std::move(producer), create_consumer, /*projection_thread_num=*/1); auto close_guard = ScopeGuard([&]() { async_key_value_producer_consumer->Close(); }); while (true) { diff --git a/src/paimon/core/mergetree/levels_test.cpp b/src/paimon/core/mergetree/levels_test.cpp index 37a347f92..d31a31c92 100644 --- a/src/paimon/core/mergetree/levels_test.cpp +++ b/src/paimon/core/mergetree/levels_test.cpp @@ -45,7 +45,7 @@ class LevelsTest : public testing::Test { max_sequence_number, /*schema_id=*/0, level, std::vector>(), Timestamp(ts_second, 0l), std::nullopt, nullptr, FileSource::Append(), std::nullopt, - std::nullopt, std::nullopt, std::nullopt); + std::nullopt, std::nullopt, std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateComparator() const { @@ -261,7 +261,8 @@ TEST_F(LevelsTest, TestUpdateDropFileCallbackExcludesUpgradeFiles) { file_level0->key_stats, file_level0->value_stats, file_level0->min_sequence_number, file_level0->max_sequence_number, file_level0->schema_id, /*level=*/1, file_level0->extra_files, file_level0->creation_time, std::nullopt, nullptr, - FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, std::nullopt); + FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); std::vector> before = {file_level0}; std::vector> after = {upgraded_file}; diff --git a/src/paimon/core/mergetree/lookup_levels.cpp b/src/paimon/core/mergetree/lookup_levels.cpp index 8b5f69e77..512f28728 100644 --- a/src/paimon/core/mergetree/lookup_levels.cpp +++ b/src/paimon/core/mergetree/lookup_levels.cpp @@ -66,14 +66,9 @@ Result>> LookupLevels::Create( .WithMemoryPool(pool); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, read_context_builder.Finish()); - // TODO(xinyu.lxy): temporarily disabled pre-buffer for parquet, which may cause high memory - // usage during compaction. Will fix via parquet format refactor. - auto new_options = options.ToMap(); - if (new_options.find("parquet.read.enable-pre-buffer") == new_options.end()) { - new_options["parquet.read.enable-pre-buffer"] = "false"; - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema, new_options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr internal_read_context, + InternalReadContext::Create(read_context, table_schema, options.ToMap())); auto split_read = std::make_unique(path_factory, internal_read_context, pool, CreateDefaultExecutor()); @@ -162,8 +157,8 @@ LookupLevels::LookupLevels( lookup_store_factory_(lookup_store_factory), lookup_file_cache_(lookup_file_cache), remote_lookup_file_manager_(remote_lookup_file_manager) { - if constexpr (std::is_same_v) { - // if T is FilePosition, only read key fields to create sst file is enough + if constexpr (std::is_same_v || std::is_same_v) { + // FilePosition and first-row lookup do not persist values, so reading key fields is enough. value_schema_ = key_schema_; } else { value_schema_ = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); @@ -332,16 +327,6 @@ std::optional LookupLevels::TryToDownloadRemoteSst( template Status LookupLevels::CreateSstFileFromDataFile(const std::shared_ptr& file, const std::string& kv_file_path) { - if constexpr (std::is_same_v) { - // Short-circuit logic: if T is bool, just write empty lookup file. - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr bloom_filter, - LookupStoreFactory::BfGenerator(file->row_count, options_, pool_.get())); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr kv_writer, - lookup_store_factory_->CreateWriter(fs_, kv_file_path, bloom_filter, pool_)); - return kv_writer->Close(); - } // Prepare reader to iterate KeyValue PAIMON_ASSIGN_OR_RAISE( std::vector> raw_readers, diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 3b6806c73..748740f88 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -27,7 +28,6 @@ #include "arrow/api.h" #include "arrow/c/abi.h" #include "arrow/c/helpers.h" -#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -37,13 +37,13 @@ #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_data_file_writer_factory.h" +#include "paimon/core/io/key_value_data_file_writer_factories.h" #include "paimon/core/io/key_value_meta_projection_consumer.h" #include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/io/row_to_arrow_array_converter.h" -#include "paimon/core/io/shredding_key_value_data_file_writer_factory.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h" +#include "paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h" #include "paimon/core/mergetree/write_buffer.h" #include "paimon/core/utils/commit_increment.h" @@ -104,8 +104,13 @@ Status MergeTreeWriter::DoClose() { // delete temporary files std::vector> delete_files; - delete_files.reserve(new_files_.size() + compact_after_.size()); + delete_files.reserve(new_files_.size() + new_changelog_files_.size() + compact_after_.size() + + compact_changelog_files_.size()); delete_files.insert(delete_files.end(), new_files_.begin(), new_files_.end()); + delete_files.insert(delete_files.end(), new_changelog_files_.begin(), + new_changelog_files_.end()); + delete_files.insert(delete_files.end(), compact_changelog_files_.begin(), + compact_changelog_files_.end()); for (const auto& file : compact_after_) { // Upgrade file is required by previous snapshot, so we should ensure that this file is // not the output of upgraded. @@ -125,9 +130,11 @@ Status MergeTreeWriter::DoClose() { write_buffer_->Clear(); new_files_.clear(); + new_changelog_files_.clear(); deleted_files_.clear(); compact_before_.clear(); compact_after_.clear(); + compact_changelog_files_.clear(); if (compact_deletion_file_) { compact_deletion_file_->Clean(); @@ -154,6 +161,65 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReadersToFiles( + 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"); + } + } + + // 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_); + raw_readers_guard.Release(); + // 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 + std::unique_ptr producer = + std::make_unique(std::move(sort_merge_reader), + options_.GetWriteBatchSize()); + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(producer), create_consumer, /*projection_thread_num=*/1); + 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); } @@ -210,7 +276,9 @@ Status MergeTreeWriter::UpdateCompactResult(const std::shared_ptr compact_after_.insert(compact_after_.end(), compact_result->After().begin(), compact_result->After().end()); - // TODO(yonghao.fyh): support compact changelog + compact_changelog_files_.insert(compact_changelog_files_.end(), + compact_result->Changelog().begin(), + compact_result->Changelog().end()); return UpdateCompactDeletionFile(compact_result->DeletionFile()); } @@ -256,49 +324,67 @@ 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> + async_changelog_producer_consumer; 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(); + changelog_writer; + std::vector> flushed_changelog_files; + ScopeGuard changelog_write_guard([&]() -> void { + if (changelog_writer) { + changelog_writer->Abort(); + } + if (async_changelog_producer_consumer) { + async_changelog_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; + if (options_.GetChangelogProducer() == ChangelogProducer::INPUT) { + PAIMON_ASSIGN_OR_RAISE(std::vector> raw_readers, + write_buffer_->CreateRawReaders()); + auto raw_sort_merge_reader = std::make_unique( + std::move(raw_readers), key_comparator_, user_defined_seq_comparator_, + /*merge_function_wrapper=*/nullptr); + std::unique_ptr producer = + std::make_unique(std::move(raw_sort_merge_reader), + options_.GetWriteBatchSize()); + async_changelog_producer_consumer = + std::make_unique>( + std::move(producer), create_consumer, /*projection_thread_num=*/1); + PAIMON_ASSIGN_OR_RAISE(changelog_writer, CreateRollingChangelogWriter()); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_changelog_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(changelog_writer->Write(std::move(key_value_batch))); } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + PAIMON_RETURN_NOT_OK(changelog_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(flushed_changelog_files, changelog_writer->GetResult()); } - 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)); + + // 1. flush write buffer to get sorted readers + // Flush write buffer to get sorted and merged data readers. + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + write_buffer_->CreateReaders()); + PAIMON_RETURN_NOT_OK(WriteSortedReadersToFiles(std::move(readers))); + if (async_changelog_producer_consumer) { + async_changelog_producer_consumer->Close(); + } + + new_changelog_files_.insert(new_changelog_files_.end(), flushed_changelog_files.begin(), + flushed_changelog_files.end()); + + changelog_write_guard.Release(); + + if (changelog_writer) { + metrics_->Merge(changelog_writer->GetMetrics()); } - metrics_->Merge(rolling_writer->GetMetrics()); } PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction)); PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction)); @@ -306,14 +392,18 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, } Result MergeTreeWriter::DrainIncrement() { - DataIncrement data_increment(std::move(new_files_), std::move(deleted_files_), {}); - CompactIncrement compact_increment(std::move(compact_before_), std::move(compact_after_), {}); + DataIncrement data_increment(std::move(new_files_), std::move(deleted_files_), + std::move(new_changelog_files_)); + CompactIncrement compact_increment(std::move(compact_before_), std::move(compact_after_), + std::move(compact_changelog_files_)); auto drain_deletion_file = compact_deletion_file_; new_files_.clear(); + new_changelog_files_.clear(); deleted_files_.clear(); compact_before_.clear(); compact_after_.clear(); + compact_changelog_files_.clear(); compact_deletion_file_ = nullptr; return CommitIncrement(data_increment, compact_increment, drain_deletion_file); @@ -321,23 +411,28 @@ Result MergeTreeWriter::DrainIncrement() { Result>>> MergeTreeWriter::CreateRollingRowWriter() const { - std::shared_ptr>> factory; PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr plan_factory, - ShreddingWritePlanFactories::SelectActive(options_, write_schema_, pool_)); - if (plan_factory != nullptr) { - factory = std::make_shared( - options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/true, plan_factory, - pool_); - } else { - factory = std::make_shared( - options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/true, pool_); - } + std::shared_ptr factory, + KeyValueDataFileWriterFactories::Create(options_, schema_id_, write_schema_, /*level=*/0, + FileSource::Append(), trimmed_primary_keys_, + path_factory_, /*create_stats_extractor=*/true, + /*is_changelog=*/false, pool_)); return std::make_unique>>( options_.GetTargetFileSize(/*has_primary_key=*/true), options_.GetTargetFileRowNum(), factory); } +Result>>> +MergeTreeWriter::CreateRollingChangelogWriter() const { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr factory, + KeyValueDataFileWriterFactories::Create(options_, schema_id_, write_schema_, /*level=*/0, + FileSource::Append(), trimmed_primary_keys_, + path_factory_, /*create_stats_extractor=*/true, + /*is_changelog=*/true, pool_)); + return std::make_unique>>( + options_.GetTargetFileSize(/*has_primary_key=*/true), + /*target_file_row_num=*/std::numeric_limits::max(), factory); +} + } // namespace paimon diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index febce2afb..575bd76f6 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 WriteSortedReadersToFiles(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; @@ -100,6 +105,9 @@ class MergeTreeWriter : public BatchWriter { Result>>> CreateRollingRowWriter() const; + Result>>> + CreateRollingChangelogWriter() const; + Status TrySyncLatestCompaction(bool blocking); Status UpdateCompactResult(const std::shared_ptr& compact_result); Status UpdateCompactDeletionFile(const std::shared_ptr& new_deletion_file); @@ -134,9 +142,11 @@ class MergeTreeWriter : public BatchWriter { std::shared_ptr metrics_; std::vector> new_files_; + std::vector> new_changelog_files_; std::vector> deleted_files_; std::vector> compact_before_; std::vector> compact_after_; + std::vector> compact_changelog_files_; std::shared_ptr compact_deletion_file_; }; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 2155647a1..c1c804f9f 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" @@ -37,11 +38,13 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/scope_guard.h" +#include "paimon/common/utils/string_utils.h" #include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/disk/io_manager.h" #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 +55,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 +69,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 { @@ -144,11 +179,13 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { } void CheckFileContent(const std::string& data_file_name, - const std::shared_ptr& expected_array) const { + const std::shared_ptr& expected_array, + const std::string& file_format_name = "orc") const { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(data_file_name)); ASSERT_TRUE(input_stream); - ASSERT_OK_AND_ASSIGN(auto file_format, FileFormatFactory::Get("orc", /*options=*/{})); + ASSERT_OK_AND_ASSIGN(auto file_format, + FileFormatFactory::Get(file_format_name, /*options=*/{})); ASSERT_OK_AND_ASSIGN(auto reader_builder, file_format->CreateReaderBuilder(/*batch_size=*/10)); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, reader_builder->Build(input_stream)); @@ -160,11 +197,13 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { void CheckShreddingFileSchema(const std::string& data_file_name, const std::shared_ptr& expected_physical_schema, int32_t field_index, - const MapSharedShreddingFieldMeta& expected_meta) const { + const MapSharedShreddingFieldMeta& expected_meta, + const std::string& file_format_name = "orc") const { ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, file_system_->Open(data_file_name)); ASSERT_TRUE(input_stream); - ASSERT_OK_AND_ASSIGN(auto file_format, FileFormatFactory::Get("orc", /*options=*/{})); + ASSERT_OK_AND_ASSIGN(auto file_format, + FileFormatFactory::Get(file_format_name, /*options=*/{})); ASSERT_OK_AND_ASSIGN(auto reader_builder, file_format->CreateReaderBuilder(/*batch_size=*/10)); ASSERT_OK_AND_ASSIGN(auto orc_batch_reader, reader_builder->Build(input_stream)); @@ -192,7 +231,8 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } Result> CreateMergeWriter( @@ -211,6 +251,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_; @@ -289,12 +346,188 @@ TEST_P(MergeTreeWriterTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); } +TEST_P(MergeTreeWriterTest, TestInputChangelog) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::CHANGELOG_FILE_PREFIX, "changes-"}, + {Options::CHANGELOG_FILE_FORMAT, "parquet"}})); + + 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(/*last_sequence_number=*/-1, dir->Str(), path_factory, + /*schema_id=*/1, options)); + + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([ + ["Bob", 20, 0, 20.0], + ["Alice", 10, 0, 10.0], + ["Alice", 11, 0, 11.0], + ["Bob", 21, 0, 21.0] + ])") + .ValueOrDie(); + WriteBatch(array, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_BEFORE, + RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE}, + merge_writer.get()); + if (GetParam()) { + ASSERT_OK(merge_writer->FlushMemory()); + } + + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, + merge_writer->PrepareCommit(/*wait_compaction=*/false)); + ASSERT_OK(merge_writer->Close()); + + const DataIncrement& data_increment = commit_increment.GetNewFilesIncrement(); + ASSERT_EQ(1, data_increment.NewFiles().size()); + ASSERT_EQ(1, data_increment.ChangelogFiles().size()); + ASSERT_EQ("data-" + uuid + "-1.orc", data_increment.NewFiles()[0]->file_name); + ASSERT_EQ("changes-" + uuid + "-0.parquet", data_increment.ChangelogFiles()[0]->file_name); + ASSERT_EQ(2, data_increment.NewFiles()[0]->row_count); + ASSERT_EQ(4, data_increment.ChangelogFiles()[0]->row_count); + + std::shared_ptr expected_data; + auto data_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [2, 2, "Alice", 11, 0, 11.0], + [3, 3, "Bob", 21, 0, 21.0] + ])"}, + &expected_data); + ASSERT_TRUE(data_status.ok()); + CheckFileContent(dir->Str() + "/" + data_increment.NewFiles()[0]->file_name, expected_data); + + std::shared_ptr expected_changelog; + auto changelog_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [1, 1, "Alice", 10, 0, 10.0], + [2, 2, "Alice", 11, 0, 11.0], + [0, 0, "Bob", 20, 0, 20.0], + [3, 3, "Bob", 21, 0, 21.0] + ])"}, + &expected_changelog); + ASSERT_TRUE(changelog_status.ok()); + CheckFileContent(dir->Str() + "/" + data_increment.ChangelogFiles()[0]->file_name, + expected_changelog, "parquet"); +} + +TEST_P(MergeTreeWriterTest, TestInputChangelogIgnoresTargetFileRowNum) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::TARGET_FILE_ROW_NUM, "1"}, + {Options::WRITE_BATCH_SIZE, "1"}})); + + 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=*/-1, dir->Str(), path_factory, + /*schema_id=*/1, options)); + + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(value_type_, R"([ + ["Bob", 20, 0, 20.0], + ["Alice", 10, 0, 10.0], + ["Alice", 11, 0, 11.0], + ["Bob", 21, 0, 21.0] + ])") + .ValueOrDie(); + WriteBatch(array, /*row_kinds=*/{}, merge_writer.get()); + if (GetParam()) { + ASSERT_OK(merge_writer->FlushMemory()); + } + + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, + merge_writer->PrepareCommit(/*wait_compaction=*/false)); + ASSERT_OK(merge_writer->Close()); + + const std::vector>& changelog_files = + commit_increment.GetNewFilesIncrement().ChangelogFiles(); + ASSERT_EQ(1, changelog_files.size()); + ASSERT_EQ(4, changelog_files[0]->row_count); +} + +TEST_P(MergeTreeWriterTest, TestInputChangelogWithSharedShredding) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({ + {Options::FILE_FORMAT, "orc"}, + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::CHANGELOG_FILE_PREFIX, "changes-"}, + {Options::CHANGELOG_FILE_FORMAT, "parquet"}, + {"fields.tags.map.storage-layout", "shared-shredding"}, + {"fields.tags.map.shared-shredding.max-columns", "3"}, + {"fields.tags.map.shared-shredding.column-placement-policy", "plain"}, + {Options::WRITE_ONLY, "true"}, + })); + + 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::vector value_fields = { + DataField(0, arrow::field("id", arrow::int32())), + DataField(1, arrow::field("tags", arrow::map(arrow::utf8(), arrow::int64()))), + }; + auto value_schema = DataField::ConvertDataFieldsToArrowSchema(value_fields); + auto value_type = DataField::ConvertDataFieldsToArrowStructType(value_fields); + auto write_schema = SpecialFields::CompleteSequenceAndValueKindField(value_schema); + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({value_fields[0]}, + /*is_ascending_order=*/true)); + ASSERT_OK_AND_ASSIGN( + auto merge_writer, + MergeTreeWriter::Create( + /*last_sequence_number=*/-1, /*trimmed_primary_keys=*/{"id"}, path_factory, + key_comparator, + /*user_defined_seq_comparator=*/nullptr, merge_function_wrapper_, /*schema_id=*/5, + value_schema, options, noop_compact_manager_, + GetParam() ? std::make_shared(dir->Str() + "/tmp", file_system_) : nullptr, + /*enable_multi_thread_spill=*/false, pool_)); + + auto array = arrow::ipc::internal::json::ArrayFromJSON(value_type, R"([ + [1, [["a", 10], ["b", 20]]], + [2, [["c", 30]]], + [1, [["a", 11], ["c", 31]]] + ])") + .ValueOrDie(); + WriteBatch(array, /*row_kinds=*/{}, merge_writer.get()); + if (GetParam()) { + ASSERT_OK(merge_writer->FlushMemory()); + } + + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, + merge_writer->PrepareCommit(/*wait_compaction=*/false)); + ASSERT_OK(merge_writer->Close()); + + const std::vector>& changelog_files = + commit_increment.GetNewFilesIncrement().ChangelogFiles(); + ASSERT_EQ(1, changelog_files.size()); + ASSERT_EQ(3, changelog_files[0]->row_count); + ASSERT_TRUE(StringUtils::EndsWith(changelog_files[0]->file_name, ".parquet")); + + std::map column_to_k = {{"tags", 3}}; + ASSERT_OK_AND_ASSIGN(auto physical_schema, MapSharedShreddingUtils::LogicalToPhysicalSchema( + write_schema, column_to_k)); + MapSharedShreddingFieldMeta expected_shredding_meta; + expected_shredding_meta.name_to_id = {{"a", 0}, {"b", 1}, {"c", 2}}; + expected_shredding_meta.field_to_columns = {{0, {0}}, {1, {1}}, {2, {0, 1}}}; + expected_shredding_meta.num_columns = 3; + expected_shredding_meta.max_row_width = 2; + CheckShreddingFileSchema(path_factory->ToPath(changelog_files[0]->file_name), physical_schema, + /*field_index=*/3, expected_shredding_meta, "parquet"); +} + TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); @@ -370,12 +603,177 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); } +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->WriteSortedReadersToFiles(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->WriteSortedReadersToFiles(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, 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->WriteSortedReadersToFiles(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReadersToFiles(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()); + bool first_mixed_reader_closed = false; + bool second_mixed_reader_closed = false; + std::vector> mixed_readers; + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &first_mixed_reader_closed)); + mixed_readers.push_back(nullptr); + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &second_mixed_reader_closed)); + Status mixed_status = merge_writer->WriteSortedReadersToFiles(std::move(mixed_readers)); + ASSERT_TRUE(mixed_status.IsInvalid()); + ASSERT_TRUE(first_mixed_reader_closed); + ASSERT_TRUE(second_mixed_reader_closed); + + 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->WriteSortedReadersToFiles(std::move(failing_readers)); + ASSERT_EQ(expected_status, failing_status); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); +} + TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({ @@ -478,7 +876,8 @@ TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { /*creation_time=*/actual_meta->creation_time, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_TRUE(expected_data_file_meta->TEST_Equal(*actual_meta)); } @@ -658,7 +1057,7 @@ TEST_P(MergeTreeWriterTest, TestWriteWithDeleteRow) { /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); @@ -769,7 +1168,7 @@ TEST_P(MergeTreeWriterTest, TestMultiplePrepareCommit) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto expected_data_file_meta2 = std::make_shared( expected_data_file_name2, /*file_size=*/data_file_status2.GetLen(), /*row_count=*/3, @@ -788,7 +1187,7 @@ TEST_P(MergeTreeWriterTest, TestMultiplePrepareCommit) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment1({expected_data_file_meta1}, /*deleted_files=*/{}, /*changelog_files=*/{}); @@ -991,7 +1390,7 @@ TEST_P(MergeTreeWriterTest, TestAutoFlush) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto expected_data_file_meta2 = std::make_shared( expected_data_file_name2, /*file_size=*/data_file_status2.GetLen(), /*row_count=*/3, @@ -1010,7 +1409,7 @@ TEST_P(MergeTreeWriterTest, TestAutoFlush) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta1, expected_data_file_meta2}, /*deleted_files=*/{}, /*changelog_files=*/{}); @@ -1124,7 +1523,7 @@ TEST_P(MergeTreeWriterTest, TestBulkData) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ASSERT_EQ(*commit_increment.GetNewFilesIncrement().NewFiles()[i], *expected_data_file_meta); } } @@ -1193,6 +1592,31 @@ TEST_P(MergeTreeWriterTest, TestUpdateCompactResultDeleteIntermediateFile) { ASSERT_EQ(merge_writer->compact_after_, std::vector>({file_y})); } +TEST_P(MergeTreeWriterTest, TestUpdateCompactResultPropagatesChangelog) { + 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)); + + auto fake_compact_manager = std::make_shared(); + ASSERT_OK_AND_ASSIGN( + auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/-1, dir->Str(), path_factory, /*schema_id=*/0, + options, /*user_defined_seq_comparator=*/nullptr, fake_compact_manager)); + + auto changelog = CreateMeta("changelog", /*level=*/0); + auto compact_result = std::make_shared( + std::vector>(), std::vector>(), + std::vector>({changelog})); + ASSERT_OK(merge_writer->UpdateCompactResult(compact_result)); + + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->DrainIncrement()); + ASSERT_EQ(commit_increment.GetCompactIncrement().ChangelogFiles(), + std::vector>({changelog})); +} + TEST_P(MergeTreeWriterTest, TestUpdateCompactResultWithFileInCompactAfter) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); diff --git a/src/paimon/core/mergetree/sorted_run_test.cpp b/src/paimon/core/mergetree/sorted_run_test.cpp index 7a19fac36..fc4dd08d6 100644 --- a/src/paimon/core/mergetree/sorted_run_test.cpp +++ b/src/paimon/core/mergetree/sorted_run_test.cpp @@ -53,7 +53,7 @@ class SortedRunTest : public testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } }; diff --git a/src/paimon/core/mergetree/write_buffer.cpp b/src/paimon/core/mergetree/write_buffer.cpp index 549975a33..d6c09a1cb 100644 --- a/src/paimon/core/mergetree/write_buffer.cpp +++ b/src/paimon/core/mergetree/write_buffer.cpp @@ -82,6 +82,10 @@ Result>> WriteBuffer::CreateRe return merged_readers; } +Result>> WriteBuffer::CreateRawReaders() { + return sort_buffer_->CreateReaders(); +} + Result WriteBuffer::FlushMemory() { return sort_buffer_->FlushMemory(); } diff --git a/src/paimon/core/mergetree/write_buffer.h b/src/paimon/core/mergetree/write_buffer.h index f4f231000..1c71f7faf 100644 --- a/src/paimon/core/mergetree/write_buffer.h +++ b/src/paimon/core/mergetree/write_buffer.h @@ -72,6 +72,10 @@ class WriteBuffer { /// @return list of KeyValueRecordReaders built from buffered data Result>> CreateReaders(); + /// Create KeyValueRecordReaders containing the raw input records without merging duplicate + /// keys. The caller should invoke Clear() after consuming the readers. + Result>> CreateRawReaders(); + /// Try to spill current buffered data. Return false when the call completed normally but the /// caller should fall back to FlushWriteBuffer before buffering more data. Result FlushMemory(); diff --git a/src/paimon/core/migrate/file_meta_utils_test.cpp b/src/paimon/core/migrate/file_meta_utils_test.cpp index ee30680d8..caca9d2fc 100644 --- a/src/paimon/core/migrate/file_meta_utils_test.cpp +++ b/src/paimon/core/migrate/file_meta_utils_test.cpp @@ -132,7 +132,7 @@ TEST_F(FileMetaUtilsTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-db2b44c0-0d73-449d-82a0-4075bd2cb6e3-0.orc", /*file_size=*/541, @@ -148,7 +148,7 @@ TEST_F(FileMetaUtilsTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); CommitMessageImpl expected(BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/-1, DataIncrement({file_meta1, file_meta2}, {}, {}), CompactIncrement({}, {}, {})); @@ -211,7 +211,7 @@ TEST_F(FileMetaUtilsTest, TestFailover) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-db2b44c0-0d73-449d-82a0-4075bd2cb6e3-0.orc", /*file_size=*/541, @@ -227,7 +227,7 @@ TEST_F(FileMetaUtilsTest, TestFailover) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); CommitMessageImpl expected(BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/-1, DataIncrement({file_meta1, file_meta2}, {}, {}), CompactIncrement({}, {}, {})); @@ -291,7 +291,7 @@ TEST_F(FileMetaUtilsTest, TestWithPartition) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1c547e5f-48b2-4917-a996-71d306377661-0.orc", /*file_size=*/589, @@ -307,7 +307,7 @@ TEST_F(FileMetaUtilsTest, TestWithPartition) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); CommitMessageImpl expected(BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()), /*bucket=*/0, /*total_buckets=*/-1, DataIncrement({file_meta1, file_meta2}, {}, {}), @@ -370,7 +370,7 @@ TEST_F(FileMetaUtilsTest, TestWithNestedType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); CommitMessageImpl expected(BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/-1, DataIncrement({file_meta1}, {}, {}), CompactIncrement({}, {}, {})); diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 2a3d9e10d..5beee34ba 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -28,12 +28,12 @@ #include "fmt/format.h" #include "paimon/common/data/blob_defs.h" #include "paimon/common/data/blob_utils.h" -#include "paimon/common/data/shredding/map_shared_shredding_file_reader.h" +#include "paimon/common/data/shredding/map_shared_shredding_read_plan_factory.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/data/shredding/shredding_file_reader.h" #include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" -#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/predicate_batch_reader.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/table/special_fields.h" @@ -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" @@ -96,7 +97,7 @@ Result>> AbstractSplitRead::CreateR PrepareReaderBuilder(data_file_identifier, extra_format_options)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr file_reader, - CreateFieldMappingReader(data_file_path, file, partition, reader_builder.get(), + CreateFieldMappingReader(data_file_path, file, partition, std::move(reader_builder), field_mapping_builder.get(), dv_factory, row_ranges, data_file_path_factory)); if (file_reader) { @@ -139,23 +140,35 @@ 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; } Result> AbstractSplitRead::CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, - int64_t data_file_size, const ReaderBuilder* reader_builder) const { + int64_t data_file_size, std::unique_ptr reader_builder) const { + if (context_->EnableLateMaterializing()) { + reader_builder = + std::make_unique(std::move(reader_builder), pool_); + } + // TODO(xinyu.lxy): test format table for mosaic format if (context_->EnablePrefetch() && file_format_identifier != "blob" && - file_format_identifier != "avro") { + file_format_identifier != "avro" && file_format_identifier != "mosaic") { PAIMON_ASSIGN_OR_RAISE( std::unique_ptr prefetch_reader, PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, reader_builder, options_.GetFileSystem(), + data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, - /*initialize_read_ranges=*/false, context_->GetPrefetchCacheMode(), - context_->GetCacheConfig(), pool_)); + /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), + context_->GetCacheConfig(), options_.PrefetchIoMetricsEnabled(), pool_)); return std::make_unique(std::move(prefetch_reader)); } else { PAIMON_ASSIGN_OR_RAISE( @@ -167,7 +180,7 @@ Result> AbstractSplitRead::CreateFileBatchReade Result> AbstractSplitRead::CreateFieldMappingReader( const std::string& data_file_path, const std::shared_ptr& file_meta, - const BinaryRow& partition, const ReaderBuilder* reader_builder, + const BinaryRow& partition, std::unique_ptr reader_builder, const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const { @@ -207,16 +220,17 @@ Result> AbstractSplitRead::CreateFieldMappingRe PAIMON_ASSIGN_OR_RAISE(std::string file_format_identifier, file_meta->FileFormat()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_reader, CreateFileBatchReader(file_format_identifier, data_file_path, - file_meta->file_size, reader_builder)); + file_meta->file_size, std::move(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; - PAIMON_ASSIGN_OR_RAISE(shared_shredding_result, ApplySharedShreddingReaderIfNeeded( - std::move(file_reader), read_schema)); - file_reader = std::move(shared_shredding_result.first); - skip_map_selected_keys_filter_field_ids = std::move(shared_shredding_result.second); - PAIMON_ASSIGN_OR_RAISE( - file_reader, ApplyVariantShreddingReaderIfNeeded(std::move(file_reader), read_schema)); + std::pair, std::set> shredding_result; + PAIMON_ASSIGN_OR_RAISE(shredding_result, + ApplyShreddingReaderIfNeeded(std::move(file_reader), read_schema)); + file_reader = std::move(shredding_result.first); + skip_map_selected_keys_filter_field_ids = std::move(shredding_result.second); } if (NeedCompleteRowTrackingFields(options_.RowTrackingEnabled(), read_schema)) { // A blob file has no self-describing schema: its physical fields are declared by the @@ -249,15 +263,16 @@ Result> AbstractSplitRead::CreateFieldMappingRe } Result, std::set>> -AbstractSplitRead::ApplySharedShreddingReaderIfNeeded( +AbstractSplitRead::ApplyShreddingReaderIfNeeded( std::unique_ptr&& file_reader, const std::shared_ptr& read_schema) const { - std::set handled_shared_shredding_field_ids; PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, file_reader->GetFileSchema()); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_arrow_schema, arrow::ImportSchema(file_schema.get())); - std::map> field_read_plans; + + std::set handled_shared_shredding_field_ids; + std::map> plans; for (const auto& read_field : read_schema->fields()) { const auto& field_name = read_field->name(); auto file_field = file_arrow_schema->GetFieldByName(field_name); @@ -275,64 +290,46 @@ AbstractSplitRead::ApplySharedShreddingReaderIfNeeded( continue; } - std::unique_ptr field_read_plan; + std::shared_ptr plan; if (is_shared_shredding_map_access) { if (is_shared_shredding_file) { PAIMON_ASSIGN_OR_RAISE(MapSharedShreddingFieldMeta meta, MapSharedShreddingUtils::DeserializeMetadata(metadata)); PAIMON_ASSIGN_OR_RAISE( - field_read_plan, - MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan(read_field, meta)); + plan, MapSharedShreddingReadPlanFactory::CreateSharedSelectedKeysReadPlan( + read_field, meta)); } else { - PAIMON_ASSIGN_OR_RAISE(field_read_plan, - MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan( - file_field, read_field)); + PAIMON_ASSIGN_OR_RAISE( + plan, MapSharedShreddingReadPlanFactory::CreateDefaultSelectedKeysReadPlan( + file_field, read_field)); } } else { PAIMON_ASSIGN_OR_RAISE(MapSharedShreddingFieldMeta meta, MapSharedShreddingUtils::DeserializeMetadata(metadata)); - PAIMON_ASSIGN_OR_RAISE(field_read_plan, - MapFieldReadPlanFactory::CreateMapReadPlan(read_field, meta)); + PAIMON_ASSIGN_OR_RAISE( + plan, MapSharedShreddingReadPlanFactory::CreateMapReadPlan(read_field, meta)); } - field_read_plans.emplace(field_name, std::move(field_read_plan)); + plans.emplace(field_name, std::move(plan)); PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(read_field)); handled_shared_shredding_field_ids.insert(field_id); } - if (!field_read_plans.empty()) { - file_reader = std::make_unique( - std::move(file_reader), std::move(field_read_plans), pool_); - } - return std::make_pair(std::move(file_reader), std::move(handled_shared_shredding_field_ids)); -} -Result> AbstractSplitRead::ApplyVariantShreddingReaderIfNeeded( - std::unique_ptr&& file_reader, - const std::shared_ptr& read_schema) const { - bool has_variant_field = false; - for (const auto& read_field : read_schema->fields()) { - // Variant columns may be nested inside struct columns; a variant-access projection also - // matches because it carries the variant extension marker itself. - if (VariantTypeUtils::ContainsVariantField(read_field)) { - has_variant_field = true; - break; + std::map> variant_plans; + PAIMON_ASSIGN_OR_RAISE(variant_plans, VariantShreddingReadPlanFactory::CreateReadPlans( + read_schema, file_arrow_schema, pool_)); + for (auto& [field_name, plan] : variant_plans) { + if (!plans.emplace(field_name, std::move(plan)).second) { + return Status::Invalid( + fmt::format("multiple shredding read plans exist for field {}", field_name)); } } - if (!has_variant_field) { - return std::move(file_reader); - } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, - file_reader->GetFileSchema()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_arrow_schema, - arrow::ImportSchema(file_schema.get())); - std::map> plans; - PAIMON_ASSIGN_OR_RAISE(plans, VariantShreddingReadPlanFactory::CreateReadPlans( - read_schema, file_arrow_schema, pool_)); + if (!plans.empty()) { file_reader = std::make_unique(std::move(file_reader), std::move(plans), pool_); } - return std::move(file_reader); + return std::make_pair(std::move(file_reader), std::move(handled_shared_shredding_field_ids)); } Result> AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index 27349fec1..a02ed5fb2 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -107,26 +107,22 @@ class AbstractSplitRead : public SplitRead { Result> CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, - int64_t data_file_size, const ReaderBuilder* reader_builder) const; + int64_t data_file_size, std::unique_ptr reader_builder) const; // return nullptr if data file is skipped by index or dv Result> CreateFieldMappingReader( const std::string& data_file_path, const std::shared_ptr& file_meta, - const BinaryRow& partition, const ReaderBuilder* reader_builder, + const BinaryRow& partition, std::unique_ptr reader_builder, const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const; + /// The returned field ID set contains MAP fields handled by shredding read plans. It tells + /// FieldMappingReader to skip its generic selected-key filtering because the plans have already + /// applied any requested key selection. Result, std::set>> - ApplySharedShreddingReaderIfNeeded(std::unique_ptr&& file_reader, - const std::shared_ptr& read_schema) const; - - /// Wraps the reader with a `ShreddingFileReader` when any read variant column needs - /// reassembly or path extraction; a plain read of an unshredded variant column is passed - /// through untouched. - Result> ApplyVariantShreddingReaderIfNeeded( - std::unique_ptr&& file_reader, - const std::shared_ptr& read_schema) const; + ApplyShreddingReaderIfNeeded(std::unique_ptr&& file_reader, + const std::shared_ptr& read_schema) const; static bool NeedCompleteRowTrackingFields(bool row_tracking_enabled, const std::shared_ptr& read_schema); 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 f319498a4..d72dcaf0d 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,8 +29,8 @@ #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" #include "paimon/core/schema/table_schema.h" #include "paimon/core/stats/simple_stats_evolution.h" @@ -44,6 +44,7 @@ #include "paimon/predicate/predicate_builder.h" #include "paimon/scan_context.h" #include "paimon/status.h" +#include "paimon/table/source/scan_metrics.h" #include "paimon/table/source/table_scan.h" #include "paimon/testing/utils/testharness.h" #include "paimon/testing/utils/timezone_guard.h" @@ -172,7 +173,20 @@ TEST(AppendOnlyFileStoreScanTest, TestScanDurationMetric) { metrics->GetCounter(ScanMetrics::LAST_SCAN_DURATION)); ASSERT_OK_AND_ASSIGN(HistogramStats stats, metrics->GetHistogramStats(ScanMetrics::SCAN_DURATION)); + ASSERT_OK_AND_ASSIGN(uint64_t manifest_read_duration, + metrics->GetCounter(ScanMetrics::LAST_MANIFEST_READ_DURATION)); + ASSERT_OK_AND_ASSIGN(HistogramStats manifest_stats, + metrics->GetHistogramStats(ScanMetrics::MANIFEST_READ_DURATION)); + ASSERT_OK_AND_ASSIGN(uint64_t scanned_rows, + metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_SCANNED_ROWS)); + ASSERT_OK_AND_ASSIGN(uint64_t materialized_rows, + metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_MATERIALIZED_ROWS)); ASSERT_EQ(stats.count, kPlanCount); + ASSERT_EQ(manifest_stats.count, kPlanCount); + ASSERT_LE(manifest_stats.min, static_cast(manifest_read_duration)); + ASSERT_LE(static_cast(manifest_read_duration), manifest_stats.max); + ASSERT_GT(scanned_rows, 0); + ASSERT_GE(scanned_rows, materialized_rows); ASSERT_LE(stats.min, stats.max); ASSERT_LE(stats.min, static_cast(last_scan_duration)); ASSERT_LE(static_cast(last_scan_duration), stats.max); @@ -186,11 +200,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 +222,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 +280,41 @@ 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()); + std::shared_ptr first_metrics = scan_first->GetScanMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t first_cache_enabled, + first_metrics->GetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_ENABLED)); + ASSERT_OK_AND_ASSIGN(uint64_t first_cache_hit, + first_metrics->GetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_HIT)); + ASSERT_OK_AND_ASSIGN(uint64_t first_cache_misses, + first_metrics->GetCounter(ScanMetrics::SNAPSHOT_CACHE_MISSES)); + ASSERT_OK(first_metrics->GetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_LOAD_DURATION)); + ASSERT_OK(first_metrics->GetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_STORE_DURATION)); + ASSERT_OK(first_metrics->GetHistogramStats(ScanMetrics::SNAPSHOT_CACHE_LOAD_DURATION)); + ASSERT_OK(first_metrics->GetHistogramStats(ScanMetrics::SNAPSHOT_CACHE_STORE_DURATION)); + ASSERT_EQ(first_cache_enabled, 1); + ASSERT_EQ(first_cache_hit, 0); + ASSERT_EQ(first_cache_misses, 1); // 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())); + std::shared_ptr second_metrics = scan_second->GetScanMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t second_cache_hit, + second_metrics->GetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_HIT)); + ASSERT_OK_AND_ASSIGN(uint64_t second_cache_hits, + second_metrics->GetCounter(ScanMetrics::SNAPSHOT_CACHE_HITS)); + ASSERT_OK_AND_ASSIGN(uint64_t scanned_rows, + second_metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_SCANNED_ROWS)); + ASSERT_OK_AND_ASSIGN( + uint64_t materialized_rows, + second_metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_MATERIALIZED_ROWS)); + ASSERT_EQ(second_cache_hit, 1); + ASSERT_EQ(second_cache_hits, 1); + ASSERT_GE(scanned_rows, materialized_rows); + ASSERT_OK(second_metrics->GetHistogramStats(ScanMetrics::SNAPSHOT_CACHE_LOAD_DURATION)); } TEST(AppendOnlyFileStoreScanTest, TestSnapshotLiveManifestCacheRebuildOnMiss) { @@ -285,6 +340,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 b2946e83c..f660093a2 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( @@ -290,14 +290,8 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); - // TODO(xinyu.lxy): temporarily disabled pre-buffer for parquet, which may cause high - // memory usage during compaction. Will fix via parquet format refactor. - auto new_options = options; - if (new_options.find("parquet.read.enable-pre-buffer") == new_options.end()) { - new_options["parquet.read.enable-pre-buffer"] = "false"; - } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_read_context, - InternalReadContext::Create(read_context, table_schema_, new_options)); + InternalReadContext::Create(read_context, table_schema_, options)); auto read = std::make_unique(file_store_path_factory_, internal_read_context, pool_, compact_executor_); 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 26411c198..4e1cef501 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_changes_provider_test.cpp b/src/paimon/core/operation/commit/commit_changes_provider_test.cpp index edbb1a48b..f47e11262 100644 --- a/src/paimon/core/operation/commit/commit_changes_provider_test.cpp +++ b/src/paimon/core/operation/commit/commit_changes_provider_test.cpp @@ -63,7 +63,7 @@ ManifestEntry CreateManifestEntry(const std::string& file_name, const FileKind& /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, CreateIntRow(partition_value), /*bucket=*/0, /*total_buckets=*/1, file_meta); diff --git a/src/paimon/core/operation/commit/commit_scanner.cpp b/src/paimon/core/operation/commit/commit_scanner.cpp index 5590dd02f..448bd7e47 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/conflict_detection_test.cpp b/src/paimon/core/operation/commit/conflict_detection_test.cpp index 66d57894b..3fb751acf 100644 --- a/src/paimon/core/operation/commit/conflict_detection_test.cpp +++ b/src/paimon/core/operation/commit/conflict_detection_test.cpp @@ -175,7 +175,7 @@ class ConflictDetectionTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, partition, bucket, total_buckets, data_file_meta); } @@ -194,7 +194,7 @@ class ConflictDetectionTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, first_row_id, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, partition, bucket, /*total_buckets=*/2, data_file_meta); } diff --git a/src/paimon/core/operation/commit/manifest_entry_changes_test.cpp b/src/paimon/core/operation/commit/manifest_entry_changes_test.cpp index 8680ff2c3..5eb1608c4 100644 --- a/src/paimon/core/operation/commit/manifest_entry_changes_test.cpp +++ b/src/paimon/core/operation/commit/manifest_entry_changes_test.cpp @@ -67,7 +67,7 @@ class ManifestEntryChangesTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateIndexFileMeta( @@ -140,7 +140,7 @@ TEST_F(ManifestEntryChangesTest, TestDropStatsOnlyForDeleteEntries) { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::vector({"f0"}), /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ASSERT_OK_AND_ASSIGN(std::shared_ptr after, before->Upgrade(/*new_level=*/1)); CompactIncrement compact_increment(/*compact_before=*/{before}, /*compact_after=*/{after}, diff --git a/src/paimon/core/operation/commit/overwrite_changes_provider_test.cpp b/src/paimon/core/operation/commit/overwrite_changes_provider_test.cpp index 8130fe65c..3552b6da1 100644 --- a/src/paimon/core/operation/commit/overwrite_changes_provider_test.cpp +++ b/src/paimon/core/operation/commit/overwrite_changes_provider_test.cpp @@ -64,7 +64,7 @@ ManifestEntry CreateManifestEntry(const std::string& file_name, const FileKind& /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, CreateIntRow(partition_value), /*bucket=*/0, /*total_buckets=*/1, file_meta); diff --git a/src/paimon/core/operation/commit/realtime_commit_properties.cpp b/src/paimon/core/operation/commit/realtime_commit_properties.cpp index 262fbbeff..c79fc968e 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 31a014b4b..4c4069200 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 68b159451..afe529e26 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/commit/row_id_column_conflict_checker_test.cpp b/src/paimon/core/operation/commit/row_id_column_conflict_checker_test.cpp index e9091b0fc..907aea44c 100644 --- a/src/paimon/core/operation/commit/row_id_column_conflict_checker_test.cpp +++ b/src/paimon/core/operation/commit/row_id_column_conflict_checker_test.cpp @@ -57,7 +57,8 @@ class RowIdColumnConflictCheckerTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/first_row_id, write_cols); + /*first_row_id=*/first_row_id, write_cols, + /*column_max_sequence_numbers=*/std::nullopt); } Result> CreateChecker( diff --git a/src/paimon/core/operation/commit/row_tracking_commit_utils_test.cpp b/src/paimon/core/operation/commit/row_tracking_commit_utils_test.cpp index 02e79fc56..7d2b0fb02 100644 --- a/src/paimon/core/operation/commit/row_tracking_commit_utils_test.cpp +++ b/src/paimon/core/operation/commit/row_tracking_commit_utils_test.cpp @@ -61,7 +61,8 @@ class RowTrackingCommitUtilsTest : public testing::Test { /*embedded_index=*/nullptr, file_source, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, write_cols); + /*first_row_id=*/std::nullopt, write_cols, + /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(FileKind::Add(), CreateIntRow(1), /*bucket=*/0, /*total_buckets=*/1, file_meta); } @@ -81,7 +82,8 @@ class RowTrackingCommitUtilsTest : public testing::Test { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, file_source, /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, first_row_id, write_cols); + /*external_path=*/std::nullopt, first_row_id, write_cols, + /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(FileKind::Add(), CreateIntRow(1), /*bucket=*/0, /*total_buckets=*/1, file_meta); } diff --git a/src/paimon/core/operation/commit/sequence_snapshot_properties.cpp b/src/paimon/core/operation/commit/sequence_snapshot_properties.cpp index 51729b991..512869f00 100644 --- a/src/paimon/core/operation/commit/sequence_snapshot_properties.cpp +++ b/src/paimon/core/operation/commit/sequence_snapshot_properties.cpp @@ -24,6 +24,7 @@ #include #include +#include "paimon/common/utils/string_utils.h" #include "paimon/core/manifest/file_kind.h" namespace paimon { @@ -40,19 +41,12 @@ Result> SequenceSnapshotProperties::MaxSequenceNumber( return std::optional(); } - try { - size_t parsed = 0; - int64_t value = std::stoll(iter->second, &parsed); - if (parsed != iter->second.size()) { - return Status::Invalid( - fmt::format("Invalid {} value '{}': trailing characters are not allowed", - kMaxSequenceNumberKey, iter->second)); - } - return std::optional(value); - } catch (const std::exception& e) { - return Status::Invalid(fmt::format("Invalid {} value '{}': {}", kMaxSequenceNumberKey, - iter->second, e.what())); + std::optional value = StringUtils::StringToValue(iter->second); + if (!value) { + return Status::Invalid( + fmt::format("Invalid {} value '{}'", kMaxSequenceNumberKey, iter->second)); } + return value; } std::optional SequenceSnapshotProperties::MaxSequenceNumberFromFiles( diff --git a/src/paimon/core/operation/commit/sequence_snapshot_properties_test.cpp b/src/paimon/core/operation/commit/sequence_snapshot_properties_test.cpp index af572b720..ca962e0e8 100644 --- a/src/paimon/core/operation/commit/sequence_snapshot_properties_test.cpp +++ b/src/paimon/core/operation/commit/sequence_snapshot_properties_test.cpp @@ -74,7 +74,7 @@ class SequenceSnapshotPropertiesTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } ManifestEntry CreateEntry(const FileKind& kind, int64_t max_sequence_number) const { @@ -115,7 +115,7 @@ TEST_F(SequenceSnapshotPropertiesTest, MaxSequenceNumberTrailingCharacters) { std::map properties{ {SequenceSnapshotProperties::kMaxSequenceNumberKey, "123abc"}}; ASSERT_NOK_WITH_MSG(SequenceSnapshotProperties::MaxSequenceNumber(MakeSnapshot(properties)), - "trailing characters are not allowed"); + "Invalid sequence.generation.max-sequence-number value '123abc'"); } TEST_F(SequenceSnapshotPropertiesTest, MaxSequenceNumberNotANumber) { diff --git a/src/paimon/core/operation/commit_metrics_test.cpp b/src/paimon/core/operation/commit_metrics_test.cpp index 240d91345..9d1cf9cf2 100644 --- a/src/paimon/core/operation/commit_metrics_test.cpp +++ b/src/paimon/core/operation/commit_metrics_test.cpp @@ -62,7 +62,7 @@ ManifestEntry CreateEntry(const FileKind& kind, int32_t partition, int32_t bucke /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, part, bucket, /*total_buckets=*/10, file_meta); } diff --git a/src/paimon/core/operation/data_evolution_file_store_scan_test.cpp b/src/paimon/core/operation/data_evolution_file_store_scan_test.cpp index b5581f85c..7d042814b 100644 --- a/src/paimon/core/operation/data_evolution_file_store_scan_test.cpp +++ b/src/paimon/core/operation/data_evolution_file_store_scan_test.cpp @@ -554,7 +554,8 @@ TEST_F(DataEvolutionFileStoreScanTest, TestFilterEntryByRowRanges) { /*creation_time=*/Timestamp(1737111915429ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/100, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); ManifestEntry entry(FileKind::Add(), BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/1, file); { @@ -574,7 +575,7 @@ TEST_F(DataEvolutionFileStoreScanTest, TestFilterEntryByRowRanges) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); ManifestEntry entry_without_first_row_id(FileKind::Add(), BinaryRow::EmptyRow(), /*bucket=*/0, /*total_buckets=*/1, file_without_first_row_id); diff --git a/src/paimon/core/operation/data_evolution_split_read.cpp b/src/paimon/core/operation/data_evolution_split_read.cpp index 52a1bb51c..9e71fe47d 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 983ca29d7..2568eae46 100644 --- a/src/paimon/core/operation/data_evolution_split_read.h +++ b/src/paimon/core/operation/data_evolution_split_read.h @@ -64,13 +64,14 @@ struct DeletionFile; /// ->(ConcatBatchReader across blob files | BlobFallbackBatchReader across blob sequence layers) /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ApplyBitmapIndexBatchReader) /// ->(CompleteRowTrackingFieldsBatchReader)->(ShreddingFileReader) -/// ->(MapSharedShreddingFileReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(VectorFileBatchReader)->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->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 809933e9f..03bc95b6f 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 d6f4b1212..a5bd25053 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 75b5fff59..238f599ed 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 d6b965ba5..ffdd30feb 100644 --- a/src/paimon/core/operation/expire_snapshots_test.cpp +++ b/src/paimon/core/operation/expire_snapshots_test.cpp @@ -148,7 +148,7 @@ class ExpireSnapshotsTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, row, bucket, /*total_buckets=*/3, data_file_meta); } @@ -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 b2ba17fa0..ad0942ade 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 5388b9e85..8f2b3c732 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 7c8044148..ca4de22f9 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_commit_impl_test.cpp b/src/paimon/core/operation/file_store_commit_impl_test.cpp index f32d86fac..b1b8677f8 100644 --- a/src/paimon/core/operation/file_store_commit_impl_test.cpp +++ b/src/paimon/core/operation/file_store_commit_impl_test.cpp @@ -198,7 +198,7 @@ class FileStoreCommitImplTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, partition, bucket, total_buckets, data_file_meta); } @@ -223,7 +223,7 @@ class FileStoreCommitImplTest : public testing::Test { /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, BinaryRow::EmptyRow(), 0, 2, data_file_meta); } @@ -278,7 +278,7 @@ class FileStoreCommitImplTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateAppendDataFileMeta(const std::string& file_name, @@ -294,7 +294,7 @@ class FileStoreCommitImplTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } bool IsStringInSet(const std::set& strSet, const std::string& target) { diff --git a/src/paimon/core/operation/file_store_scan.cpp b/src/paimon/core/operation/file_store_scan.cpp index 865e006ff..befe7eb97 100644 --- a/src/paimon/core/operation/file_store_scan.cpp +++ b/src/paimon/core/operation/file_store_scan.cpp @@ -43,7 +43,6 @@ #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/manifest/snapshot_live_manifest_entries.h" -#include "paimon/core/operation/metrics/scan_metrics.h" #include "paimon/core/partition/partition_info.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/stats/simple_stats_evolution.h" @@ -58,6 +57,7 @@ #include "paimon/predicate/predicate_builder.h" #include "paimon/predicate/predicate_utils.h" #include "paimon/scan_context.h" +#include "paimon/table/source/scan_metrics.h" namespace paimon { enum class FieldType; @@ -136,6 +136,9 @@ Result> FileStoreScan::ReadPartitionEntries() const Result> FileStoreScan::CreatePlan() const { Duration duration; + Duration manifest_read_duration; + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_LOAD_DURATION, 0); + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_STORE_DURATION, 0); std::optional snapshot; std::vector all_manifest_file_metas; std::vector filtered_manifest_file_metas; @@ -148,9 +151,13 @@ Result> FileStoreScan::CreatePlan() cons core_options_.GetScanManifestEntryCacheMaxSnapshots() > 0 && core_options_.GetCache() != nullptr && !table_path_.empty() && !row_range_index_.has_value() && bucket_filter_.has_value(); + uint64_t lazy_decode_scanned_rows = 0; + bool snapshot_cache_hit = false; if (use_snapshot_live_manifest_cache) { - PAIMON_RETURN_NOT_OK(ReadManifestEntriesWithCache( - snapshot.value(), all_manifest_file_metas, bucket_filter_.value(), &manifest_entries)); + PAIMON_RETURN_NOT_OK(ReadManifestEntriesWithCache(snapshot.value(), all_manifest_file_metas, + bucket_filter_.value(), &manifest_entries, + &snapshot_cache_hit)); + lazy_decode_scanned_rows = manifest_entries.size(); std::vector filtered_entries; filtered_entries.reserve(manifest_entries.size()); for (auto& entry : manifest_entries) { @@ -161,8 +168,15 @@ Result> FileStoreScan::CreatePlan() cons } manifest_entries = std::move(filtered_entries); } else { + lazy_decode_scanned_rows = std::accumulate( + filtered_manifest_file_metas.begin(), filtered_manifest_file_metas.end(), uint64_t{0}, + [](uint64_t sum, const ManifestFileMeta& meta) { + return sum + static_cast(meta.NumAddedFiles() + meta.NumDeletedFiles()); + }); PAIMON_RETURN_NOT_OK(ReadManifestEntries(filtered_manifest_file_metas, &manifest_entries)); } + const uint64_t lazy_decode_materialized_rows = manifest_entries.size(); + const uint64_t manifest_read_duration_ms = manifest_read_duration.Get(); PAIMON_ASSIGN_OR_RAISE(manifest_entries, PostFilterManifestEntries(std::move(manifest_entries))); @@ -208,6 +222,22 @@ Result> FileStoreScan::CreatePlan() cons ScanMetrics::LAST_SCAN_SKIPPED_TABLE_FILES, std::max(int64_t{0}, all_data_files - static_cast(manifest_entries.size()))); metrics_->SetCounter(ScanMetrics::LAST_SCAN_RESULTED_TABLE_FILES, manifest_entries.size()); + metrics_->SetCounter(ScanMetrics::LAST_MANIFEST_READ_DURATION, manifest_read_duration_ms); + metrics_->ObserveHistogram(ScanMetrics::MANIFEST_READ_DURATION, + static_cast(manifest_read_duration_ms)); + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_ENABLED, + use_snapshot_live_manifest_cache ? 1 : 0); + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_HIT, snapshot_cache_hit ? 1 : 0); + Result cache_hits = metrics_->GetCounter(ScanMetrics::SNAPSHOT_CACHE_HITS); + Result cache_misses = metrics_->GetCounter(ScanMetrics::SNAPSHOT_CACHE_MISSES); + metrics_->SetCounter(ScanMetrics::SNAPSHOT_CACHE_HITS, + (cache_hits.ok() ? cache_hits.value() : 0) + (snapshot_cache_hit ? 1 : 0)); + metrics_->SetCounter(ScanMetrics::SNAPSHOT_CACHE_MISSES, + (cache_misses.ok() ? cache_misses.value() : 0) + + (use_snapshot_live_manifest_cache && !snapshot_cache_hit ? 1 : 0)); + metrics_->SetCounter(ScanMetrics::LAST_LAZY_DECODE_SCANNED_ROWS, lazy_decode_scanned_rows); + metrics_->SetCounter(ScanMetrics::LAST_LAZY_DECODE_MATERIALIZED_ROWS, + lazy_decode_materialized_rows); return std::make_shared(scan_mode_, snapshot, std::move(manifest_entries)); } @@ -245,6 +275,8 @@ Status FileStoreScan::ReadManifestsWithSnapshot(const Snapshot& snapshot, return manifest_list_->ReadDataManifests(snapshot, manifests); case ScanMode::DELTA: return manifest_list_->ReadDeltaManifests(snapshot, manifests); + case ScanMode::CHANGELOG: + return manifest_list_->ReadChangelogManifests(snapshot, manifests); default: return Status::NotImplemented("Unknown scan mode ", std::to_string(static_cast(scan_mode_))); @@ -298,15 +330,22 @@ Status FileStoreScan::ReadManifestEntries(const std::vector& m // snapshot's data manifests. Status FileStoreScan::ReadManifestEntriesWithCache( const Snapshot& snapshot, const std::vector& all_manifest_metas, - int32_t bucket, std::vector* manifest_entries) const { + int32_t bucket, std::vector* manifest_entries, bool* cache_hit) const { + Duration cache_load_duration; PAIMON_ASSIGN_OR_RAISE(SnapshotLiveManifestEntries cached_entries, LoadSnapshotLiveManifestEntries(bucket)); + const uint64_t cache_load_duration_ms = cache_load_duration.Get(); + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_LOAD_DURATION, cache_load_duration_ms); + metrics_->ObserveHistogram(ScanMetrics::SNAPSHOT_CACHE_LOAD_DURATION, + static_cast(cache_load_duration_ms)); std::optional cached = cached_entries.LatestBeforeOrEqual(snapshot.Id()); if (cached && cached->snapshot_id == snapshot.Id()) { + *cache_hit = true; *manifest_entries = *cached->entries; return Status::OK(); } + *cache_hit = false; // Rebuild the target snapshot bucket from all manifests and write the live entries back to the // cache. @@ -320,7 +359,12 @@ Status FileStoreScan::ReadManifestEntriesWithCache( ReadAndMergeBucketFileEntries(bucket_manifest_metas, bucket, manifest_entries)); std::vector cache_entries = *manifest_entries; cached_entries.Put(snapshot.Id(), std::move(cache_entries)); + Duration cache_store_duration; PAIMON_RETURN_NOT_OK(StoreSnapshotLiveManifestEntries(bucket, cached_entries)); + const uint64_t cache_store_duration_ms = cache_store_duration.Get(); + metrics_->SetCounter(ScanMetrics::LAST_SNAPSHOT_CACHE_STORE_DURATION, cache_store_duration_ms); + metrics_->ObserveHistogram(ScanMetrics::SNAPSHOT_CACHE_STORE_DURATION, + static_cast(cache_store_duration_ms)); return Status::OK(); } @@ -365,6 +409,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_scan.h b/src/paimon/core/operation/file_store_scan.h index 155f36e74..ff53e0c9e 100644 --- a/src/paimon/core/operation/file_store_scan.h +++ b/src/paimon/core/operation/file_store_scan.h @@ -161,7 +161,9 @@ class FileStoreScan { } std::shared_ptr GetScanMetrics() const { - return metrics_; + auto snapshot = std::make_shared(); + snapshot->Overwrite(metrics_); + return snapshot; } static Result> CreatePartitionPredicate( @@ -260,7 +262,8 @@ class FileStoreScan { Status ReadManifestEntriesWithCache(const Snapshot& snapshot, const std::vector& bucket_manifest_metas, int32_t bucket, - std::vector* manifest_entries) const; + std::vector* manifest_entries, + bool* cache_hit) const; std::shared_ptr SnapshotLiveManifestEntriesCacheKey(int32_t bucket) const; Result LoadSnapshotLiveManifestEntries(int32_t bucket) const; Status StoreSnapshotLiveManifestEntries(int32_t bucket, diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 940608c72..fa1294d1b 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -57,6 +57,27 @@ struct KeyValue; template class MergeFunctionWrapper; +namespace { + +Status RestoreRealtimeCommittedProgress(const std::shared_ptr& realtime_context, + const std::shared_ptr& snapshot_manager, + const CoreOptions& options) { + 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(realtime_context)); + PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( + latest_snapshot->Id(), realtime_committed_offsets)); + } + return Status::OK(); +} + +} // namespace + Result> FileStoreWrite::PrepareCommitWithProgress(int64_t) { return Status::Invalid("prepare commit with progress requires a real-time writer"); } @@ -120,6 +141,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(); @@ -140,17 +164,8 @@ Result> FileStoreWrite::Create(std::unique_ptr 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)); - } + PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(), + snapshot_manager, options)); } std::shared_ptr write_schema = arrow_schema; const auto& write_field_names = ctx->GetWriteSchema(); @@ -194,7 +209,16 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - return Status::Invalid("real-time write currently supports append tables only"); + PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *schema)); + if (ignore_previous_files) { + return Status::NotImplemented( + "PK realtime requires restore from the latest snapshot"); + } + if (!ctx->GetWriteSchema().empty()) { + return Status::NotImplemented("PK realtime does not support a custom write schema"); + } + PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(), + snapshot_manager, options)); } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( @@ -250,7 +274,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 f33b7a359..8e773cdcd 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -74,6 +74,9 @@ class InternalReadContext { bool EnablePrefetch() const { return read_context_->EnablePrefetch(); } + bool EnableLateMaterializing() const { + return read_context_->EnableLateMaterializing(); + } uint32_t GetPrefetchBatchCount() const { return read_context_->GetPrefetchBatchCount(); } @@ -96,8 +99,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/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index 286188570..20838122b 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -20,6 +20,7 @@ #include +#include "arrow/c/bridge.h" #include "arrow/type.h" #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" diff --git a/src/paimon/core/operation/key_value_file_store_scan.cpp b/src/paimon/core/operation/key_value_file_store_scan.cpp index 54197f871..54af98a0c 100644 --- a/src/paimon/core/operation/key_value_file_store_scan.cpp +++ b/src/paimon/core/operation/key_value_file_store_scan.cpp @@ -177,8 +177,13 @@ Result KeyValueFileStoreScan::IsValueFilterEnabled() const { return value_filter_force_enabled_; case ScanMode::DELTA: return false; + case ScanMode::CHANGELOG: { + ChangelogProducer producer = core_options_.GetChangelogProducer(); + return producer == ChangelogProducer::LOOKUP || + producer == ChangelogProducer::FULL_COMPACTION; + } default: - return Status::NotImplemented("only support ALL and DELTA scan mode"); + return Status::NotImplemented("unknown scan mode"); } } diff --git a/src/paimon/core/operation/key_value_file_store_scan_test.cpp b/src/paimon/core/operation/key_value_file_store_scan_test.cpp index b221715f5..f9fd2ea2f 100644 --- a/src/paimon/core/operation/key_value_file_store_scan_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_scan_test.cpp @@ -33,7 +33,6 @@ #include "paimon/core/manifest/file_source.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" -#include "paimon/core/operation/metrics/scan_metrics.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" @@ -50,6 +49,7 @@ #include "paimon/predicate/literal.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/scan_context.h" +#include "paimon/table/source/scan_metrics.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/testharness.h" @@ -353,7 +353,8 @@ TEST_F(KeyValueFileStoreScanTest, TestNoOverlapping) { /*embedded_index=*/nullptr, /*file_source=*/FileSource::Append(), /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt)); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt)); } return entries; }; @@ -405,7 +406,7 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithValueStatsCols) { /*value_stats_cols=*/value_stats_cols, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); // max(v0)=50 > 30.1, should be kept. SimpleStats value_stats_keep = BinaryRowGenerator::GenerateStats( @@ -430,7 +431,7 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithValueStatsCols) { /*value_stats_cols=*/value_stats_cols, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); // max(v0)=20 <= 30.1, should be filtered out. ASSERT_OK_AND_ASSIGN(bool keep, scan->FilterByStats(entry)); @@ -482,7 +483,7 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithSchemaEvolution) { /*value_stats_cols=*/value_stats_cols, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); SimpleStats value_stats_keep = BinaryRowGenerator::GenerateStats( /*min=*/{40}, /*max=*/{50}, /*null=*/{0}, pool.get()); @@ -506,7 +507,7 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithSchemaEvolution) { /*value_stats_cols=*/value_stats_cols, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(bool keep, scan->FilterByStats(entry)); ASSERT_FALSE(keep); @@ -555,7 +556,7 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithNewFieldUsesNullSta /*value_stats_cols=*/value_stats_cols, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(bool keep, scan->FilterByStats(old_schema_entry)); ASSERT_FALSE(keep); 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 08c5ea0c3..e8e8cd3d9 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/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_reader.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,14 +77,19 @@ 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), compact_manager_factory_(std::make_unique( 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")) {} + file_store_path_factory_, root_path_, ignore_previous_files, pool_)), + logger_(Logger::GetLogger("KeyValueFileStoreWrite")) { + if (realtime_context_) { + writer_memory_manager_ = std::make_unique(); + } +} Result> KeyValueFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { @@ -106,22 +121,68 @@ Result> KeyValueFileStoreWrite::CreateWriter( file_store_path_factory_->CreateDataFilePathFactory(partition, bucket)); PAIMON_ASSIGN_OR_RAISE(std::vector trimmed_primary_keys, table_schema_->TrimmedPrimaryKeys()); - 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; + std::shared_ptr transport_schema; + 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_)); + transport_schema = RealtimePrimaryKeyLayout::CreateSchema(schema_->fields()); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*transport_schema, c_write_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(c_write_schema), options_.ToMap(), pool_, + RealtimeStoreMode::PRIMARY_KEY}, + RealtimePartitionBucket(partition_map, bucket))); + realtime_store_state = std::move(store_state); + compact_manager = std::make_shared(); + } else { + 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(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_, transport_schema, + trimmed_primary_keys, key_comparator_, options_, + 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 14457590f..66c362f2e 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 35d938af7..88b848eea 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 @@ -40,10 +44,13 @@ #include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/data/shredding/map_shredding_defs.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/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/realtime/realtime_primary_key_reader.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 +59,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 +69,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(); +}; + +} // namespace class KeyValueFileStoreWriteTest : public ::testing::Test { protected: @@ -127,14 +180,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 +247,67 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { EXPECT_NE(nullptr, metadata); return MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy()).value(); } + + Result>> + ReadRealtimePrimaryKeyTransportRows( + 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"); + } + arrow::FieldVector value_fields = {DataField::ConvertDataFieldToArrowField(DataField( + 0, arrow::field("id", arrow::int64(), false))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}; + std::shared_ptr transport_schema = + RealtimePrimaryKeyLayout::CreateSchema(value_fields); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*transport_schema, c_schema.get())); + RealtimeQueryContext query_context{c_schema.get(), 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 realtime primary-key transport 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 realtime primary-key transport 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 +418,197 @@ 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 RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector transport_rows, + ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{ + {0, 1, "old", 0, 0}, {2, 1, "new", 2, 2}, {3, 2, "two", 1, 1}}), + transport_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, 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 RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector retained_rows, + ReadRealtimePrimaryKeyTransportRows(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, + ReadRealtimePrimaryKeyTransportRows(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 RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector transport_rows, + ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), + transport_rows); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[2, "overflow"]])")), + "real-time offset range exceeds INT64_MAX"); + ASSERT_OK_AND_ASSIGN(transport_rows, ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), + transport_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/manifest_file_merger_test.cpp b/src/paimon/core/operation/manifest_file_merger_test.cpp index bff1659e3..d2617a654 100644 --- a/src/paimon/core/operation/manifest_file_merger_test.cpp +++ b/src/paimon/core/operation/manifest_file_merger_test.cpp @@ -101,7 +101,7 @@ class ManifestFileMergerTest : public testing::Test { nullptr, // not used FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); } ManifestFileMeta MakeManifest(const std::vector& entries) { diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index b753ea431..9b913edfb 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -36,6 +36,7 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/object_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/apply_deletion_vector_batch_reader.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" @@ -78,6 +79,185 @@ struct KeyValue; template class MergeFunctionWrapper; +namespace { + +class SortMergeKeyValueRecordReader : public KeyValueRecordReader { + public: + explicit SortMergeKeyValueRecordReader(std::unique_ptr&& reader) + : reader_(std::move(reader)) {} + + class Iterator : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(std::unique_ptr&& iterator) + : iterator_(std::move(iterator)) {} + + Result HasNext() const override { + return iterator_->HasNext(); + } + + Result Next() override { + return std::move(iterator_->Next()); + } + + private: + std::unique_ptr iterator_; + }; + + Result> NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + reader_->NextBatch()); + if (!iterator) { + return std::unique_ptr(); + } + return std::make_unique(std::move(iterator)); + } + + void Close() override { + reader_->Close(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + private: + std::unique_ptr reader_; +}; + +} // namespace + +class MergeFileSplitRead::RealtimeReaderBuilder { + public: + static Result> Create( + const std::vector>& disk_splits, + std::vector>&& additional_readers, + MergeFileSplitRead* owner) { + ScopeGuard additional_readers_guard([&additional_readers]() { + for (const std::unique_ptr& reader : additional_readers) { + if (reader) { + reader->Close(); + } + } + }); + 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)); + } + additional_readers_guard.Release(); + 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::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) { + return Status::Invalid("merge input disk split is not a data split"); + } + if (!first_split) { + first_split = data_split; + } + 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()); + } + } + 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)); + + DeletionVector::Factory dv_factory; + std::vector> disk_sections; + PAIMON_RETURN_NOT_OK( + owner_->CreateDiskSections(data_files, deletion_files, &dv_factory, &disk_sections)); + if (disk_sections.empty()) { + return Status::OK(); + } + std::vector> section_readers; + ScopeGuard section_readers_guard([§ion_readers]() { + for (const std::unique_ptr& reader : section_readers) { + reader->Close(); + } + }); + section_readers.reserve(disk_sections.size()); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> merge_function_wrapper, + MergeFileSplitRead::CreateMergeFunctionWrapper(owner_->options_, + owner_->context_->GetTableSchema(), + owner_->value_schema_, owner_->pool_)); + for (const std::vector& section : disk_sections) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr section_reader, + owner_->CreateSortMergeReaderForSection( + section, partition, dv_factory, owner_->predicate_for_keys_, + data_file_path_factory, /*drop_delete=*/false, merge_function_wrapper)); + section_readers.push_back( + std::make_unique(std::move(section_reader))); + } + std::unique_ptr concat_reader = + std::make_unique(std::move(section_readers)); + section_readers_guard.Release(); + readers->push_back(std::move(concat_reader)); + return Status::OK(); + } + + Result> CreateMergedReader( + std::vector>&& record_readers) { + ScopeGuard record_readers_guard([&record_readers]() { + for (const std::unique_ptr& reader : record_readers) { + if (reader) { + reader->Close(); + } + } + }); + if (record_readers.empty()) { + record_readers_guard.Release(); + return std::make_unique(std::vector>{}, + owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + owner_->CreateSortMergeReader(std::move(record_readers))); + record_readers_guard.Release(); + ScopeGuard sort_merge_reader_guard([&sort_merge_reader]() { + if (sort_merge_reader) { + sort_merge_reader->Close(); + } + }); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr result, + owner_->CreateProjectedReader(std::move(sort_merge_reader), + owner_->context_->GetPredicate(), + /*complete_row_kind=*/true)); + sort_merge_reader_guard.Release(); + return result; + } + + MergeFileSplitRead* owner_; +}; + Result> MergeFileSplitRead::Create( const std::shared_ptr& path_factory, const std::shared_ptr& context, @@ -158,6 +338,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(disk_splits, std::move(additional_readers), this); +} + void MergeFileSplitRead::SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper) { merge_function_wrapper_ = merge_function_wrapper; @@ -236,13 +422,10 @@ Result> MergeFileSplitRead::ApplyIndexAndDvRead Result> MergeFileSplitRead::CreateMergeReader( const std::shared_ptr& data_split, const std::shared_ptr& data_file_path_factory) { - auto dv_factory = DeletionVector::CreateFactory( - options_.GetFileSystem(), - DeletionVector::CreateDeletionFileMap(data_split->DataFiles(), data_split->DeletionFiles()), - pool_); - - std::vector> sections = - IntervalPartition(data_split->DataFiles(), key_comparator_).Partition(); + DeletionVector::Factory dv_factory; + std::vector> sections; + PAIMON_RETURN_NOT_OK(CreateDiskSections(data_split->DataFiles(), data_split->DeletionFiles(), + &dv_factory, §ions)); std::vector> batch_readers; batch_readers.reserve(sections.size()); // no overlap through multiple sections @@ -453,42 +636,121 @@ Result> MergeFileSplitRead::CreateReaderForSection( } else { predicate = context_->GetPredicate(); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, - CreateSortMergeReaderForSection(section, partition, dv_factory, - predicate, data_file_path_factory, - /*drop_delete=*/!force_keep_delete_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr sort_merge_reader, + CreateSortMergeReaderForSection(section, partition, dv_factory, predicate, + data_file_path_factory, /*drop_delete=*/false)); + return CreateProjectedReader(std::move(sort_merge_reader), /*predicate=*/nullptr, + /*complete_row_kind=*/false); +} + +Status MergeFileSplitRead::CreateDiskSections( + const std::vector>& data_files, + const std::vector>& deletion_files, + DeletionVector::Factory* dv_factory, std::vector>* sections) const { + *dv_factory = DeletionVector::CreateFactory( + options_.GetFileSystem(), DeletionVector::CreateDeletionFileMap(data_files, deletion_files), + pool_); + *sections = IntervalPartition(data_files, key_comparator_).Partition(); + return Status::OK(); +} + +Result>> +MergeFileSplitRead::CreateRecordReadersForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory) const { + std::vector> record_readers; + record_readers.reserve(section.size()); + for (const SortedRun& run : section) { + // no overlap in a run + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr run_reader, + CreateReaderForRun(partition, run, dv_factory, predicate, data_file_path_factory)); + record_readers.emplace_back(std::move(run_reader)); + } + return record_readers; +} + +Result> MergeFileSplitRead::CreateProjectedReader( + std::unique_ptr&& sort_merge_reader, + const std::shared_ptr& predicate, bool complete_row_kind) { + if (!force_keep_delete_) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } // KeyValueProjectionReader converts KeyValue objects to arrow array according to projection + std::unique_ptr projection_reader; if (!context_->EnableMultiThreadRowToBatch()) { - return KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, - projection_, options_.GetReadBatchSize(), pool_); - } - int32_t thread_number = context_->GetRowToBatchThreadNumber(); - assert(thread_number > 0); - return std::make_unique( - std::move(sort_merge_reader), raw_read_schema_, projection_, options_.GetReadBatchSize(), - thread_number, pool_); + PAIMON_ASSIGN_OR_RAISE( + projection_reader, + KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, + projection_, options_.GetReadBatchSize(), pool_)); + } else { + const int32_t thread_number = context_->GetRowToBatchThreadNumber(); + assert(thread_number > 0); + projection_reader = std::make_unique( + std::move(sort_merge_reader), raw_read_schema_, projection_, + options_.GetReadBatchSize(), thread_number, pool_); + } + ScopeGuard projection_reader_guard([&projection_reader]() { + if (projection_reader) { + projection_reader->Close(); + } + }); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr filtered_reader, + ApplyPredicateFilterIfNeeded(std::move(projection_reader), predicate)); + projection_reader_guard.Release(); + projection_reader = std::move(filtered_reader); + if (complete_row_kind) { + return std::make_unique(std::move(projection_reader), pool_); + } + return projection_reader; } Result> MergeFileSplitRead::CreateSortMergeReaderForSection( const std::vector& section, const BinaryRow& partition, DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, bool drop_delete) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr> merge_function_wrapper, + GetMergeFunctionWrapper()); + return CreateSortMergeReaderForSection(section, partition, dv_factory, predicate, + data_file_path_factory, drop_delete, + merge_function_wrapper); +} + +Result> MergeFileSplitRead::CreateSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, bool drop_delete, + const std::shared_ptr>& merge_function_wrapper) { // with overlap in one section + PAIMON_ASSIGN_OR_RAISE(std::vector> record_readers, + CreateRecordReadersForSection(section, partition, dv_factory, predicate, + data_file_path_factory)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr sort_merge_reader, + CreateSortMergeReader(std::move(record_readers), merge_function_wrapper)); + if (drop_delete) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } + return sort_merge_reader; +} + +Result> MergeFileSplitRead::CreateRawSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory) { std::vector> record_readers; record_readers.reserve(section.size()); for (const auto& run : section) { - // no overlap in a run PAIMON_ASSIGN_OR_RAISE( std::unique_ptr run_reader, CreateReaderForRun(partition, run, dv_factory, predicate, data_file_path_factory)); record_readers.emplace_back(std::move(run_reader)); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, - CreateSortMergeReader(std::move(record_readers))); - if (drop_delete) { - sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); - } - return sort_merge_reader; + return std::make_unique(std::move(record_readers), key_comparator_, + user_defined_seq_comparator_, + /*merge_function_wrapper=*/nullptr); } Result> MergeFileSplitRead::CreateReaderForRun( @@ -519,6 +781,12 @@ Result> MergeFileSplitRead::CreateSortMergeRead std::vector>&& record_readers) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr> merge_function_wrapper, GetMergeFunctionWrapper()); + return CreateSortMergeReader(std::move(record_readers), merge_function_wrapper); +} + +Result> MergeFileSplitRead::CreateSortMergeReader( + std::vector>&& record_readers, + const std::shared_ptr>& merge_function_wrapper) const { auto sort_engine = options_.GetSortEngine(); if (sort_engine == SortEngine::MIN_HEAP) { return std::make_unique( diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index d4bfa727c..07b5e70b2 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -74,8 +74,8 @@ class MergeFunctionWrapper; /// files->KeyValueProjectionReader/AsyncKeyValueProjectionReader /// ->DropDeleteReader->SortMergeReader->ConcatKeyValueRecordReader->KeyValueDataFileRecordReader /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ShreddingFileReader) -/// ->(MapSharedShreddingFileReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->FormatReader class MergeFileSplitRead : public AbstractSplitRead { public: static Result> Create( @@ -104,6 +104,12 @@ class MergeFileSplitRead : public AbstractSplitRead { DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, bool drop_delete); + /// Creates a min-heap reader which only sorts records and preserves duplicate keys. + Result> CreateRawSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory); + std::shared_ptr GetPathFactory() const { return path_factory_; } @@ -117,10 +123,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); @@ -134,6 +150,26 @@ class MergeFileSplitRead : public AbstractSplitRead { DeletionVector::Factory dv_factory, const std::shared_ptr& data_file_path_factory); + Status CreateDiskSections(const std::vector>& data_files, + const std::vector>& deletion_files, + DeletionVector::Factory* dv_factory, + std::vector>* sections) const; + + Result>> CreateRecordReadersForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory) const; + + Result> CreateProjectedReader( + std::unique_ptr&& sort_merge_reader, + const std::shared_ptr& predicate, bool complete_row_kind); + + Result> CreateSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, bool drop_delete, + const std::shared_ptr>& merge_function_wrapper); + Result> CreateReaderForRun( const BinaryRow& partition, const SortedRun& sorted_run, DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, @@ -142,6 +178,10 @@ class MergeFileSplitRead : public AbstractSplitRead { Result> CreateSortMergeReader( std::vector>&& record_readers); + Result> CreateSortMergeReader( + std::vector>&& record_readers, + const std::shared_ptr>& merge_function_wrapper) const; + Result>> GetMergeFunctionWrapper(); MergeFileSplitRead(const std::shared_ptr& path_factory, diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index 911cf7961..d02120de4 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -40,6 +40,7 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/key_value_in_memory_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/schema/schema_manager.h" @@ -51,7 +52,6 @@ #include "paimon/executor.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" -#include "paimon/metrics.h" #include "paimon/predicate/literal.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" @@ -66,6 +66,7 @@ class FileSystem; } // namespace paimon namespace paimon::test { + // Parameter: min_heap/loser_tree; enable/disable IO prefetch; enable/disable multi thread row to // batch class MergeFileSplitReadTest : public ::testing::Test, @@ -137,7 +138,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta1_2 = std::make_shared( "data-c80ccf0f-6387-4cbc-8889-ade8cef54c43-1.parquet", /*file_size=*/3370, /*row_count=*/4, @@ -154,7 +155,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta1_3 = std::make_shared( "data-c80ccf0f-6387-4cbc-8889-ade8cef54c43-2.parquet", /*file_size=*/3252, /*row_count=*/1, @@ -173,7 +174,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1(BinaryRowGenerator::GenerateRow({0, 0}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -200,7 +201,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta2_2 = std::make_shared( "data-24f8588c-d950-4e44-9d99-a023ea65a136-1.parquet", /*file_size=*/3229, /*row_count=*/1, @@ -218,7 +219,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder2(BinaryRowGenerator::GenerateRow({0, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -246,7 +247,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta3_2 = std::make_shared( "data-184f2304-49fd-4916-ba07-037757e904eb-1.parquet", /*file_size=*/3259, /*row_count=*/1, @@ -264,7 +265,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder3(BinaryRowGenerator::GenerateRow({1, 0}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -296,7 +297,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta1_2 = std::make_shared( "data-d03e13e5-5e2e-463a-b53a-8d44e4dc9141-1.parquet", /*file_size=*/2623, /*row_count=*/ @@ -314,7 +315,7 @@ class MergeFileSplitReadTest : public ::testing::Test, /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1( /*partition=*/BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/ @@ -328,9 +329,8 @@ class MergeFileSplitReadTest : public ::testing::Test, return {data_split1}; } - Result> CreateReader( - const std::shared_ptr& internal_context, - const std::vector>& data_splits) { + Result> CreateMergeFileSplitRead( + const std::shared_ptr& internal_context) { const auto& core_options = internal_context->GetCoreOptions(); const auto& table_schema = internal_context->GetTableSchema(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); @@ -347,9 +347,14 @@ class MergeFileSplitReadTest : public ::testing::Test, core_options.DataFilePrefix(), core_options.LegacyPartitionNameEnabled(), external_paths, global_index_external_path, core_options.IndexFileInDataFileDir(), pool_)); - PAIMON_ASSIGN_OR_RAISE(auto split_read, - MergeFileSplitRead::Create(path_factory, std::move(internal_context), - pool_, executor_)); + return MergeFileSplitRead::Create(path_factory, internal_context, pool_, executor_); + } + + Result> CreateReader( + const std::shared_ptr& internal_context, + const std::vector>& data_splits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr split_read, + CreateMergeFileSplitRead(internal_context)); std::vector> batch_readers; batch_readers.reserve(data_splits.size()); for (const auto& split : data_splits) { @@ -666,6 +671,73 @@ TEST_P(MergeFileSplitReadTest, TestSimple) { CheckResult(result_array, expected_array, read_schema); } +TEST_P(MergeFileSplitReadTest, TestRealtimeReadConcatenatesOrderedDiskSections) { + std::string path = + paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; + ReadContextBuilder context_builder(path); + std::vector raw_read_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("k1", arrow::int32())), + DataField(5, arrow::field("s1", arrow::utf8())), + DataField(6, arrow::field("v0", arrow::float64()))}; + std::shared_ptr read_schema = + DataField::ConvertDataFieldsToArrowSchema(raw_read_fields); + ASSERT_TRUE(read_schema); + + context_builder.SetReadFieldNames({"k0", "k1", "s1", "v0"}); + context_builder.SetOptions( + {{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::MERGE_ENGINE, "deduplicate"}}); + AddOptions(&context_builder); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); + std::shared_ptr internal_context = CreateInternalReadContext(read_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr split_read, + CreateMergeFileSplitRead(internal_context)); + + std::shared_ptr memory_type = + arrow::struct_(split_read->GetValueSchema()->fields()); + std::shared_ptr memory_array = + std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(memory_type, R"([ + [100, 200, "memory-late", 10000.0, "zzzz"], + [1, 1, "memory-delete", 1100.0, "zzzz"], + [0, 0, "memory-first", 1000.0, "zzzz"], + [50, 0, "memory-middle", 5000.0, "zzzz"] + ])") + .ValueOrDie()); + std::vector> memory_readers; + memory_readers.push_back(std::make_unique( + /*last_sequence_num=*/9, memory_array, + std::vector( + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT}), + std::vector({"k0", "k1"}), std::vector({"s0", "s1"}), + /*sequence_fields_ascending=*/true, split_read->GetKeyComparator(), pool_)); + + std::vector> disk_splits = {PrepareDataSplit().front()}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + split_read->CreateRealtimeReader(disk_splits, std::move(memory_readers))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, + ReadResultCollector::CollectResult(batch_reader.get())); + + arrow::FieldVector fields_with_row_kind = read_schema->fields(); + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + std::shared_ptr expected_array; + auto expected_status = + arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(fields_with_row_kind), {R"([ + [0, 0, 0, "memory-first", 1000.0], + [0, 0, 1, "you", 11.1], + [0, 1, 0, "later", 12.2], + [0, 1, 2, "!", 13.3], + [0, 50, 0, "memory-middle", 5000.0], + [0, 100, 200, "memory-late", 10000.0] + ])"}, + &expected_array); + ASSERT_TRUE(expected_status.ok()); + CheckResult(result_array, expected_array, read_schema); + ASSERT_TRUE(batch_reader->GetReaderMetrics()); + batch_reader->Close(); +} + TEST_P(MergeFileSplitReadTest, TestLookUp) { std::string path = paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; @@ -798,6 +870,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { context_builder.SetOptions({{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::MERGE_ENGINE, "deduplicate"}, {Options::IGNORE_DELETE, "true"}}); + context_builder.EnableLateMaterializing(false); AddOptions(&context_builder); // less_than will be ignore as it is partition predicate @@ -842,6 +915,63 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { CheckResult(result_array, expected_array, read_schema); } +TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { + std::string path = + paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; + ReadContextBuilder context_builder(path); + + std::vector raw_read_fields = {DataField(1, arrow::field("k1", arrow::int32())), + DataField(3, arrow::field("p1", arrow::int32())), + DataField(5, arrow::field("s1", arrow::utf8())), + DataField(4, arrow::field("s0", arrow::utf8())), + DataField(6, arrow::field("v0", arrow::float64())), + DataField(7, arrow::field("v1", arrow::boolean()))}; + auto read_schema = DataField::ConvertDataFieldsToArrowSchema(raw_read_fields); + ASSERT_TRUE(read_schema); + + context_builder.SetReadFieldNames({"k1", "p1", "s1", "s0", "v0", "v1"}); + context_builder.SetOptions({{Options::SEQUENCE_FIELD, "s0,s1"}, + {Options::MERGE_ENGINE, "deduplicate"}, + {Options::IGNORE_DELETE, "true"}}); + AddOptions(&context_builder); + context_builder.EnableLateMaterializing(true); + // key predicate, always pushed down into the data files + auto greater_or_equal = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k1", + FieldType::INT, Literal(1)); + // value predicate, only pushed down when a section holds a single sorted run + auto greater_than = PredicateBuilder::GreaterThan(/*field_index=*/4, /*field_name=*/"v0", + FieldType::DOUBLE, Literal(12.0)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate_result, + PredicateBuilder::And({greater_or_equal, greater_than})); + context_builder.SetPredicate(predicate_result); + context_builder.EnablePredicateFilter(true).EnableLateMaterializing(true); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); + + auto internal_context = CreateInternalReadContext(read_context); + ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, + ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_schema->fields(); + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + + // Only the merged rows with k1 >= 1 and v0 > 12.0 remain. + std::shared_ptr expected_array; + auto array_status = + arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(fields_with_row_kind), {R"([ + [0, 1, 0, "!", "driver", 13.3, false], + [0, 2, 0, "!", "driver", 13.3, false], + [0, 200, 0, "number", "max", 140.4, false], + [0, 1, 1, "you", "zoo", 130.0, false] + + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + CheckResult(result_array, expected_array, read_schema); +} + TEST_P(MergeFileSplitReadTest, TestReadWithAlterTable) { std::string path = paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; @@ -1235,7 +1365,7 @@ TEST_P(MergeFileSplitReadTest, Test09VersionWithoutInlineFieldId) { /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto meta2 = std::make_shared( "data-6871b960-edd9-40fc-9859-aaca9ea205cf-0.orc", /*file_size=*/887, /*row_count=*/5, /*min_key=*/BinaryRowGenerator::GenerateRow({std::string("Alex"), 0}, pool_.get()), @@ -1253,7 +1383,7 @@ TEST_P(MergeFileSplitReadTest, Test09VersionWithoutInlineFieldId) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( BinaryRowGenerator::GenerateRow({10}, pool_.get()), /*bucket=*/1, /*bucket_path=*/ diff --git a/src/paimon/core/operation/metrics/commit_stats_test.cpp b/src/paimon/core/operation/metrics/commit_stats_test.cpp index a18f88a61..b22f7a60f 100644 --- a/src/paimon/core/operation/metrics/commit_stats_test.cpp +++ b/src/paimon/core/operation/metrics/commit_stats_test.cpp @@ -59,7 +59,7 @@ ManifestEntry CreateEntry(const FileKind& kind, int32_t partition, int32_t bucke /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); return ManifestEntry(kind, part, bucket, /*total_buckets=*/10, file_meta); } diff --git a/src/paimon/core/operation/orphan_files_cleaner_impl.cpp b/src/paimon/core/operation/orphan_files_cleaner_impl.cpp index 6b8e5ba04..015d13408 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 ac211b257..646f24ac7 100644 --- a/src/paimon/core/operation/raw_file_split_read.h +++ b/src/paimon/core/operation/raw_file_split_read.h @@ -54,8 +54,9 @@ struct DeletionFile; /// splits)->CompleteRowKindBatchReader->(PredicateBatchReader) /// ->ConcatBatchReader across /// files->FieldMappingReader->(ApplyBitmapIndexBatchReader)->(CompleteRowTrackingFieldsBatchReader) -/// ->(ShreddingFileReader)->(MapSharedShreddingFileReader)->(DelegatingPrefetchReader) -/// ->(PrefetchFileBatchReader)->FormatReader +/// ->(ShreddingFileReader)->(VectorFileBatchReader) +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->FormatReader class RawFileSplitRead : public AbstractSplitRead { public: diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 6f7ae9781..51c478325 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -67,7 +67,7 @@ class RawFileSplitReadTest : public ::testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1(BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -91,7 +91,7 @@ class RawFileSplitReadTest : public ::testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder2(BinaryRowGenerator::GenerateRow({20, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -115,7 +115,7 @@ class RawFileSplitReadTest : public ::testing::Test { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder3(BinaryRowGenerator::GenerateRow({10, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + @@ -462,7 +462,7 @@ TEST_F(RawFileSplitReadTest, TestMatch) { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ paimon::test::GetDataDir() + diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index eb3d8826d..deacfa78b 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -35,14 +35,14 @@ ReadContext::ReadContext( const std::string& path, const std::string& branch, const std::vector& read_field_names, const std::vector& read_field_ids, const std::shared_ptr& predicate, bool enable_predicate_filter, bool enable_prefetch, - uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, - bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, - const std::optional& table_schema, const std::shared_ptr& memory_pool, - const std::shared_ptr& executor, + bool enable_late_materializing, uint32_t prefetch_batch_count, + uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, + uint32_t row_to_batch_thread_number, const std::optional& table_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& executor, 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), @@ -51,6 +51,7 @@ ReadContext::ReadContext( predicate_(predicate), enable_predicate_filter_(enable_predicate_filter), enable_prefetch_(enable_prefetch), + enable_late_materializing_(enable_late_materializing), prefetch_batch_count_(prefetch_batch_count), prefetch_max_parallel_num_(prefetch_max_parallel_num), enable_multi_thread_row_to_batch_(enable_multi_thread_row_to_batch), @@ -62,7 +63,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 +98,8 @@ class ReadContextBuilder::Impl { predicate_.reset(); enable_predicate_filter_ = false; enable_prefetch_ = false; - prefetch_cache_mode_ = PrefetchCacheMode::ALWAYS; + enable_late_materializing_ = false; + read_ahead_cache_enabled_ = true; prefetch_batch_count_ = 600; prefetch_max_parallel_num_ = 3; enable_multi_thread_row_to_batch_ = false; @@ -122,6 +124,7 @@ class ReadContextBuilder::Impl { std::shared_ptr predicate_; bool enable_predicate_filter_ = false; bool enable_prefetch_ = false; + bool enable_late_materializing_ = false; uint32_t prefetch_batch_count_ = 600; uint32_t prefetch_max_parallel_num_ = 3; bool enable_multi_thread_row_to_batch_ = false; @@ -131,7 +134,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_; }; @@ -191,6 +194,11 @@ ReadContextBuilder& ReadContextBuilder::EnablePrefetch(bool enabled) { return *this; } +ReadContextBuilder& ReadContextBuilder::EnableLateMaterializing(bool enabled) { + impl_->enable_late_materializing_ = enabled; + return *this; +} + ReadContextBuilder& ReadContextBuilder::SetPrefetchBatchCount(uint32_t batch_count) { impl_->prefetch_batch_count_ = batch_count; return *this; @@ -250,8 +258,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; } @@ -297,11 +305,12 @@ Result> ReadContextBuilder::Finish() { auto ctx = std::make_unique( impl_->path_, impl_->branch_, impl_->read_field_names_, impl_->read_field_ids_, impl_->predicate_, impl_->enable_predicate_filter_, impl_->enable_prefetch_, - impl_->prefetch_batch_count_, impl_->prefetch_max_parallel_num_, - 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_->enable_late_materializing_, impl_->prefetch_batch_count_, + impl_->prefetch_max_parallel_num_, 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_->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 8d5a78e14..c686cccb5 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/operation/write_restore_test.cpp b/src/paimon/core/operation/write_restore_test.cpp index 43f070154..754dc6959 100644 --- a/src/paimon/core/operation/write_restore_test.cpp +++ b/src/paimon/core/operation/write_restore_test.cpp @@ -40,7 +40,8 @@ std::shared_ptr CreateDataFileMeta(const std::string& file_name) { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, /*file_source=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } ManifestEntry CreateManifestEntry(int32_t total_buckets, const std::string& file_name) { diff --git a/src/paimon/core/postpone/postpone_bucket_writer.cpp b/src/paimon/core/postpone/postpone_bucket_writer.cpp index 47fd3bcb8..71dfbc074 100644 --- a/src/paimon/core/postpone/postpone_bucket_writer.cpp +++ b/src/paimon/core/postpone/postpone_bucket_writer.cpp @@ -31,10 +31,8 @@ #include "arrow/c/helpers.h" #include "arrow/scalar.h" #include "fmt/format.h" -#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/metrics/metrics_impl.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" @@ -43,8 +41,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_data_file_writer_factory.h" -#include "paimon/core/io/shredding_key_value_data_file_writer_factory.h" +#include "paimon/core/io/key_value_data_file_writer_factories.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/utils/commit_increment.h" #include "paimon/format/file_format.h" @@ -55,27 +52,12 @@ namespace paimon { class InternalRow; class MemoryPool; -namespace { - -std::shared_ptr BuildPostponeBucketWriteSchema( - const std::shared_ptr& value_schema) { - arrow::FieldVector target_fields; - target_fields.push_back( - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())); - target_fields.push_back(DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())); - target_fields.insert(target_fields.end(), value_schema->fields().begin(), - value_schema->fields().end()); - return arrow::schema(target_fields); -} - -} // namespace - Result> PostponeBucketWriter::Create( const std::vector& trimmed_primary_keys, const std::shared_ptr& path_factory, int64_t schema_id, const std::shared_ptr& value_schema, const CoreOptions& options, const std::shared_ptr& pool) { - auto write_schema = BuildPostponeBucketWriteSchema(value_schema); + auto write_schema = SpecialFields::CompleteSequenceAndValueKindField(value_schema); return std::unique_ptr(new PostponeBucketWriter( trimmed_primary_keys, path_factory, schema_id, value_schema, write_schema, options, pool)); } @@ -258,20 +240,12 @@ PostponeBucketWriter::PrepareMinMaxKey( Result>>> PostponeBucketWriter::CreateRollingRowWriter() const { - std::shared_ptr>> factory; PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr plan_factory, - ShreddingWritePlanFactories::SelectActive(options_, write_schema_, pool_)); - if (plan_factory != nullptr) { - factory = std::make_shared( - options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/false, plan_factory, - pool_); - } else { - factory = std::make_shared( - options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/false, pool_); - } + std::shared_ptr factory, + KeyValueDataFileWriterFactories::Create(options_, schema_id_, write_schema_, /*level=*/0, + FileSource::Append(), trimmed_primary_keys_, + path_factory_, /*create_stats_extractor=*/false, + /*is_changelog=*/false, pool_)); return std::make_unique>>( options_.GetTargetFileSize(/*has_primary_key=*/true), options_.GetTargetFileRowNum(), factory); diff --git a/src/paimon/core/postpone/postpone_bucket_writer_test.cpp b/src/paimon/core/postpone/postpone_bucket_writer_test.cpp index e932d40da..30357327a 100644 --- a/src/paimon/core/postpone/postpone_bucket_writer_test.cpp +++ b/src/paimon/core/postpone/postpone_bucket_writer_test.cpp @@ -226,7 +226,7 @@ TEST_P(PostponeBucketWriterTest, TestSimple) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); @@ -308,7 +308,7 @@ TEST_P(PostponeBucketWriterTest, TestNestedType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); @@ -474,7 +474,7 @@ TEST_P(PostponeBucketWriterTest, TestWriteMultiBatch) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); @@ -624,7 +624,7 @@ TEST_P(PostponeBucketWriterTest, TestMultiplePrepareCommit) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment1({expected_data_file_meta1}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment1, commit_increment1.GetNewFilesIncrement()); @@ -644,7 +644,7 @@ TEST_P(PostponeBucketWriterTest, TestMultiplePrepareCommit) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement expected_data_increment2({expected_data_file_meta2}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment2, commit_increment2.GetNewFilesIncrement()); diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index cf1e37aac..1136243e8 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -25,32 +25,43 @@ #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/arrow_utils.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" namespace paimon { namespace { -uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { - uint64_t result = 0; - for (const std::shared_ptr& buffer : data->buffers) { - if (buffer) { - result += static_cast(buffer->size()); - } - } - for (const std::shared_ptr& child : data->child_data) { - result += GetArrayMemoryUsage(child); - } - if (data->dictionary) { - result += GetArrayMemoryUsage(data->dictionary); +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; } - return result; } } // namespace @@ -165,11 +176,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 +206,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 +234,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 +271,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 +371,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()); + uint64_t memory_usage = ArrowUtils::GetArrayMemoryUsage(struct_array->data()); + if (statistics) { + memory_usage += ArrowUtils::GetArrayMemoryUsage(statistics->min_values->data()) + + ArrowUtils::GetArrayMemoryUsage(statistics->max_values->data()) + + ArrowUtils::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 +443,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 bd9ecfbf7..97339852f 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 0eecb5d02..dff12b589 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -25,24 +25,37 @@ #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, request.memory_pool)); + 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 cc08cc8e1..864b1f810 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,69 @@ 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_, 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, 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/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 000000000..421eb0c8d --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,309 @@ +/* + * 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/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.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/memory/memory_pool.h" + +namespace paimon { + +namespace { + +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, + std::shared_ptr arrow_pool) + : arrow_pool_(std::move(arrow_pool)), + data_(batch.data), + metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!data_) { + return MakeEofBatch(); + } + auto array = std::make_unique(); + auto schema = std::make_unique(); + ScopeGuard export_guard([array_ptr = array.get(), schema_ptr = schema.get()]() { + ArrowArrayRelease(array_ptr); + ArrowSchemaRelease(schema_ptr); + }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr record_batch, + arrow::RecordBatch::FromStructArray(data_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr normalized_batch, + ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(array.get(), arrow_pool_)); + data_.reset(); + arrow_pool_.reset(); + export_guard.Release(); + return ReadBatch(std::move(array), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + void Close() override { + data_.reset(); + arrow_pool_.reset(); + } + + private: + std::shared_ptr arrow_pool_; + std::shared_ptr data_; + std::shared_ptr metrics_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + Impl(std::shared_ptr transport_schema, + std::shared_ptr arrow_pool) + : transport_schema_(std::move(transport_schema)), arrow_pool_(std::move(arrow_pool)) {} + + 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_(transport_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time transport batch is not a StructArray"); + } + std::shared_ptr transport = + checked_pointer_cast(array); + std::lock_guard lock(mutex_); + building_.push_back(StoredBatch{transport, write_batch.offset_range, + ArrowUtils::GetArrayMemoryUsage(transport->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, arrow_pool_)); + } + 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::make_shared(std::move(segments)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t, + const RealtimeQueryContext& context) { + 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"); + } + if (context.read_schema == nullptr || context.read_schema->release == nullptr) { + return Status::Invalid("PK real-time query read schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, + arrow::ImportSchema(context.read_schema)); + std::vector> readers; + for (const std::shared_ptr& segment : typed->Segments()) { + for (const StoredBatch& batch : segment->Batches()) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr projected, + NestedProjectionUtils::AlignArrayToReadType( + batch.data, arrow::struct_(read_schema->fields()), arrow_pool_.get())); + if (!projected || projected->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + "PK memory query projection did not produce a StructArray"); + } + StoredBatch query_batch{checked_pointer_cast(projected), + batch.offset_range, /*memory_usage=*/0}; + readers.push_back(std::make_unique(query_batch, arrow_pool_)); + } + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_end_offset) { + std::lock_guard lock(mutex_); + auto first_retained = std::find_if( + sealed_.begin(), sealed_.end(), [committed_end_offset](const auto& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + sealed_.erase(sealed_.begin(), first_retained); + 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 transport_schema_; + std::shared_ptr arrow_pool_; + 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& transport_schema, + const std::shared_ptr& memory_pool) { + if (!memory_pool) { + return Status::Invalid("PK real-time store memory pool is null"); + } + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + return std::shared_ptr(new PrimaryKeyRealtimeStore( + std::make_unique(transport_schema, std::move(arrow_pool)))); +} +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 000000000..46c0fe8f7 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,61 @@ +/* + * 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 MemoryPool; + +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& memory_pool); + + ~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 000000000..ed80db275 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,406 @@ +/* + * 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 + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.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/realtime/realtime_primary_key_reader.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr FieldWithId(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))) + ->WithNullable(nullable); +} + +std::shared_ptr TransportSchema() { + return RealtimePrimaryKeyLayout::CreateSchema( + {FieldWithId("id", arrow::int64(), 0), FieldWithId("value", arrow::utf8(), 1)}); +} + +std::shared_ptr NestedTransportSchema() { + return RealtimePrimaryKeyLayout::CreateSchema( + {FieldWithId("id", arrow::int64(), 0), + FieldWithId("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}), + 1)}); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(TransportSchema()->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 MakeSlicedBatch(const std::shared_ptr& schema, + const std::string& json, int64_t offset, + int64_t length) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie() + ->Slice(offset, length); + 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(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + 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, TestCommitReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + 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(); + } +} + +void AssertSlicedBatch(BatchReader* reader) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + std::shared_ptr values = checked_pointer_cast(array); + ASSERT_EQ(2, checked_pointer_cast(values->field(3))->Value(0)); + ASSERT_EQ(3, checked_pointer_cast(values->field(3))->Value(1)); + std::shared_ptr nested = + checked_pointer_cast(values->field(4)); + ASSERT_EQ("two", checked_pointer_cast(nested->field(0))->GetString(0)); + ASSERT_EQ("three", checked_pointer_cast(nested->field(0))->GetString(1)); + std::shared_ptr items = + checked_pointer_cast(nested->field(1)); + std::shared_ptr first_items = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(3, first_items->Value(0)); + ASSERT_EQ(4, first_items->Value(1)); + std::shared_ptr second_items = + checked_pointer_cast(items->value_slice(1)); + ASSERT_EQ(5, second_items->Value(0)); + ASSERT_EQ(6, second_items->Value(1)); + ASSERT_OK_AND_ASSIGN(batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { + std::shared_ptr schema = NestedTransportSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + schema, + R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]], [0, 3, 2, 3, ["three", [5, 6]]], [0, 4, 3, 4, ["four", [7, 8]]]])", + 1, 2), + 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()); + AssertSlicedBatch(readers[0].get()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(readers, store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + AssertSlicedBatch(readers[0].get()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + 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_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 5, 2, "two"]])"), OffsetRange(5, 6)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 2, 6, 3, "three"]])"), OffsetRange(6, 7)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr retained_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), retained_view->GetOffsetRange()); + + const uint64_t initial_memory_usage = store->GetMemoryUsage(); + ASSERT_GT(initial_memory_usage, 0); + ASSERT_OK(store->AdvanceCommittedOffset(4)); + ASSERT_EQ(initial_memory_usage, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_LT(store->GetMemoryUsage(), initial_memory_usage); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(5, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(6)); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(6, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(7)); + ASSERT_EQ(0, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_FALSE(current_view->GetOffsetRange().has_value()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(retained_view, /*offset_begin=*/0, context)); + 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\"")); + ASSERT_NE(std::string::npos, actual.find("\"three\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + 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()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{/*read_schema=*/c_schema.get(), /*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\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryPoolOutlivesStoreReaderAndExport) { + const std::shared_ptr stored_schema = TransportSchema(); + std::shared_ptr pool = GetMemoryPool(); + std::weak_ptr pool_lifetime = pool; + auto write_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, write_schema.get()).ok()); + ArrowRealtimeStoreFactory factory; + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + factory.Create(RealtimeStoreCreateRequest{ + std::move(write_schema), + /*options=*/{}, pool, RealtimeStoreMode::PRIMARY_KEY})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 7, "seven"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + view.reset(); + store.reset(); + pool.reset(); + ASSERT_FALSE(pool_lifetime.expired()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr imported = std::move(import_result).ValueOrDie(); + readers.clear(); + ASSERT_FALSE(pool_lifetime.expired()); + imported.reset(); + ASSERT_TRUE(pool_lifetime.expired()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { + const std::shared_ptr stored_profile_a = + FieldWithId("profile_a", arrow::int32(), 30); + const std::shared_ptr stored_a = FieldWithId("a", arrow::int32(), 10); + const std::shared_ptr stored_b = FieldWithId("b", arrow::int32(), 11); + const std::shared_ptr stored_x = FieldWithId("x", arrow::int32(), 20); + const std::shared_ptr stored_y = FieldWithId("y", arrow::int32(), 21); + arrow::FieldVector stored_value_fields = { + FieldWithId("id", arrow::int64(), 0), + FieldWithId("profile", arrow::struct_({stored_profile_a}), 1), + FieldWithId("items", arrow::list(arrow::struct_({stored_a, stored_b})), 2), + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_x, stored_y})), 3)}; + std::shared_ptr stored_schema = + RealtimePrimaryKeyLayout::CreateSchema(stored_value_fields); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(stored_schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + stored_schema, + R"([[0, 1, 0, 6, [5], [[1, 2]], [["before", [3, 4]]]], [0, 2, 1, 7, [50], [[100, 200], null], [["k1", [7, 8]], ["k2", null]]], [0, 3, 2, 8, [500], [[9, 10]], [["after", [11, 12]]]]])", + 1, 1), + OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + arrow::FieldVector requested_value_fields; + requested_value_fields.push_back(FieldWithId("profile", arrow::struct_({stored_profile_a}), 1)); + requested_value_fields.push_back( + FieldWithId("items", arrow::list(arrow::struct_({stored_b, stored_a})), 2)); + requested_value_fields.push_back( + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_y, stored_x})), 3)); + std::shared_ptr requested_schema = + RealtimePrimaryKeyLayout::CreateSchema(requested_value_fields); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*requested_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(requested_schema->fields()))); + std::shared_ptr projected = checked_pointer_cast(array); + const std::shared_ptr profile = + checked_pointer_cast(projected->field(3)); + ASSERT_EQ(50, checked_pointer_cast(profile->field(0))->Value(0)); + const std::shared_ptr items = + checked_pointer_cast(projected->field(4)); + const std::shared_ptr item_values = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(200, checked_pointer_cast(item_values->field(0))->Value(0)); + ASSERT_EQ(100, checked_pointer_cast(item_values->field(1))->Value(0)); + ASSERT_TRUE(item_values->IsNull(1)); + + const std::shared_ptr attrs = + checked_pointer_cast(projected->field(5)); + const int64_t attr_offset = attrs->value_offset(0); + const int64_t attr_length = attrs->value_length(0); + const std::shared_ptr attr_keys = + checked_pointer_cast(attrs->keys()->Slice(attr_offset, attr_length)); + ASSERT_EQ("k1", attr_keys->GetString(0)); + const std::shared_ptr attr_values = + checked_pointer_cast(attrs->items()->Slice(attr_offset, attr_length)); + ASSERT_EQ(8, checked_pointer_cast(attr_values->field(0))->Value(0)); + ASSERT_EQ(7, checked_pointer_cast(attr_values->field(1))->Value(0)); + ASSERT_TRUE(attr_values->IsNull(1)); +} + +} // 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 af48b289b..632d64e16 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, + 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)); + std::move(request), RealtimePartitionBucket(partition, bucket))); 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 29c198c64..a6190d3b0 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 1c7a7cf55..2b9d7d07c 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -34,14 +34,33 @@ #include #include +#include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "fmt/format.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 +97,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, const RealtimePartitionBucket& partition_bucket) { + 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); + auto iter = stores_.find(partition_bucket); int64_t initial_offset = 0; - auto offset_iter = committed_offsets_.find(key); + auto offset_iter = committed_offsets_.find(partition_bucket); 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(fmt::format( + "real-time store schema or mode mismatch for partition {}, bucket {}; recreate " + "the RealtimeContext", + PartitionToString(partition_bucket.partition), partition_bucket.bucket)); } 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 +141,54 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset}; + return RealtimeStoreState{iter->second.store, initial_offset}; + } + 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; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, + factory_->Create(std::move(request))); + if (!store) { + return Status::Invalid("real-time store factory returned a null store"); } - Result> store_result = - factory_->Create(std::move(write_schema), options, memory_pool); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); - stores_.emplace(key, store); + stores_.emplace(partition_bucket, StoreEntry{store, requested_schema, mode}); if (offset_iter != committed_offsets_.end()) { - reclaimed_offsets_.emplace(key, offset_iter->second); + reclaimed_offsets_.emplace(partition_bucket, 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(fmt::format("real-time store not found for partition {}, bucket {}", + PartitionToString(partition_bucket.partition), + 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()); + if (!read_view) { + return Status::Invalid("real-time store returned a null read view"); + } result.push_back( - RealtimePartitionBucketView{partition_bucket, store, std::move(read_view)}); + RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; } @@ -229,12 +281,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 +321,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 13fec77ab..ea069a5cd 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; @@ -65,10 +70,10 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { 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); + RealtimeStoreCreateRequest&& request, const RealtimePartitionBucket& partition_bucket); + + Result AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number); Result> AcquireReadViews(); @@ -78,6 +83,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 +97,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 +113,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 ea050094f..c9531eda5 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,24 @@ 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; + if (return_null_read_view) { + return std::shared_ptr(); + } 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 +74,6 @@ class TestingRealtimeStore : public RealtimeStore { committed_offsets.push_back(committed_offset); return Status::OK(); } - uint64_t GetMemoryUsage() const override { return 0; } @@ -86,32 +81,37 @@ class TestingRealtimeStore : public RealtimeStore { int32_t acquire_count = 0; int32_t advance_count = 0; bool fail_next_advance = false; + bool return_null_read_view = false; std::vector committed_offsets; }; 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()); + if (return_null_store) { + return std::shared_ptr(); + } auto store = std::make_shared(); stores.push_back(store); return store; } + bool return_null_store = false; 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 +121,42 @@ 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, + RealtimeStoreMode::APPEND_ONLY, statistics_mode}, + RealtimePartitionBucket(partition, bucket)); +} + +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 +164,87 @@ 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(), RealtimeStoreMode::PRIMARY_KEY}, + RealtimePartitionBucket(partition, 0)), + "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(context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{MakeWriteSchema(), /*options=*/{}, GetDefaultPool(), + RealtimeStoreMode::PRIMARY_KEY}, + partition_bucket)); + + 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, 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 +261,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 +280,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 +328,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 +338,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 +398,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()); @@ -286,5 +420,20 @@ TEST(RealtimeContextTest, TestRejectsNullFactory) { "real-time store factory is null"); } +TEST(RealtimeContextTest, TestRejectsNullPluginResults) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + factory->return_null_store = true; + ASSERT_NOK_WITH_MSG(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, + MakeWriteSchema(), {}, GetDefaultPool()), + "real-time store factory returned a null store"); + + factory->return_null_store = false; + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); + factory->stores[0]->return_null_read_view = true; + ASSERT_NOK_WITH_MSG(context->AcquireReadViews(), "real-time store returned a null read view"); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.cpp b/src/paimon/core/realtime/realtime_primary_key_reader.cpp new file mode 100644 index 000000000..19ba0a985 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader.cpp @@ -0,0 +1,510 @@ +/* + * 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_reader.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.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/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" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { + +namespace { + +template +void CloseReaders(const std::vector>& readers) { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + +class RealtimeOffsetCoverage { + public: + static Result> Create(const OffsetRange& offsets, + size_t reader_count, + bool allow_committed_prefix) { + if (offsets.begin < 0 || offsets.end < offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid offset range"); + } + return std::shared_ptr( + new RealtimeOffsetCoverage(offsets, reader_count, allow_committed_prefix)); + } + + Status Add(const arrow::Int64Array& offsets) { + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (allow_committed_prefix_ && offset < 0) { + return Status::Invalid("PK real-time store reader offset must be non-negative"); + } + if (allow_committed_prefix_ && offset < offsets_.begin) { + continue; + } + if (offset < offsets_.begin || offset >= offsets_.end) { + return Status::Invalid( + allow_committed_prefix_ + ? "PK real-time store query reader offset is outside the visible range" + : "PK real-time store commit reader offset is outside the sealed range"); + } + if (!seen_offsets_.CheckedAdd(offset)) { + return CoverageError(); + } + } + return Status::OK(); + } + + Status FinishReader() { + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && + seen_offsets_.Cardinality() != offsets_.Count()) { + return CoverageError(); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& offsets, size_t reader_count, + bool allow_committed_prefix) + : offsets_(offsets), + reader_count_(reader_count), + allow_committed_prefix_(allow_committed_prefix) {} + + Status CoverageError() const { + return Status::Invalid( + allow_committed_prefix_ + ? "PK real-time store query readers did not cover the visible range" + : "PK real-time store commit readers did not cover the sealed range"); + } + + OffsetRange offsets_; + size_t reader_count_; + bool allow_committed_prefix_; + RoaringBitmap64 seen_offsets_; + size_t finished_reader_count_ = 0; +}; + +Status CheckTransportField(const std::shared_ptr& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid( + fmt::format("realtime primary-key transport schema is missing 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( + "realtime primary-key transport 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> ResolveFieldIndexes( + const std::shared_ptr& transport_schema, + const std::unordered_map& field_indexes, + const std::shared_ptr& row_schema) { + std::vector result; + result.reserve(row_schema->num_fields()); + for (const std::shared_ptr& row_field : row_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(row_field)); + auto field_index = field_indexes.find(field_id); + if (field_index == field_indexes.end()) { + return Status::Invalid(fmt::format( + "cannot find field id {} in realtime primary-key transport schema", field_id)); + } + const std::shared_ptr& transport_field = + transport_schema->field(field_index->second); + if (!transport_field->type()->Equals(row_field->type())) { + return Status::Invalid(fmt::format( + "realtime primary-key transport field id {} type {} does not match row type {}", + field_id, transport_field->type()->ToString(), row_field->type()->ToString())); + } + result.push_back(field_index->second); + } + return result; +} + +class RealtimePrimaryKeyReaderPlan { + public: + static Result> Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema) { + std::unordered_map field_indexes; + field_indexes.reserve(transport_schema->num_fields() - + RealtimePrimaryKeyLayout::kValueStartIndex); + for (int32_t i = RealtimePrimaryKeyLayout::kValueStartIndex; + i < transport_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId( + transport_schema->field(i))); + if (!field_indexes.emplace(field_id, i).second) { + return Status::Invalid(fmt::format( + "duplicate field id {} in realtime primary-key transport schema", field_id)); + } + } + PAIMON_ASSIGN_OR_RAISE(std::vector key_field_indexes, + ResolveFieldIndexes(transport_schema, field_indexes, key_schema)); + PAIMON_ASSIGN_OR_RAISE(std::vector value_field_indexes, + ResolveFieldIndexes(transport_schema, field_indexes, value_schema)); + return std::shared_ptr(new RealtimePrimaryKeyReaderPlan( + transport_schema, std::move(key_field_indexes), std::move(value_field_indexes))); + } + + const std::shared_ptr& TransportSchema() const { + return transport_schema_; + } + + const std::vector& KeyFieldIndexes() const { + return key_field_indexes_; + } + + const std::vector& ValueFieldIndexes() const { + return value_field_indexes_; + } + + private: + RealtimePrimaryKeyReaderPlan(const std::shared_ptr& schema, + std::vector&& key_indexes, + std::vector&& value_indexes) + : transport_schema_(schema), + key_field_indexes_(std::move(key_indexes)), + value_field_indexes_(std::move(value_indexes)) {} + + const std::shared_ptr transport_schema_; + const std::vector key_field_indexes_; + const std::vector value_field_indexes_; +}; + +class RealtimePrimaryKeyReader final : public KeyValueRecordReader { + public: + RealtimePrimaryKeyReader(std::unique_ptr&& reader, + const std::shared_ptr& plan, + const std::optional& visible_offsets, + const std::shared_ptr& pool, + const std::shared_ptr& offset_coverage) + : reader_(std::move(reader)), + plan_(plan), + visible_offsets_(visible_offsets), + pool_(pool), + offset_coverage_(offset_coverage) {} + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(RealtimePrimaryKeyReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->RowCount(); + } + + Result Next() override { + if (cursor_ >= reader_->RowCount()) { + return Status::Invalid("No more realtime primary-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: + RealtimePrimaryKeyReader* reader_; + int64_t cursor_ = 0; + }; + + Result> NextBatch() override { + return NextBatchImpl(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + ResetBatchState(); + reader_->Close(); + } + + private: + Result> NextBatchImpl() { + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } + return std::unique_ptr(); + } + auto& [batch, selection] = batch_with_bitmap; + 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 realtime primary-key transport batch to StructArray"); + } + std::shared_ptr data_batch = + checked_pointer_cast(arrow_array); + PAIMON_RETURN_NOT_OK(ValidateTransportBatch(data_batch)); + + std::shared_ptr> offset_array = + checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex)); + if (offset_coverage_) { + PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array)); + } + + row_kind_array_ = checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kValueKindIndex)); + sequence_number_array_ = checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kSequenceNumberIndex)); + arrow::ArrayVector key_fields; + key_fields.reserve(plan_->KeyFieldIndexes().size()); + for (int32_t index : plan_->KeyFieldIndexes()) { + key_fields.push_back(data_batch->field(index)); + } + arrow::ArrayVector value_fields; + value_fields.reserve(plan_->ValueFieldIndexes().size()); + for (int32_t index : plan_->ValueFieldIndexes()) { + value_fields.push_back(data_batch->field(index)); + } + key_ctx_ = std::make_shared(key_fields, pool_); + value_ctx_ = std::make_shared(value_fields, pool_); + PAIMON_ASSIGN_OR_RAISE(bool has_selected_rows, + SelectRows(*offset_array, std::move(selection))); + if (!has_selected_rows) { + continue; + } + ArrowUtils::TraverseArray(data_batch); + return std::make_unique(this); + } + } + + Status ValidateTransportBatch(const std::shared_ptr& data_batch) const { + if (data_batch->num_fields() != plan_->TransportSchema()->num_fields()) { + return Status::Invalid(fmt::format( + "realtime primary-key transport batch field count {} does not match schema field " + "count {}", + data_batch->num_fields(), plan_->TransportSchema()->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(plan_->TransportSchema()->field(i), true)) { + return Status::Invalid(fmt::format( + "realtime primary-key transport batch field {} does not match declared schema", + i)); + } + } + if (data_batch->field(RealtimePrimaryKeyLayout::kValueKindIndex)->null_count() != 0 || + data_batch->field(RealtimePrimaryKeyLayout::kSequenceNumberIndex)->null_count() != 0 || + data_batch->field(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex)->null_count() != 0) { + return Status::Invalid("realtime primary-key transport columns must not contain nulls"); + } + return Status::OK(); + } + + Result SelectRows(const arrow::Int64Array& offsets, RoaringBitmap32&& selection) { + for (auto iter = selection.Begin(); iter != selection.End(); ++iter) { + const uint32_t row = *iter; + if (static_cast(row) >= offsets.length()) { + return Status::Invalid( + fmt::format("selected row id {} is out of bounds for realtime primary-key " + "transport batch length {}", + row, offsets.length())); + } + } + if (selection.Cardinality() != offsets.length()) { + return Status::Invalid( + "PK real-time store reader bitmap must cover every raw " + "transport row"); + } + selected_rows_.reserve(offsets.length()); + for (int64_t row = 0; row < offsets.length(); ++row) { + if (!visible_offsets_.has_value() || (offsets.Value(row) >= visible_offsets_->begin && + offsets.Value(row) < visible_offsets_->end)) { + selected_rows_.push_back(row); + } + } + return !selected_rows_.empty(); + } + + int64_t RowCount() const { + return static_cast(selected_rows_.size()); + } + + int64_t RowAt(int64_t ordinal) const { + return selected_rows_[ordinal]; + } + + void ResetBatchState() { + key_ctx_.reset(); + value_ctx_.reset(); + row_kind_array_.reset(); + sequence_number_array_.reset(); + selected_rows_.clear(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr plan_; + std::optional visible_offsets_; + std::shared_ptr 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::vector selected_rows_; +}; + +} // namespace + +std::shared_ptr RealtimePrimaryKeyLayout::CreateSchema( + const std::vector>& value_fields) { + arrow::FieldVector fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + fields.insert(fields.end(), value_fields.begin(), value_fields.end()); + return arrow::schema(std::move(fields)); +} + +Status RealtimePrimaryKeyLayout::ValidateSchema( + const std::shared_ptr& transport_schema) { + if (!transport_schema || transport_schema->num_fields() < kValueStartIndex) { + return Status::Invalid( + "realtime primary-key transport schema must contain transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckTransportField(transport_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK(CheckTransportField(transport_schema, kSequenceNumberIndex, + SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK(CheckTransportField(transport_schema, kRealtimeOffsetIndex, + SpecialFields::RealtimeOffset())); + return Status::OK(); +} + +Result>> +RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard remaining_raw_readers_guard([&readers]() { CloseReaders(readers); }); + if (readers.empty() && visible_offsets.begin < visible_offsets.end) { + return Status::Invalid( + "PK real-time store returned no query readers for a non-empty visible range"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null query reader"); + } + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr plan, + RealtimePrimaryKeyReaderPlan::Create(transport_schema, key_schema, value_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(visible_offsets, readers.size(), + /*allow_committed_prefix=*/true)); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + adapted_readers.push_back(std::make_unique( + std::move(reader), plan, visible_offsets, memory_pool, offset_coverage)); + } + remaining_raw_readers_guard.Release(); + return adapted_readers; +} + +Result>> +RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::vector>&& readers, + const std::shared_ptr& transport_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 remaining_raw_readers_guard([&readers]() { CloseReaders(readers); }); + if (readers.empty()) { + return Status::Invalid( + "PK real-time store returned no commit readers for a sealed segment"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr plan, + RealtimePrimaryKeyReaderPlan::Create(transport_schema, key_schema, value_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(sealed_offsets, readers.size(), + /*allow_committed_prefix=*/false)); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + adapted_readers.push_back(std::make_unique( + std::move(reader), plan, std::nullopt, memory_pool, offset_coverage)); + } + remaining_raw_readers_guard.Release(); + return adapted_readers; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.h b/src/paimon/core/realtime/realtime_primary_key_reader.h new file mode 100644 index 000000000..d175c3b63 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader.h @@ -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. + */ + +#pragma once + +#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; + +/// Defines the Arrow field layout for PK realtime transport batches. +class RealtimePrimaryKeyLayout { + public: + RealtimePrimaryKeyLayout() = delete; + ~RealtimePrimaryKeyLayout() = delete; + + static constexpr int32_t kValueKindIndex = 0; + static constexpr int32_t kSequenceNumberIndex = 1; + static constexpr int32_t kRealtimeOffsetIndex = 2; + static constexpr int32_t kValueStartIndex = 3; + + static std::shared_ptr CreateSchema( + const std::vector>& value_fields); + + static Status ValidateSchema(const std::shared_ptr& transport_schema); +}; + +class RealtimePrimaryKeyReaderFactory { + public: + RealtimePrimaryKeyReaderFactory() = delete; + ~RealtimePrimaryKeyReaderFactory() = delete; + + static Result>> CreateForQuery( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + + static Result>> CreateForCommit( + std::vector>&& readers, + const std::shared_ptr& transport_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/realtime_primary_key_reader_test.cpp b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp new file mode 100644 index 000000000..89a39cd63 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp @@ -0,0 +1,688 @@ +/* + * 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_reader.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/array_nested.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.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/utils/key_value_checker.h" +#include "paimon/testing/utils/read_result_collector.h" +#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 MakeTransportSchema(const arrow::FieldVector& value_fields) { + return RealtimePrimaryKeyLayout::CreateSchema(value_fields); +} + +Result> CreateRealtimePrimaryKeyQueryReaderForTest( + std::unique_ptr&& reader, const std::shared_ptr& transport_schema, + const OffsetRange& visible_offsets, const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> readers; + readers.push_back(std::move(reader)); + PAIMON_ASSIGN_OR_RAISE(std::vector> adapted_readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(readers), transport_schema, visible_offsets, key_schema, + value_schema, memory_pool)); + return std::move(adapted_readers[0]); +} + +Result> CreateRealtimePrimaryKeyCommitReaderForTest( + std::unique_ptr&& reader, const std::shared_ptr& transport_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> readers; + readers.push_back(std::move(reader)); + PAIMON_ASSIGN_OR_RAISE(std::vector> adapted_readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(readers), transport_schema, sealed_offsets, key_schema, + value_schema, memory_pool)); + return std::move(adapted_readers[0]); +} + +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_; +}; + +class MalformedBitmapBatchReader : public BatchReader { + public: + MalformedBitmapBatchReader(std::unique_ptr&& delegate, int32_t row_id) + : delegate_(std::move(delegate)), row_id_(row_id) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + Result NextBatchWithBitmap() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, delegate_->NextBatchWithBitmap()); + if (!IsEofBatch(batch)) { + batch.second.Add(row_id_); + } + return batch; + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + int32_t row_id_; +}; + +} // namespace + +class RealtimePrimaryKeyReaderTest : public testing::Test { + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +TEST_F(RealtimePrimaryKeyReaderTest, TestTransportSchemaLayout) { + arrow::FieldVector value_fields = {arrow::field("key", arrow::int64(), false), + arrow::field("value", arrow::utf8())}; + std::shared_ptr schema = MakeTransportSchema(value_fields); + + ASSERT_EQ(RealtimePrimaryKeyLayout::kValueKindIndex, 0); + ASSERT_EQ(RealtimePrimaryKeyLayout::kSequenceNumberIndex, 1); + ASSERT_EQ(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex, 2); + ASSERT_EQ(RealtimePrimaryKeyLayout::kValueStartIndex, 3); + ASSERT_EQ(schema->field(0)->name(), "_VALUE_KIND"); + ASSERT_EQ(schema->field(1)->name(), "_SEQUENCE_NUMBER"); + ASSERT_EQ(schema->field(2)->name(), "_REALTIME_OFFSET"); + ASSERT_EQ(schema->field(3)->name(), "key"); + ASSERT_EQ(schema->field(4)->name(), "value"); + ASSERT_FALSE(schema->field(0)->nullable()); + ASSERT_FALSE(schema->field(1)->nullable()); + ASSERT_EQ(schema->field(2)->nullable(), SpecialFields::RealtimeOffset().Nullable()); + ASSERT_FALSE(schema->field(3)->nullable()); + ASSERT_TRUE(schema->field(4)->nullable()); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestTransportSchemaValidation) { + const std::shared_ptr valid = MakeTransportSchema({}); + 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(RealtimePrimaryKeyLayout::ValidateSchema(arrow::schema(fields)), + "transport schema field"); + } +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsCommittedPrefix) { + 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 transport_schema = MakeTransportSchema(value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ + [0, 100, 0, 1, 10], + [0, 101, 1, 2, 20], + [0, 102, 2, 4, 40], + [0, 103, 3, 6, 60] + ])") + .ValueOrDie()); + + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(transport_array, transport_type, 2)); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(2, 4), + key_schema, value_schema, pool_)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(readers[0].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, {102, 103}, levels, {{4}, {6}}, {{4, 40}, {6, 60}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsNegativeOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, -1, 1]])") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(1, 2), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "reader offset must be non-negative"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryOffsetCoverageAcrossReadersAndBatches) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 2, 1], [0, 11, 0, 2]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 12, 3, 3], [0, 13, 1, 4]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 4), + value_schema, value_schema, pool_)); + int64_t row_count = 0; + for (const std::unique_ptr& reader : readers) { + ASSERT_OK_AND_ASSIGN( + std::vector rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(reader.get()))); + row_count += static_cast(rows.size()); + } + ASSERT_EQ(4, row_count); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsMissingVisibleOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 2, 2]])") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(0, 3), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsDuplicateVisibleOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 11, 1, 2], [0, 12, 1, 3]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector first_rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(readers[0].get()))); + ASSERT_EQ(1, first_rows.size()); + ASSERT_NOK_WITH_MSG((ReadResultCollector::CollectKeyValueResult( + readers[1].get())), + "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsEmptyEofForVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([])").ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(0, 1), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsEmptyReadersForVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_NOK_WITH_MSG( + RealtimePrimaryKeyReaderFactory::CreateForQuery(std::move(batch_readers), transport_schema, + OffsetRange(0, 1), value_schema, + value_schema, pool_), + "PK real-time store returned no query readers for a non-empty visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsEmptyReadersForEmptyVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(1, 1), + value_schema, value_schema, pool_)); + ASSERT_TRUE(readers.empty()); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryBitmapBounds) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1]])") + .ValueOrDie(); + auto batch_reader = std::make_unique( + std::make_unique(transport_array, transport_type, /*batch_size=*/1), + /*row_id=*/1); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + Result> result = + ReadResultCollector::CollectKeyValueResult(reader.get()); + ASSERT_TRUE(result.status().IsInvalid()); + ASSERT_NOK_WITH_MSG(result, + "selected row id 1 is out of bounds for realtime primary-key transport " + "batch length 1"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsPartialBitmap) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 1, 2]])") + .ValueOrDie(); + RoaringBitmap32 partial_bitmap; + partial_bitmap.Add(0); + auto batch_reader = std::make_unique( + transport_array, transport_type, partial_bitmap, /*read_batch_size=*/2); + batch_reader->EnableRandomizeBatchSize(false); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "must cover every raw transport row"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitRejectsPartialBitmap) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 1, 2]])") + .ValueOrDie(); + RoaringBitmap32 partial_bitmap; + partial_bitmap.Add(0); + auto batch_reader = std::make_unique( + transport_array, transport_type, partial_bitmap, /*read_batch_size=*/2); + batch_reader->EnableRandomizeBatchSize(false); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyCommitReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "must cover every raw transport row"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryProjection) { + 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 transport_schema = MakeTransportSchema({key, extra}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1, 2]])") + .ValueOrDie()); + + auto query_batch_reader = + std::make_unique(transport_array, transport_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr query_reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(query_batch_reader), transport_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); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitOffsetCoverage) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 2, 1], [0, 11, 0, 3]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 12, 1, 2], [0, 13, 3, 4]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 4), + value_schema, value_schema, pool_)); + int64_t row_count = 0; + for (const std::unique_ptr& reader : readers) { + ASSERT_OK_AND_ASSIGN( + std::vector rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(reader.get()))); + row_count += static_cast(rows.size()); + } + ASSERT_EQ(4, row_count); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitRejectsEmptyReaders) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_), + "PK real-time store returned no commit readers for a sealed segment"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestRejectsDuplicateCommitOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON( + transport_type, R"([[0, 10, 0, 1], [0, 11, 0, 2], [0, 12, 2, 3]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back(std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 3), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG((ReadResultCollector::CollectKeyValueResult( + readers[0].get())), + "did not cover the sealed range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, 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 transport_schema = MakeTransportSchema({key, value}); + std::shared_ptr actual_schema = MakeTransportSchema({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, + CreateRealtimePrimaryKeyCommitReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + arrow::schema({key}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "field count"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestSafeDecode) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + + arrow::FieldVector invalid_fields = transport_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, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "transport batch field"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestNestedValues) { + std::shared_ptr id = MakeField("id", arrow::int32(), 0); + std::shared_ptr key_schema = arrow::schema({id}); + 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}); + std::shared_ptr transport_schema = + MakeTransportSchema(query_value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON( + transport_type, + R"([[0, 10, 0, 1, [[200, 100], [400, 300]], [["k1", [8, 7]], ["k2", [10, 9]]], [[[12, 11], 13], [[22, 21], 23]]]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(transport_array, transport_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_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(RealtimePrimaryKeyReaderTest, TestFactoryFailureClosesReaders) { + 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 transport_schema = MakeTransportSchema(value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ + [0, 10, 0, 1, 100] + ])") + .ValueOrDie()); + + int32_t factory_failure_close_count = 0; + std::vector> batch_readers; + batch_readers.push_back(std::make_unique( + std::make_unique(transport_array, transport_type, 1), + &factory_failure_close_count)); + batch_readers.push_back(nullptr); + ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_), + "PK real-time store returned a null query reader"); + ASSERT_EQ(factory_failure_close_count, 1); +} + +} // namespace paimon::test 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 000000000..61cd14009 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,293 @@ +/* + * 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/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/core_options.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/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/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/core/utils/primary_key_table_utils.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +Result> CreateRealtimePrimaryKeyTransportBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& transport_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 transport, + arrow::StructArray::Make(std::move(columns), transport_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(transport), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(transport), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + return checked_pointer_cast(sorted.make_array()); +} + +} // namespace + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + 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 (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + 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)); + } + 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, + transport_schema, arrow::schema(std::move(key_fields)), trimmed_primary_keys, + key_comparator, options, 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& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + 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), + transport_schema_(transport_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + options_(options), + 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 transport, + CreateRealtimePrimaryKeyTransportBatch(std::move(batch), write_schema_, transport_schema_, + trimmed_primary_keys_, first_sequence, next_offset_, + arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport, output.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(transport_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_)); + 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> realtime_primary_key_readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit(std::move(readers), transport_schema_, + sealed_offsets, key_schema_, write_schema_, + memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(realtime_primary_key_readers.size()); + for (std::unique_ptr& realtime_primary_key_reader : + realtime_primary_key_readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + write_schema_, trimmed_primary_keys_, options_, memory_pool_)); + sorted_readers.push_back(std::make_unique( + std::move(realtime_primary_key_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); + } + return merge_tree_writer_->WriteSortedReadersToFiles(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 000000000..cdd3d889f --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,105 @@ +/* + * 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 + +#include "paimon/core/core_options.h" +#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::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + 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& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const CoreOptions& options, 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 transport_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; + CoreOptions options_; + 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 ec37cfed4..ded060989 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 f78e15506..0db1145f7 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 0363dff6a..ed56b6370 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 7b8947ea1..90f508b47 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" @@ -51,6 +52,7 @@ #include "paimon/core/schema/table_schema.h" #include "paimon/core/table/bucket_mode.h" #include "paimon/defs.h" +#include "paimon/format/file_format.h" #include "paimon/result.h" namespace paimon { @@ -98,6 +100,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::EqualsIgnoreCase(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) { @@ -153,14 +164,8 @@ Status SchemaValidation::ValidateTableSchema(const TableSchema& schema) { PAIMON_RETURN_NOT_OK(ValidateFieldsPrefix(schema, options)); PAIMON_RETURN_NOT_OK(ValidateSequenceField(schema, options)); PAIMON_RETURN_NOT_OK(ValidateSequenceGroup(schema, options)); + PAIMON_RETURN_NOT_OK(ValidateChangelogProducer(schema, options)); - ChangelogProducer changelog_producer = options.GetChangelogProducer(); - if (schema.PrimaryKeys().empty() && changelog_producer != ChangelogProducer::NONE) { - return Status::Invalid( - fmt::format("Can not set {} on table without primary keys, please define primary keys.", - Options::CHANGELOG_PRODUCER)); - } - PAIMON_RETURN_NOT_OK(ValidateChangelogProducer(options)); PAIMON_RETURN_NOT_OK(Preconditions::CheckState( options.GetExpireConfig().GetSnapshotRetainMin() > 0, std::string(Options::SNAPSHOT_NUM_RETAINED_MIN) + " should be at least 1")); @@ -187,7 +192,9 @@ Status SchemaValidation::ValidateTableSchema(const TableSchema& schema) { PAIMON_RETURN_NOT_OK(ValidateRowTracking(schema, options)); PAIMON_RETURN_NOT_OK(ValidateBlobFields(schema, options)); + PAIMON_RETURN_NOT_OK(ValidateMosaicDataFields(schema, options)); PAIMON_RETURN_NOT_OK(ValidateMapStorageLayout(schema, options)); + PAIMON_RETURN_NOT_OK(ValidateVectorFields(schema, options)); return Status::OK(); } @@ -312,10 +319,41 @@ Status SchemaValidation::ValidateBucket(const TableSchema& schema, const CoreOpt return Status::OK(); } -Status SchemaValidation::ValidateChangelogProducer(const CoreOptions& options) { - return Preconditions::CheckState(options.GetChangelogProducer() == ChangelogProducer::NONE, - "C++ Paimon does not support changelog-producer yet. Please " - "keep changelog-producer as 'none'."); +Status SchemaValidation::ValidateChangelogProducer(const TableSchema& schema, + const CoreOptions& options) { + ChangelogProducer changelog_producer = options.GetChangelogProducer(); + if (schema.PrimaryKeys().empty() && changelog_producer != ChangelogProducer::NONE) { + return Status::Invalid( + fmt::format("Can not set {} on table without primary keys, please define primary keys.", + Options::CHANGELOG_PRODUCER)); + } + + bool row_deduplicate = options.ChangelogRowDeduplicate(); + const std::vector& ignore_fields = + options.GetChangelogRowDeduplicateIgnoreFields(); + PAIMON_RETURN_NOT_OK(Preconditions::CheckState( + ignore_fields.empty() || row_deduplicate, "'{}' is only valid when '{}' is true.", + Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, + Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE)); + PAIMON_RETURN_NOT_OK(Preconditions::CheckState( + ObjectUtils::ContainsAll(schema.FieldNames(), ignore_fields), + "Fields {} configured in '{}' can not be found in table schema.", ignore_fields, + Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS)); + PAIMON_RETURN_NOT_OK(Preconditions::CheckState( + !row_deduplicate || changelog_producer == ChangelogProducer::LOOKUP || + changelog_producer == ChangelogProducer::FULL_COMPACTION, + "'{}' is only valid for 'lookup' or 'full-compaction' changelog producer.", + Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE)); + PAIMON_RETURN_NOT_OK(Preconditions::CheckState( + changelog_producer == ChangelogProducer::NONE || + changelog_producer == ChangelogProducer::INPUT || + changelog_producer == ChangelogProducer::LOOKUP, + "C++ Paimon only supports 'none', 'input' and 'lookup' changelog-producer now.")); + return Preconditions::CheckState( + options.GetMergeEngine() != MergeEngine::FIRST_ROW || + changelog_producer == ChangelogProducer::NONE || + changelog_producer == ChangelogProducer::LOOKUP, + "Only support 'none' and 'lookup' changelog-producer on FIRST_ROW merge engine"); } Status SchemaValidation::ValidateForDeletionVectors(const CoreOptions& options) { @@ -550,6 +588,76 @@ Status SchemaValidation::ValidateBlobFields(const TableSchema& schema, const Cor return Status::OK(); } +Status SchemaValidation::ValidateMosaicDataField(const std::shared_ptr& field) { + if (VariantTypeUtils::IsVariantField(field)) { + return Status::Invalid("Mosaic file format does not support type VARIANT"); + } + if (BlobUtils::IsBlobField(field)) { + return Status::Invalid("Mosaic file format does not support type BLOB"); + } + + const std::shared_ptr& type = field->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::DATE32: + case arrow::Type::STRING: + case arrow::Type::BINARY: + case arrow::Type::TIME32: + case arrow::Type::DECIMAL128: + return Status::OK(); + case arrow::Type::TIMESTAMP: { + const auto& timestamp_type = checked_cast(*type); + if (timestamp_type.unit() == arrow::TimeUnit::SECOND) { + return Status::Invalid("Mosaic file format does not support TIMESTAMP(0)"); + } + return Status::OK(); + } + case arrow::Type::LIST: + return ValidateMosaicDataField(type->field(0)); + case arrow::Type::MAP: { + const auto& map_type = checked_cast(*type); + PAIMON_RETURN_NOT_OK(ValidateMosaicDataField(map_type.key_field())); + return ValidateMosaicDataField(map_type.item_field()); + } + case arrow::Type::FIXED_SIZE_LIST: + return Status::Invalid("Mosaic file format does not support type VECTOR"); + case arrow::Type::STRUCT: + return Status::Invalid("Mosaic file format does not support type ROW"); + default: + break; + } + return Status::Invalid( + fmt::format("Mosaic file format does not support type {}", type->ToString())); +} + +Status SchemaValidation::ValidateMosaicDataFields(const TableSchema& schema, + const CoreOptions& options) { + if (StringUtils::ToLowerCase(options.GetFileFormat()->Identifier()) != "mosaic") { + return Status::OK(); + } + + const std::vector inline_blob_fields = options.GetBlobInlineFields(); + const std::set inline_blob_field_set(inline_blob_fields.begin(), + inline_blob_fields.end()); + // Match Java SchemaValidation by validating only fields stored in the normal data file. C++ + // permits BLOB only as a top-level field; descriptor and view fields are inline, so Mosaic + // must reject them here. + for (const DataField& field : schema.Fields()) { + if (BlobUtils::IsBlobField(field.ArrowField()) && + inline_blob_field_set.count(field.Name()) == 0) { + continue; + } + PAIMON_RETURN_NOT_OK(ValidateMosaicDataField(field.ArrowField())); + } + return Status::OK(); +} + Status SchemaValidation::ValidateMapStorageLayout(const TableSchema& schema, const CoreOptions& options) { // Extract all field names that have map.storage-layout configured from options @@ -622,6 +730,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 @@ -640,12 +751,48 @@ Status SchemaValidation::ValidateMapStorageLayout(const TableSchema& schema, options.GetFileFormat()->Identifier())); PAIMON_RETURN_NOT_OK(ValidatePerLevelOption(options_map, Options::FILE_FORMAT_PER_LEVEL, ValidateSharedShreddingFileFormat)); + std::shared_ptr changelog_format = options.GetChangelogFileFormat(); + if (changelog_format) { + PAIMON_RETURN_NOT_OK(ValidateSharedShreddingFileFormat(Options::CHANGELOG_FILE_FORMAT, + changelog_format->Identifier())); + } PAIMON_RETURN_NOT_OK(ValidateSharedShreddingCompression(Options::FILE_COMPRESSION, options.GetFileCompression())); PAIMON_RETURN_NOT_OK(ValidatePerLevelOption(options_map, Options::FILE_COMPRESSION_PER_LEVEL, ValidateSharedShreddingCompression)); + std::optional changelog_compression = options.GetChangelogFileCompression(); + if (changelog_compression) { + PAIMON_RETURN_NOT_OK(ValidateSharedShreddingCompression(Options::CHANGELOG_FILE_COMPRESSION, + changelog_compression.value())); + } 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 613372ff8..388ab42e2 100644 --- a/src/paimon/core/schema/schema_validation.h +++ b/src/paimon/core/schema/schema_validation.h @@ -66,15 +66,21 @@ class SchemaValidation { static Status ValidateFieldsPrefix(const TableSchema& schema, const CoreOptions& options); static Status ValidateSequenceField(const TableSchema& schema, const CoreOptions& options); static Status ValidateSequenceGroup(const TableSchema& schema, const CoreOptions& options); - static Status ValidateChangelogProducer(const CoreOptions& options); + static Status ValidateChangelogProducer(const TableSchema& schema, const CoreOptions& options); static Status ValidateForDeletionVectors(const CoreOptions& options); static Status ValidateRowTracking(const TableSchema& table_schema, const CoreOptions& options); static Status ValidateBlobFields(const TableSchema& schema, const CoreOptions& options); + static Status ValidateMosaicDataFields(const TableSchema& schema, const CoreOptions& options); + + static Status ValidateMosaicDataField(const std::shared_ptr& field); + 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 4c878113c..2137970a3 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -24,6 +24,7 @@ #include "arrow/api.h" #include "gtest/gtest.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/core/schema/table_schema.h" #include "paimon/defs.h" #include "paimon/testing/utils/testharness.h" @@ -46,6 +47,152 @@ TEST(SchemaValidationTest, TestSimple) { ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestRealtimeOffsetIsGloballyReserved) { + auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(0, schema, {}, {}, {})); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "field name '_REALTIME_OFFSET' in schema cannot be special field"); +} + +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)); + + parquet_options[Options::FILE_FORMAT] = "PARQUET"; + ASSERT_OK_AND_ASSIGN(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."); +} + +#ifdef PAIMON_ENABLE_MOSAIC +TEST(SchemaValidationTest, TestMosaicDataTypes) { + std::map options = {{Options::BUCKET, "-1"}, + {Options::FILE_FORMAT, "mosaic"}}; + arrow::FieldVector supported_fields = { + arrow::field("f0", arrow::boolean()), + arrow::field("f1", arrow::int8()), + arrow::field("f2", arrow::int16()), + arrow::field("f3", arrow::int32()), + arrow::field("f4", arrow::int64()), + arrow::field("f5", arrow::float32()), + arrow::field("f6", arrow::float64()), + arrow::field("f7", arrow::utf8()), + arrow::field("f8", arrow::binary()), + arrow::field("f9", arrow::date32()), + arrow::field("f10", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("f11", arrow::decimal128(38, 2)), + arrow::field("f12", arrow::list(arrow::float32())), + arrow::field("f13", arrow::map(arrow::int8(), arrow::int16())), + }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, arrow::schema(supported_fields), + /*partition_keys=*/{}, /*primary_keys=*/{}, options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); + + arrow::FieldVector unsupported_fields = { + arrow::field("row", arrow::struct_({arrow::field("value", arrow::int32())})), + VariantTypeUtils::ToArrowField("variant"), + arrow::field("vector", arrow::fixed_size_list(arrow::float32(), 3)), + arrow::field("timestamp", arrow::timestamp(arrow::TimeUnit::SECOND)), + arrow::field("nested_row", + arrow::list(arrow::struct_({arrow::field("value", arrow::int32())}))), + }; + std::vector expected_errors = {"type ROW", "type VARIANT", "type VECTOR", + "TIMESTAMP(0)", "type ROW"}; + for (size_t i = 0; i < unsupported_fields.size(); ++i) { + SCOPED_TRACE("field=" + unsupported_fields[i]->name()); + ASSERT_OK_AND_ASSIGN( + table_schema, + TableSchema::Create(/*schema_id=*/0, arrow::schema({unsupported_fields[i]}), + /*partition_keys=*/{}, /*primary_keys=*/{}, options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + expected_errors[i]); + } + + std::shared_ptr blob_field = BlobUtils::ToArrowField("blob", false); + std::map blob_options = { + {Options::BUCKET, "-1"}, + {Options::FILE_FORMAT, "mosaic"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}, + }; + ASSERT_OK_AND_ASSIGN( + table_schema, + TableSchema::Create(/*schema_id=*/0, + arrow::schema({arrow::field("id", arrow::int32()), blob_field}), + /*partition_keys=*/{}, /*primary_keys=*/{}, blob_options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); + + blob_options[Options::BLOB_DESCRIPTOR_FIELD] = "blob"; + ASSERT_OK_AND_ASSIGN( + table_schema, + TableSchema::Create(/*schema_id=*/0, + arrow::schema({arrow::field("id", arrow::int32()), blob_field}), + /*partition_keys=*/{}, /*primary_keys=*/{}, blob_options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), "type BLOB"); +} +#endif + TEST(SchemaValidationTest, TestRowTracking) { auto f0 = arrow::field("f0", arrow::utf8()); auto f1 = arrow::field("f1", arrow::int32()); @@ -575,7 +722,18 @@ TEST(SchemaValidationTest, ValidateDeletionVector) { std::shared_ptr table_schema, TableSchema::Create(/*schema_id=*/0, schema, partition_keys, primary_keys, options)); ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), - "C++ Paimon does not support changelog-producer yet"); + "C++ Paimon only supports 'none', 'input' and 'lookup' " + "changelog-producer now"); + } + { + std::map options = {{Options::BUCKET, "2"}, + {Options::BUCKET_KEY, "f0"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {Options::CHANGELOG_PRODUCER, "input"}}; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, partition_keys, primary_keys, options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } { std::map options = {{Options::BUCKET, "2"}, @@ -727,6 +885,41 @@ TEST(SchemaValidationTest, ValidateInvalidConfiguration) { "Can not set changelog-producer on table without primary keys, please " "define primary keys."); } + { + std::map options = { + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, "true"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"f0"}, options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "'changelog-producer.row-deduplicate' is only valid for 'lookup' or " + "'full-compaction' changelog producer"); + } + { + std::map options = { + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, "f1"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"f0"}, options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "'changelog-producer.row-deduplicate-ignore-fields' is only valid when " + "'changelog-producer.row-deduplicate' is true"); + } + { + std::map options = { + {Options::CHANGELOG_PRODUCER, "input"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, "true"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, "missing"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"f0"}, options)); + ASSERT_NOK_WITH_MSG( + SchemaValidation::ValidateTableSchema(*table_schema), + "Fields [\"missing\"] configured in " + "'changelog-producer.row-deduplicate-ignore-fields' can not be found in table schema"); + } { auto invalid_field = arrow::field("_SEQUENCE_NUMBER", arrow::int64()); arrow::FieldVector invalid_fields = fields; @@ -758,15 +951,18 @@ TEST(SchemaValidationTest, ValidateInvalidConfiguration) { TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, /*primary_keys=*/{"f0"}, options)); ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), - "C++ Paimon does not support changelog-producer yet"); + "Only support 'none' and 'lookup' changelog-producer on FIRST_ROW " + "merge engine"); } { - std::map options = {{Options::CHANGELOG_PRODUCER, "lookup"}}; + std::map options = { + {Options::CHANGELOG_PRODUCER, "full-compaction"}}; ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, /*primary_keys=*/{"f0"}, options)); - ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), - "C++ Paimon does not support changelog-producer yet"); + ASSERT_NOK_WITH_MSG( + SchemaValidation::ValidateTableSchema(*table_schema), + "C++ Paimon only supports 'none', 'input' and 'lookup' changelog-producer now."); } // test for row tracking { @@ -1043,6 +1239,16 @@ TEST(SchemaValidationTest, TestMapSharedShreddingCompression) { "MAP shared-shredding only supports none/lz4/zstd compression, but " "file.compression.per.level.1 is snappy."); } + { + auto options = base_options; + options[Options::CHANGELOG_FILE_COMPRESSION] = "snappy"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "MAP shared-shredding only supports none/lz4/zstd compression, but " + "changelog-file.compression is snappy."); + } { auto options = base_options; options.erase("fields.f1.map.storage-layout"); @@ -1094,6 +1300,16 @@ TEST(SchemaValidationTest, TestMapSharedShreddingFileFormat) { "MAP shared-shredding only supports parquet/orc file formats, but " "file.format.per.level.1 is avro."); } + { + auto options = base_options; + options[Options::CHANGELOG_FILE_FORMAT] = "avro"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "MAP shared-shredding only supports parquet/orc file formats, but " + "changelog-file.format is avro."); + } } TEST(SchemaValidationTest, TestMapSharedShreddingRejectsPostponeBucketMode) { diff --git a/src/paimon/core/schema/table_schema.cpp b/src/paimon/core/schema/table_schema.cpp index d7be7f115..6e2d8747e 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/stats/simple_stats_converter.h b/src/paimon/core/stats/simple_stats_converter.h index 903cf6e08..91eb04937 100644 --- a/src/paimon/core/stats/simple_stats_converter.h +++ b/src/paimon/core/stats/simple_stats_converter.h @@ -24,6 +24,7 @@ #include "paimon/format/column_stats.h" #include "paimon/result.h" +#include "paimon/visibility.h" namespace paimon { @@ -31,7 +32,7 @@ class SimpleStats; class ColumnStats; class MemoryPool; -class SimpleStatsConverter { +class PAIMON_EXPORT SimpleStatsConverter { public: static Result ToBinary(const std::vector>& stats, MemoryPool* pool); diff --git a/src/paimon/core/table/bucket_mode.cpp b/src/paimon/core/table/bucket_mode.cpp index d46739f89..429a787d2 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 f87ab3467..e6e6491c6 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 2a77e5ea9..8feeb6e8c 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/sink/commit_message_serializer.cpp b/src/paimon/core/table/sink/commit_message_serializer.cpp index 1b0645f4e..e3424d4a2 100644 --- a/src/paimon/core/table/sink/commit_message_serializer.cpp +++ b/src/paimon/core/table/sink/commit_message_serializer.cpp @@ -38,6 +38,7 @@ #include "paimon/core/io/data_file_meta_12_serializer.h" #include "paimon/core/io/data_file_meta_first_row_id_legacy_serializer.h" #include "paimon/core/io/data_file_meta_serializer.h" +#include "paimon/core/io/data_file_meta_write_cols_legacy_serializer.h" #include "paimon/core/io/data_increment.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/utils/object_serializer.h" @@ -46,7 +47,7 @@ namespace paimon { class MemoryPool; -const int32_t CommitMessageSerializer::CURRENT_VERSION = 12; +const int32_t CommitMessageSerializer::CURRENT_VERSION = 13; CommitMessageSerializer::CommitMessageSerializer(const std::shared_ptr& pool) : memory_pool_(pool), @@ -204,16 +205,25 @@ Result> CommitMessageSerializer::Deserialize(int3 DataInputStream* in) { if (version == CURRENT_VERSION) { return Deserialize(version, data_file_serializer_.get(), index_entry_serializer_.get(), in); + } else if (version == 12) { + auto data_file_meta_write_cols_legacy_serializer = + std::make_unique(memory_pool_); + return Deserialize(version, data_file_meta_write_cols_legacy_serializer.get(), + index_entry_serializer_.get(), in); } else if (version == 11) { + auto data_file_meta_write_cols_legacy_serializer = + std::make_unique(memory_pool_); auto index_entry_v4_deserializer = std::make_unique(memory_pool_); - return Deserialize(version, data_file_serializer_.get(), index_entry_v4_deserializer.get(), - in); + return Deserialize(version, data_file_meta_write_cols_legacy_serializer.get(), + index_entry_v4_deserializer.get(), in); } else if (version == 9 || version == 10) { + auto data_file_meta_write_cols_legacy_serializer = + std::make_unique(memory_pool_); auto index_entry_v3_deserializer = std::make_unique(memory_pool_); - return Deserialize(version, data_file_serializer_.get(), index_entry_v3_deserializer.get(), - in); + return Deserialize(version, data_file_meta_write_cols_legacy_serializer.get(), + index_entry_v3_deserializer.get(), in); } else if (version == 8) { auto data_file_meta_first_row_id_legacy_serializer = std::make_unique(memory_pool_); diff --git a/src/paimon/core/table/sink/commit_message_test.cpp b/src/paimon/core/table/sink/commit_message_test.cpp index d602da227..f372361cb 100644 --- a/src/paimon/core/table/sink/commit_message_test.cpp +++ b/src/paimon/core/table/sink/commit_message_test.cpp @@ -65,6 +65,44 @@ TEST(CommitMessageTest, TestCurrentVersion) { ASSERT_EQ(CommitMessageSerializer::CURRENT_VERSION, CommitMessage::CurrentVersion()); } +TEST(CommitMessageTest, TestDeserializeVersion13GeneratedByJava) { + // Generated by CommitMessageSerializer::serialize from Apache Paimon Java master. + std::string data_path = paimon::test::GetDataDir() + "/compatibility/commit_message-v13"; + auto file_system = std::make_shared(); + auto buffer_length = file_system->GetFileStatus(data_path).value().GetLen(); + ASSERT_GT(buffer_length, 0); + + std::vector buffer(buffer_length, 0); + ASSERT_OK_AND_ASSIGN(auto in_stream, file_system->Open(data_path)); + ASSERT_OK(in_stream->Read(buffer.data(), buffer.size())); + ASSERT_OK(in_stream->Close()); + + auto pool = GetDefaultPool(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + CommitMessage::Deserialize(CommitMessage::CurrentVersion(), buffer.data(), + buffer.size(), pool)); + auto result_message = std::dynamic_pointer_cast(result); + ASSERT_NE(result_message, nullptr); + ASSERT_EQ(result_message->Partition(), + BinaryRowGenerator::GenerateRow({std::string("aaaaa")}, pool.get())); + ASSERT_EQ(result_message->Bucket(), 20); + ASSERT_EQ(result_message->TotalBuckets(), std::optional(32)); + ASSERT_EQ(result_message->GetNewFilesIncrement().NewFiles().size(), 1); + ASSERT_TRUE(result_message->GetCompactIncrement().IsEmpty()); + + const std::shared_ptr& data_file = + result_message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(data_file->file_name, "my_file"); + ASSERT_TRUE(data_file->write_cols.has_value()); + ASSERT_EQ(data_file->write_cols.value(), (std::vector{"a", "b", "c", "f"})); + ASSERT_TRUE(data_file->column_max_sequence_numbers.has_value()); + ASSERT_EQ(data_file->column_max_sequence_numbers.value(), + (std::vector{15, 100, 150, 200})); + + ASSERT_OK_AND_ASSIGN(std::string serialized_bytes, CommitMessage::Serialize(result, pool)); + ASSERT_EQ(serialized_bytes, std::string(buffer.data(), buffer.size())); +} + TEST(CommitMessageTest, TestCompatibleWithVersion12) { // index file meta: add global index meta source meta int32_t version = 12; @@ -95,7 +133,13 @@ TEST(CommitMessageTest, TestCompatibleWithVersion12) { // check result ASSERT_OK_AND_ASSIGN(std::string serialized_bytes, CommitMessage::Serialize(ret, pool)); - ASSERT_EQ(serialized_bytes, std::string(reinterpret_cast(buffer.data()), buffer.size())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr roundtrip, + CommitMessage::Deserialize(CommitMessage::CurrentVersion(), serialized_bytes.data(), + serialized_bytes.size(), pool)); + auto roundtrip_message = std::dynamic_pointer_cast(roundtrip); + ASSERT_NE(roundtrip_message, nullptr); + ASSERT_EQ(*roundtrip_message, *res_msg); } TEST(CommitMessageTest, TestCompatibleWithVersion11) { @@ -185,7 +229,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion10) { /*creation_time=*/Timestamp(1761242383412ll, 0), /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); LinkedHashMap dv_ranges; dv_ranges.insert_or_assign( @@ -249,7 +294,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion9) { /*creation_time=*/Timestamp(1757349273600ll, 0), /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); LinkedHashMap dv_ranges; dv_ranges.insert_or_assign( @@ -316,7 +362,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion9WithExternalPathForIndex) { /*value_stats_cols=*/std::nullopt, /*external_path=*/ "FILE:/tmp/external/f1=10/bucket-1/data-72b62a5f-d698-4db5-b51a-04c0dc027702-1.orc", - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); LinkedHashMap dv_ranges; dv_ranges.insert_or_assign( @@ -379,7 +426,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion8) { /*creation_time=*/Timestamp(1754068646844ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); expected_msgs.emplace_back(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*total_bucket=*/2, DataIncrement({file_meta}, {}, {}), CompactIncrement({}, {}, {})); @@ -398,7 +446,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion8) { /*creation_time=*/Timestamp(1754068646864ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); expected_msgs.emplace_back(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/2, DataIncrement({file_meta2}, {}, {}), CompactIncrement({}, {}, {})); @@ -448,7 +497,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion7) { /*creation_time=*/Timestamp(1743525392885ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); expected_msgs.emplace_back(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*total_bucket=*/2, DataIncrement({file_meta}, {}, {}), CompactIncrement({}, {}, {})); @@ -469,7 +519,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion7) { /*creation_time=*/Timestamp(1743525392921ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); expected_msgs.emplace_back(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/2, DataIncrement({file_meta2}, {}, {}), CompactIncrement({}, {}, {})); @@ -515,7 +566,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion6) { /*creation_time=*/Timestamp(1737052260143ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); expected_msgs.emplace_back(BinaryRow::EmptyRow(), /*bucket=*/0, /*total_bucket=*/std::nullopt, DataIncrement({file_meta}, {}, {}), CompactIncrement({}, {}, {})); // check result @@ -562,7 +614,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion5) { /*creation_time=*/Timestamp(1734707236040ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta1_after_compact = std::make_shared( "data-0d0f29cc-63c6-4fab-a594-71bd7d06fcde-1.orc", /*file_size=*/859, /*row_count=*/1, BinaryRowGenerator::GenerateRow({std::string("Alice"), 1}, pool.get()), @@ -576,7 +629,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion5) { /*creation_time=*/Timestamp(1734707236040ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}, {}, {}); LinkedHashMap dv_metas1; @@ -605,7 +659,8 @@ TEST(CommitMessageTest, TestCompatibleWithVersion5) { /*creation_time=*/Timestamp(1734707236109ll, 0), /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment2({file_meta2}, {}, {}, {}, {}); LinkedHashMap dv_metas2; dv_metas2.insert_or_assign( @@ -663,7 +718,7 @@ TEST(CommitMessageTest, TestCompatibleWithVersion4) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment1, @@ -682,7 +737,7 @@ TEST(CommitMessageTest, TestCompatibleWithVersion4) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment2({file_meta2}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/std::nullopt, data_increment2, @@ -701,7 +756,7 @@ TEST(CommitMessageTest, TestCompatibleWithVersion4) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment3({file_meta3}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({20}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment3, @@ -751,7 +806,8 @@ TEST(CommitMessageTest, TestCompatibleWithJavaPaimon10WithStatsDenseStore) { /*creation_time=*/Timestamp(1731412938869ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment1, @@ -768,7 +824,8 @@ TEST(CommitMessageTest, TestCompatibleWithJavaPaimon10WithStatsDenseStore) { /*creation_time=*/Timestamp(1731412938891ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment2({file_meta2}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/std::nullopt, data_increment2, @@ -785,7 +842,8 @@ TEST(CommitMessageTest, TestCompatibleWithJavaPaimon10WithStatsDenseStore) { /*creation_time=*/Timestamp(1731412938908ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment3({file_meta3}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({20}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment3, @@ -835,7 +893,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon1) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment1, @@ -854,7 +912,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon1) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment2({file_meta2}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/std::nullopt, data_increment2, @@ -873,7 +931,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon1) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment3({file_meta3}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({20}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment3, @@ -923,7 +981,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon2) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/std::nullopt, data_increment1, @@ -943,7 +1001,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon2) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment2({file_meta2}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({20}, pool.get()), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment2, @@ -993,7 +1051,7 @@ TEST(CommitMessageTest, TestCompatibleWith09JavaPaimon3) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*total_bucket=*/std::nullopt, data_increment1, @@ -1046,7 +1104,7 @@ TEST(CommitMessageTest, TestPkTableCompatibleWithJavaPaimon09) { /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta_with_level = std::make_shared( "data-2eb2a766-97e4-4fe4-88ce-eb606675c101-0.orc", /*file_size=*/789, /*row_count=*/1, /*min_key=*/BinaryRowGenerator::GenerateRow({std::string("Bob"), 0}, pool.get()), @@ -1063,7 +1121,7 @@ TEST(CommitMessageTest, TestPkTableCompatibleWithJavaPaimon09) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta}, {}, {}, {}, {}); LinkedHashMap dv_ranges; @@ -1141,7 +1199,7 @@ TEST(CommitMessageTest, TestCompatibleWithComplexDataType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment1({file_meta1}, {}, {}); expected_msgs.emplace_back(BinaryRow::EmptyRow(), /*bucket=*/0, /*total_bucket=*/std::nullopt, data_increment1, CompactIncrement({}, {}, {})); diff --git a/src/paimon/core/table/source/abstract_table_scan.h b/src/paimon/core/table/source/abstract_table_scan.h index 363fe0dce..f8dd15aac 100644 --- a/src/paimon/core/table/source/abstract_table_scan.h +++ b/src/paimon/core/table/source/abstract_table_scan.h @@ -39,6 +39,10 @@ class AbstractTableScan : public TableScan { const std::shared_ptr& snapshot_reader) : core_options_(core_options), snapshot_reader_(snapshot_reader) {} + std::shared_ptr GetMetrics() const override { + return snapshot_reader_->GetMetrics(); + } + protected: Result> CreateStartingScanner(bool is_streaming) const { const auto& snapshot_manager = snapshot_reader_->GetSnapshotManager(); diff --git a/src/paimon/core/table/source/append_count_reader.cpp b/src/paimon/core/table/source/append_count_reader.cpp index 8f9af90ac..5d684af67 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 20786c9f0..12b3823a7 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/data_evolution_batch_scan_test.cpp b/src/paimon/core/table/source/data_evolution_batch_scan_test.cpp index 14fae3972..5fbd5f093 100644 --- a/src/paimon/core/table/source/data_evolution_batch_scan_test.cpp +++ b/src/paimon/core/table/source/data_evolution_batch_scan_test.cpp @@ -49,7 +49,7 @@ std::shared_ptr NewAppendFile(const std::string& file_name, int64_ /*max_sequence_number=*/first_row_id + row_count - 1, /*schema_id=*/0, /*level=*/0, std::vector>(), Timestamp(0l, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), std::nullopt, std::nullopt, first_row_id, - std::nullopt); + std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr NewDataPlan(std::vector> files) { diff --git a/src/paimon/core/table/source/data_split_impl.cpp b/src/paimon/core/table/source/data_split_impl.cpp index 5b9128a73..be0ef4030 100644 --- a/src/paimon/core/table/source/data_split_impl.cpp +++ b/src/paimon/core/table/source/data_split_impl.cpp @@ -223,6 +223,8 @@ DataSplitImpl::GetFileMetaSerializer(int32_t version, const std::shared_ptr(pool); } else if (version == 7) { return std::make_unique(pool); + } else if (version == 8) { + return std::make_unique(pool); } else if (version == VERSION) { return std::make_unique(pool); } else { diff --git a/src/paimon/core/table/source/data_split_impl.h b/src/paimon/core/table/source/data_split_impl.h index 01c1231c8..2d357a163 100644 --- a/src/paimon/core/table/source/data_split_impl.h +++ b/src/paimon/core/table/source/data_split_impl.h @@ -35,6 +35,7 @@ #include "paimon/core/io/data_file_meta_12_serializer.h" #include "paimon/core/io/data_file_meta_first_row_id_legacy_serializer.h" #include "paimon/core/io/data_file_meta_serializer.h" +#include "paimon/core/io/data_file_meta_write_cols_legacy_serializer.h" #include "paimon/core/table/source/deletion_file.h" #include "paimon/table/source/data_split.h" @@ -44,7 +45,7 @@ namespace paimon { class DataSplitImpl : public DataSplit { public: static constexpr int64_t MAGIC = -2394839472490812314L; - static constexpr int32_t VERSION = 8; + static constexpr int32_t VERSION = 9; int64_t SnapshotId() const { return snapshot_id_; diff --git a/src/paimon/core/table/source/data_split_test.cpp b/src/paimon/core/table/source/data_split_test.cpp index b65988526..bb631be9a 100644 --- a/src/paimon/core/table/source/data_split_test.cpp +++ b/src/paimon/core/table/source/data_split_test.cpp @@ -46,6 +46,47 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +TEST(DataSplitTest, TestDeserializeVersion9GeneratedByJava) { + // Generated by DataSplit::serialize from Apache Paimon Java master. + std::string file_name = paimon::test::GetDataDir() + "/compatibility/data_split-v9"; + auto file_system = std::make_unique(); + + ASSERT_OK_AND_ASSIGN(auto input_stream, file_system->Open(file_name)); + std::vector split_bytes(input_stream->Length().value_or(0), 0); + ASSERT_GT(split_bytes.size(), 0); + ASSERT_OK(input_stream->Read(split_bytes.data(), split_bytes.size())); + ASSERT_OK(input_stream->Close()); + + auto pool = GetDefaultPool(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + Split::Deserialize(split_bytes.data(), split_bytes.size(), pool)); + auto result_data_split = std::dynamic_pointer_cast(result); + ASSERT_NE(result_data_split, nullptr); + ASSERT_EQ(result_data_split->SnapshotId(), 18); + ASSERT_EQ(result_data_split->Partition(), + BinaryRowGenerator::GenerateRow({std::string("aaaaa")}, pool.get())); + ASSERT_EQ(result_data_split->Bucket(), 20); + ASSERT_EQ(result_data_split->BucketPath(), "my path"); + ASSERT_EQ(result_data_split->TotalBuckets(), std::optional(32)); + ASSERT_EQ(result_data_split->DataFiles().size(), 1); + ASSERT_EQ( + result_data_split->DeletionFiles(), + (std::vector>{DeletionFile("deletion_file", 100, 22, 33)})); + ASSERT_FALSE(result_data_split->IsStreaming()); + ASSERT_FALSE(result_data_split->RawConvertible()); + + const std::shared_ptr& data_file = result_data_split->DataFiles()[0]; + ASSERT_EQ(data_file->file_name, "my_file"); + ASSERT_TRUE(data_file->write_cols.has_value()); + ASSERT_EQ(data_file->write_cols.value(), (std::vector{"a", "b", "c", "f"})); + ASSERT_TRUE(data_file->column_max_sequence_numbers.has_value()); + ASSERT_EQ(data_file->column_max_sequence_numbers.value(), + (std::vector{15, 100, 150, 200})); + + ASSERT_OK_AND_ASSIGN(std::string serialized_bytes, Split::Serialize(result, pool)); + ASSERT_EQ(serialized_bytes, std::string(split_bytes.data(), split_bytes.size())); +} + TEST(DataSplitTest, TestDeserializeVersion8WithWriteColsAndExternalPath) { std::string file_name = paimon::test::GetDataDir() + "/orc/pk_dv_index_in_data_with_external.db/" @@ -86,7 +127,8 @@ TEST(DataSplitTest, TestDeserializeVersion8WithWriteColsAndExternalPath) { /*value_stats_cols=*/std::nullopt, /*external_path=*/ "FILE:/tmp/external/f1=10/bucket-1/data-72b62a5f-d698-4db5-b51a-04c0dc027702-0.orc", - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -108,7 +150,10 @@ TEST(DataSplitTest, TestDeserializeVersion8WithWriteColsAndExternalPath) { .value()); ASSERT_EQ(*result_data_split, *expected_data_split) << result_data_split->ToString(); ASSERT_OK_AND_ASSIGN(std::string serialize_bytes, Split::Serialize(result_data_split, pool)); - ASSERT_EQ(serialize_bytes, std::string(split_bytes.data(), split_bytes.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr roundtrip, + Split::Deserialize(serialize_bytes.data(), serialize_bytes.size(), pool)); + auto roundtrip_data_split = std::dynamic_pointer_cast(roundtrip); + ASSERT_EQ(*roundtrip_data_split, *expected_data_split) << roundtrip_data_split->ToString(); } TEST(DataSplitTest, TestDeserializeVersion8WithWriteCols) { @@ -149,7 +194,8 @@ TEST(DataSplitTest, TestDeserializeVersion8WithWriteCols) { /*creation_time=*/Timestamp(1757349273246ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -172,7 +218,10 @@ TEST(DataSplitTest, TestDeserializeVersion8WithWriteCols) { .value()); ASSERT_EQ(*result_data_split, *expected_data_split) << result_data_split->ToString(); ASSERT_OK_AND_ASSIGN(std::string serialize_bytes, Split::Serialize(result_data_split, pool)); - ASSERT_EQ(serialize_bytes, std::string(split_bytes.data(), split_bytes.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr roundtrip, + Split::Deserialize(serialize_bytes.data(), serialize_bytes.size(), pool)); + auto roundtrip_data_split = std::dynamic_pointer_cast(roundtrip); + ASSERT_EQ(*roundtrip_data_split, *expected_data_split) << roundtrip_data_split->ToString(); } TEST(DataSplitTest, TestDeserializeVersion7WithFirstRowId) { @@ -209,7 +258,8 @@ TEST(DataSplitTest, TestDeserializeVersion7WithFirstRowId) { /*creation_time=*/Timestamp(1754073518741ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/5, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/5, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), @@ -261,7 +311,8 @@ TEST(DataSplitTest, TestDeserializeVersion7WithNullFirstRowId) { /*creation_time=*/Timestamp(1754068646844ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -312,7 +363,8 @@ TEST(DataSplitTest, TestDeserializeVersion6PkWithTotalBuckets) { /*creation_time=*/Timestamp(1743525392885ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -358,7 +410,8 @@ TEST(DataSplitTest, TestDeserializeVersion5PkWithExternalPath) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/"file:/tmp/bucket-0/data-80110e15-97b5-4bcf-ac09-6ca2659a4950-0.orc", - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), @@ -407,7 +460,8 @@ TEST(DataSplitTest, TestDeserializeVersion5PkWithEmptyExternalPath) { /*creation_time=*/Timestamp(1737052260143ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRow::EmptyRow(), @@ -463,7 +517,8 @@ TEST(DataSplitTest, TestDeserializeVersion4PkWithSnapshot4WithDvCardinality) { /*creation_time=*/Timestamp(1734707235578ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*bucket_path=*/ @@ -516,7 +571,7 @@ TEST(DataSplitTest, TestDeserializeVersion3AppendWithSnapshot1) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*bucket_path=*/ "data/append_10.db/append_10/f1=10/bucket-1", {file_meta}); @@ -555,7 +610,8 @@ TEST(DataSplitTest, TestDeserializeVersion3AppendWithSnapshot1WithStatsDenseStor /*creation_time=*/Timestamp(1731412938891ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*bucket_path=*/ @@ -596,7 +652,7 @@ TEST(DataSplitTest, TestDeserializeAppendWithSnapshot1) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*bucket_path=*/ @@ -636,7 +692,7 @@ TEST(DataSplitTest, TestDeserializeAppendWithSnapshot3) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-b913a160-a4d1-4084-af2a-18333c35668e-0.orc", /*file_size=*/506, /*row_count=*/1, @@ -651,7 +707,7 @@ TEST(DataSplitTest, TestDeserializeAppendWithSnapshot3) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({20}, pool.get()), @@ -709,7 +765,7 @@ TEST(DataSplitTest, TestDeserializeAppendWithSnapshot5) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -759,7 +815,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot2) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*bucket_path=*/ @@ -808,7 +864,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot6OfSingleFile) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/1, /*bucket_path=*/ "data/pk_09.db/pk_09/f1=10/bucket-1", {file_meta}); @@ -866,7 +922,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot6OfMultiFiles) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1c7a85f1-55bd-424f-b503-34a33be0fb96-0.orc", /*file_size=*/1148, /*row_count=*/2, @@ -888,7 +944,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot6OfMultiFiles) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-8cdb8b8d-5830-4b3b-aa94-8a30c449277a-0.orc", /*file_size=*/810, /*row_count=*/1, @@ -907,7 +963,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot6OfMultiFiles) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -967,7 +1023,7 @@ TEST(DataSplitTest, TestDeserializePkWithSnapshot8) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(/*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*bucket_path=*/ "data/pk_09.db/pk_09/f1=10/bucket-0", {file_meta}); @@ -1018,7 +1074,7 @@ TEST(DataSplitTest, TestDeserializePk10WithSnapshot6) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-d6d370f3-242b-45c9-8739-44bf31b2b449-0.orc", /*file_size=*/924, /*row_count=*/1, /*min_key=*/BinaryRowGenerator::GenerateRow({52}, pool.get()), /*max_key=*/ @@ -1035,7 +1091,7 @@ TEST(DataSplitTest, TestDeserializePk10WithSnapshot6) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({1, 1}, pool.get()), /*bucket=*/0, @@ -1068,7 +1124,7 @@ TEST(DataSplitTest, TestPartialMergedRowCount) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/2, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1078,7 +1134,7 @@ TEST(DataSplitTest, TestPartialMergedRowCount) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1112,7 +1168,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountRawConvertibleWithoutDeletionFiles) /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/4, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1122,7 +1178,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountRawConvertibleWithoutDeletionFiles) /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1147,7 +1203,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountRawConvertibleWithCardinality) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/2, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1157,7 +1213,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountRawConvertibleWithCardinality) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-2.orc", /*file_size=*/100, /*row_count=*/3, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1167,7 +1223,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountRawConvertibleWithCardinality) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1202,7 +1258,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountMixedCardinalityReturnsNullopt) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/2, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1212,7 +1268,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountMixedCardinalityReturnsNullopt) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1245,7 +1301,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountUnknownDeleteRowCountDoesNotBlockRa /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/2, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1255,7 +1311,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountUnknownDeleteRowCountDoesNotBlockRa /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1284,7 +1340,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountFallsBackToDataEvolution) { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/100, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/5, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1294,7 +1350,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountFallsBackToDataEvolution) { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/100, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-2.orc", /*file_size=*/100, /*row_count=*/2, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1304,7 +1360,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountFallsBackToDataEvolution) { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/200, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1335,7 +1391,7 @@ TEST(DataSplitTest, TestDataEvolutionMergedRowCountSubtractsDeletionFileCardinal /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/100, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/5, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1345,7 +1401,7 @@ TEST(DataSplitTest, TestDataEvolutionMergedRowCountSubtractsDeletionFileCardinal /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/100, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1380,7 +1436,7 @@ TEST(DataSplitTest, TestSerializeDataEvolutionSplitWithDeletionFiles) { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/10, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1390,7 +1446,7 @@ TEST(DataSplitTest, TestSerializeDataEvolutionSplitWithDeletionFiles) { /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1438,7 +1494,7 @@ TEST(DataSplitTest, TestDataEvolutionMergedRowCountUnavailableWithoutCardinality /*delete_row_count=*/std::nullopt, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/100, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), @@ -1471,7 +1527,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountResolvesMissingCardinalityViaFactor /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-1.orc", /*file_size=*/100, /*row_count=*/5, DataFileMeta::EmptyMinKey(), DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), @@ -1481,7 +1537,7 @@ TEST(DataSplitTest, TestPartialMergedRowCountResolvesMissingCardinalityViaFactor /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({10}, pool.get()), diff --git a/src/paimon/core/table/source/data_table_stream_scan.cpp b/src/paimon/core/table/source/data_table_stream_scan.cpp index 4c6d4df11..1657c9c7b 100644 --- a/src/paimon/core/table/source/data_table_stream_scan.cpp +++ b/src/paimon/core/table/source/data_table_stream_scan.cpp @@ -26,6 +26,7 @@ #include "paimon/core/options/changelog_producer.h" #include "paimon/core/table/bucket_mode.h" #include "paimon/core/table/source/plan_impl.h" +#include "paimon/core/table/source/snapshot/changelog_follow_up_scanner.h" #include "paimon/core/table/source/snapshot/delta_follow_up_scanner.h" #include "paimon/core/table/source/snapshot/follow_up_scanner.h" #include "paimon/core/table/source/snapshot/snapshot_reader.h" @@ -55,10 +56,15 @@ Result> DataTableStreamScan::CreatePlan() { Result> DataTableStreamScan::TryFirstPlan() { std::shared_ptr scan_result; - if (core_options_.GetChangelogProducer() == ChangelogProducer::LOOKUP) { - return Status::NotImplemented("do not support lookup changelog producer"); - } else if (core_options_.GetChangelogProducer() == ChangelogProducer::FULL_COMPACTION) { + if (core_options_.GetChangelogProducer() == ChangelogProducer::FULL_COMPACTION) { return Status::NotImplemented("do not support full compaction changelog producer"); + } else if (core_options_.GetChangelogProducer() == ChangelogProducer::LOOKUP) { + // Level-0 files will be compacted later to produce changelog records. Exclude them from + // the initial full scan so that the same changes are not emitted both in the full phase + // and again in the incremental changelog phase. + snapshot_reader_->WithLevelFilter([](int32_t level) -> bool { return level > 0; }); + PAIMON_ASSIGN_OR_RAISE(scan_result, starting_scanner_->Scan(snapshot_reader_)); + snapshot_reader_->WithLevelFilter([](int32_t) -> bool { return true; }); } else { PAIMON_ASSIGN_OR_RAISE(scan_result, starting_scanner_->Scan(snapshot_reader_)); } @@ -85,6 +91,9 @@ Result> DataTableStreamScan::NextPlan() { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, follow_up_scanner_->Scan(snapshot.value(), snapshot_reader_)); next_snapshot_id_.value()++; + if (plan->Splits().empty()) { + continue; + } return plan; } else { next_snapshot_id_.value()++; @@ -122,8 +131,19 @@ Result> DataTableStreamScan::GetNextSnapshot( Status DataTableStreamScan::InitScanner() { PAIMON_ASSIGN_OR_RAISE(starting_scanner_, CreateStartingScanner(/*is_streaming=*/true)); - follow_up_scanner_ = std::make_shared(); - return Status::OK(); + switch (core_options_.GetChangelogProducer()) { + case ChangelogProducer::NONE: + follow_up_scanner_ = std::make_shared(); + return Status::OK(); + case ChangelogProducer::INPUT: + case ChangelogProducer::LOOKUP: + follow_up_scanner_ = std::make_shared(); + return Status::OK(); + case ChangelogProducer::FULL_COMPACTION: + return Status::NotImplemented("do not support full compaction changelog producer"); + default: + return Status::NotImplemented("unknown changelog producer"); + } } } // namespace paimon diff --git a/src/paimon/core/table/source/fallback_data_split_test.cpp b/src/paimon/core/table/source/fallback_data_split_test.cpp index 80d290560..5ae97e19e 100644 --- a/src/paimon/core/table/source/fallback_data_split_test.cpp +++ b/src/paimon/core/table/source/fallback_data_split_test.cpp @@ -115,7 +115,8 @@ TEST(FallbackDataSplitTest, TestDeserialize) { /*creation_time=*/Timestamp(1755880762233ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-43880011-d066-4255-ad65-891d79cde23b-0.parquet", /*file_size=*/891, /*row_count=*/1, /*min_key=*/BinaryRow::EmptyRow(), @@ -129,7 +130,8 @@ TEST(FallbackDataSplitTest, TestDeserialize) { /*creation_time=*/Timestamp(1755884315482ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( /*partition=*/BinaryRowGenerator::GenerateRow({1}, pool.get()), @@ -185,7 +187,8 @@ TEST(FallbackDataSplitTest, TestDeserialize2) { /*creation_time=*/Timestamp(1755880762585ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-625b3277-84d3-4320-80b9-89a5075bf5fd-0.parquet", /*file_size=*/891, /*row_count=*/1, /*min_key=*/BinaryRow::EmptyRow(), @@ -199,7 +202,8 @@ TEST(FallbackDataSplitTest, TestDeserialize2) { /*creation_time=*/Timestamp(1755884315889ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); std::vector file_list; file_list.emplace_back( 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 208807493..afb852a6b 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -19,13 +19,31 @@ #include "paimon/core/table/source/key_value_table_read.h" +#include #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/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/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/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.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/core/utils/nested_projection_utils.h" +#include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/status.h" namespace paimon { @@ -34,16 +52,79 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +struct ColumnarBatchContext; -KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, - const std::shared_ptr& path_factory, - const std::shared_ptr& context, - const std::shared_ptr& memory_pool, - const std::shared_ptr& executor) +namespace { + +Result> CreateRealtimePrimaryKeyQueryTransportSchema( + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema) { + arrow::FieldVector transport_value_fields; + transport_value_fields.reserve(key_schema->num_fields() + value_schema->num_fields()); + std::unordered_set field_ids; + for (const std::shared_ptr& field : key_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field_ids.insert(field_id).second) { + transport_value_fields.push_back(field); + } + } + for (const std::shared_ptr& field : value_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field_ids.insert(field_id).second) { + transport_value_fields.push_back(field); + } + } + return RealtimePrimaryKeyLayout::CreateSchema(transport_value_fields); +} + +Result>> CreateMemoryReaders( + const std::shared_ptr& split, const RealtimePartitionBucketView& memory, + const std::shared_ptr& transport_schema, + 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) { + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*transport_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)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> realtime_primary_key_readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, + OffsetRange(split->CommittedEndOffset(), split->MemoryEndOffset()), key_schema, + value_schema, memory_pool)); + std::vector> result; + result.reserve(realtime_primary_key_readers.size()); + for (std::unique_ptr& realtime_primary_key_reader : + realtime_primary_key_readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge, + PrimaryKeyTableUtils::CreateMergeFunction( + value_schema, context->GetTableSchema()->PrimaryKeys(), + context->GetCoreOptions(), memory_pool)); + result.push_back(std::make_unique( + std::move(realtime_primary_key_reader), key_comparator, + std::make_shared(std::move(merge)))); + } + return result; +} + +} // namespace + +KeyValueTableRead::KeyValueTableRead( + std::vector>&& split_reads, + const std::shared_ptr& path_factory, + const std::shared_ptr& context, + const std::shared_ptr& realtime_primary_key_transport_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& executor) : TableRead(memory_pool), split_reads_(std::move(split_reads)), path_factory_(path_factory), context_(context), + realtime_primary_key_transport_schema_(realtime_primary_key_transport_schema), executor_(executor) {} Result> KeyValueTableRead::Create( @@ -57,10 +138,18 @@ Result> KeyValueTableRead::Create( PAIMON_ASSIGN_OR_RAISE( std::unique_ptr merge_file_split_read, MergeFileSplitRead::Create(path_factory, context, memory_pool, executor)); + std::shared_ptr realtime_primary_key_transport_schema; + if (context->GetRealtimeContext()) { + PAIMON_ASSIGN_OR_RAISE( + realtime_primary_key_transport_schema, + CreateRealtimePrimaryKeyQueryTransportSchema(merge_file_split_read->GetKeySchema(), + merge_file_split_read->GetValueSchema())); + } split_reads.emplace_back(std::move(merge_file_split_read)); - return std::unique_ptr(new KeyValueTableRead(std::move(split_reads), path_factory, - context, memory_pool, executor)); + return std::unique_ptr( + new KeyValueTableRead(std::move(split_reads), path_factory, context, + realtime_primary_key_transport_schema, memory_pool, executor)); } void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { @@ -75,6 +164,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 +220,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, realtime_primary_key_transport_schema_, + 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 d6a1c83d3..1dd59b016 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -22,6 +22,7 @@ #include #include +#include "arrow/type_fwd.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/operation/split_read.h" #include "paimon/core/utils/file_store_path_factory.h" @@ -35,6 +36,7 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +class RealtimeSplit; class KeyValueTableRead : public TableRead { public: @@ -45,6 +47,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; @@ -54,12 +59,17 @@ class KeyValueTableRead : public TableRead { KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, const std::shared_ptr& context, + const std::shared_ptr& realtime_primary_key_transport_schema, 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_; + std::shared_ptr realtime_primary_key_transport_schema_; std::shared_ptr executor_; bool force_keep_delete_ = false; }; diff --git a/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp b/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp index a56469656..4126232a0 100644 --- a/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp +++ b/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp @@ -220,7 +220,8 @@ class PrimaryKeySortedIndexScanTest : public ::testing::Test { /*creation_time=*/Timestamp(1721643142456LL, 0), delete_row_count, /*embedded_index=*/nullptr, file_source, /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, - /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); } Result> BuildPayload(std::vector ordinals, diff --git a/src/paimon/core/table/source/realtime_split.h b/src/paimon/core/table/source/realtime_split.h index 2e264e827..7f5714534 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 5a76e81ff..4c3968dc3 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 7d036d420..6a4ad3e91 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, @@ -47,6 +47,10 @@ class RealtimeTableScan : public TableScan { Result> CreatePlan() override; + std::shared_ptr GetMetrics() const override { + return disk_scan_->GetMetrics(); + } + private: using MemoryViewMap = std::map; @@ -67,6 +71,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/scan_mode.h b/src/paimon/core/table/source/scan_mode.h index c236aad69..a01fde506 100644 --- a/src/paimon/core/table/source/scan_mode.h +++ b/src/paimon/core/table/source/scan_mode.h @@ -26,10 +26,10 @@ enum class ScanMode { ALL = 0, /// Only scan newly changed files of a snapshot. - DELTA = 1 + DELTA = 1, /// Only scan changelog files of a snapshot. - /* CHANGELOG = 2 */ + CHANGELOG = 2 }; } // namespace paimon diff --git a/src/paimon/core/table/source/snapshot/changelog_follow_up_scanner.h b/src/paimon/core/table/source/snapshot/changelog_follow_up_scanner.h new file mode 100644 index 000000000..ee2b73dfd --- /dev/null +++ b/src/paimon/core/table/source/snapshot/changelog_follow_up_scanner.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 "paimon/core/table/source/snapshot/follow_up_scanner.h" +#include "paimon/logging.h" + +namespace paimon { + +/// Follow-up scanner for snapshots containing changelog manifests. +class ChangelogFollowUpScanner : public FollowUpScanner { + public: + ChangelogFollowUpScanner() : logger_(Logger::GetLogger("ChangelogFollowUpScanner")) {} + + bool NeedScanSnapshot(const Snapshot& snapshot) const override { + if (snapshot.ChangelogManifestList()) { + return true; + } + PAIMON_LOG_DEBUG(logger_, "Snapshot #%ld has no changelog, check the next snapshot.", + snapshot.Id()); + return false; + } + + Result> Scan( + const Snapshot& snapshot, + const std::shared_ptr& snapshot_reader) const override { + return snapshot_reader->WithMode(ScanMode::CHANGELOG)->WithSnapshot(snapshot)->Read(); + } + + private: + std::unique_ptr logger_; +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/snapshot/snapshot_reader.h b/src/paimon/core/table/source/snapshot/snapshot_reader.h index c425fd863..35134d22a 100644 --- a/src/paimon/core/table/source/snapshot/snapshot_reader.h +++ b/src/paimon/core/table/source/snapshot/snapshot_reader.h @@ -108,6 +108,10 @@ class SnapshotReader { return scan_->GetPartitionPredicate(); } + std::shared_ptr GetMetrics() const { + return scan_->GetScanMetrics(); + } + /// Get splits from `FileKind::ADD` files. Result> Read() const; diff --git a/src/paimon/core/table/source/snapshot/snapshot_reader_test.cpp b/src/paimon/core/table/source/snapshot/snapshot_reader_test.cpp index cebc644cf..427d68fe2 100644 --- a/src/paimon/core/table/source/snapshot/snapshot_reader_test.cpp +++ b/src/paimon/core/table/source/snapshot/snapshot_reader_test.cpp @@ -30,6 +30,8 @@ #include "paimon/core/index/index_file_handler.h" #include "paimon/core/index/index_file_meta.h" #include "paimon/core/io/data_file_meta.h" +#include "paimon/core/snapshot.h" +#include "paimon/core/table/source/snapshot/changelog_follow_up_scanner.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/index_file_path_factories.h" #include "paimon/fs/local/local_file_system.h" @@ -52,7 +54,7 @@ class SnapshotReaderTest : public testing::Test { /*min_sequence_number=*/0, /*max_sequence_number=*/0, /*schema_id=*/0, DataFileMeta::DUMMY_LEVEL, std::vector>{}, Timestamp(0, 0), std::nullopt, nullptr, FileSource::Append(), std::nullopt, std::nullopt, std::nullopt, - std::nullopt); + std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateIndexFileMeta(const std::string& index_file_name, @@ -112,4 +114,23 @@ TEST_F(SnapshotReaderTest, GetDeletionFilesOverwritesDuplicateDataFileName) { EXPECT_EQ(deletion_files[0]->cardinality, std::optional(4)); } +TEST_F(SnapshotReaderTest, ChangelogFollowUpScannerSkipsSnapshotsWithoutChangelog) { + auto create_snapshot = [](const std::optional& changelog_manifest_list) { + return Snapshot( + /*id=*/1, /*schema_id=*/0, /*base_manifest_list=*/"", + /*base_manifest_list_size=*/std::nullopt, /*delta_manifest_list=*/"", + /*delta_manifest_list_size=*/std::nullopt, changelog_manifest_list, + /*changelog_manifest_list_size=*/std::nullopt, + /*index_manifest=*/std::nullopt, /*commit_user=*/"user", /*commit_identifier=*/1, + Snapshot::CommitKind::Append(), /*time_millis=*/0, /*total_record_count=*/0, + /*delta_record_count=*/0, /*changelog_record_count=*/std::nullopt, + /*watermark=*/std::nullopt, /*statistics=*/std::nullopt, + /*properties=*/std::nullopt, /*next_row_id=*/std::nullopt); + }; + + ChangelogFollowUpScanner scanner; + ASSERT_FALSE(scanner.NeedScanSnapshot(create_snapshot(std::nullopt))); + ASSERT_TRUE(scanner.NeedScanSnapshot(create_snapshot("changelog-manifest-list"))); +} + } // namespace paimon::test diff --git a/src/paimon/core/table/source/split.cpp b/src/paimon/core/table/source/split.cpp index 007df3cb1..d6117a899 100644 --- a/src/paimon/core/table/source/split.cpp +++ b/src/paimon/core/table/source/split.cpp @@ -23,6 +23,7 @@ #include "paimon/common/data/binary_row.h" #include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/common/utils/math.h" #include "paimon/common/utils/serialization_utils.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta_serializer.h" @@ -159,8 +160,8 @@ Result Split::Serialize(const std::shared_ptr& split, if (!scores.empty()) { out.WriteValue(true); out.WriteValue(scores.size()); - for (const auto& score : scores) { - out.WriteValue(score); + for (float score : scores) { + out.WriteValue(CanonicalizeFloatingPoint(score)); } } else { out.WriteValue(false); diff --git a/src/paimon/core/table/source/split_generator_test.cpp b/src/paimon/core/table/source/split_generator_test.cpp index ac80eb3f1..9f1368700 100644 --- a/src/paimon/core/table/source/split_generator_test.cpp +++ b/src/paimon/core/table/source/split_generator_test.cpp @@ -74,7 +74,7 @@ class SplitGeneratorTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::optional(), /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateDataFileMeta(const std::string& file_name, int32_t level, @@ -99,7 +99,7 @@ class SplitGeneratorTest : public testing::Test { /*embedded_index=*/nullptr, FileSource::Append(), /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateDataFileMeta(const std::string& file_name, int32_t min_key, @@ -117,7 +117,7 @@ class SplitGeneratorTest : public testing::Test { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*external_path=*/std::nullopt, /*value_stats_cols=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } std::shared_ptr CreateDataFileMetaWithRowId(const std::string& file_name, @@ -132,7 +132,7 @@ class SplitGeneratorTest : public testing::Test { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, first_row_id, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); } static void CheckResult(const std::vector& result_groups, diff --git a/src/paimon/core/table/source/table_read.cpp b/src/paimon/core/table/source/table_read.cpp index e7e3f02f9..3d6a36766 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 a543a051c..95af2a23b 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -27,6 +27,7 @@ #include #include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/predicate/predicate_validator.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/fields_comparator.h" @@ -63,6 +64,7 @@ #include "paimon/core/utils/field_mapping.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/index_file_path_factories.h" +#include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/format/file_format.h" #include "paimon/realtime/realtime_context.h" @@ -187,6 +189,10 @@ Result> NewDataTableScan(const std::shared_ptr TableScan::GetMetrics() const { + return std::make_shared(); +} + Result> TableScan::Create(std::unique_ptr context) { if (context == nullptr) { return Status::Invalid("scan context is null pointer"); @@ -218,12 +224,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 +355,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/source/table_scan_test.cpp b/src/paimon/core/table/source/table_scan_test.cpp index 358dfd469..dbc5b9926 100644 --- a/src/paimon/core/table/source/table_scan_test.cpp +++ b/src/paimon/core/table/source/table_scan_test.cpp @@ -26,11 +26,36 @@ #include "gtest/gtest.h" #include "paimon/defs.h" +#include "paimon/metrics.h" #include "paimon/scan_context.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +class DefaultMetricsTableScan : public TableScan { + public: + Result> CreatePlan() override { + return Status::NotImplemented("not implemented"); + } +}; + +} // namespace + +TEST(TableScanTest, TestDefaultMetricsSnapshot) { + DefaultMetricsTableScan table_scan; + std::shared_ptr metrics = table_scan.GetMetrics(); + ASSERT_TRUE(metrics); + metrics->SetCounter("external", 1); + + std::shared_ptr second_metrics = table_scan.GetMetrics(); + ASSERT_TRUE(second_metrics); + Result external_counter = second_metrics->GetCounter("external"); + ASSERT_FALSE(external_counter.ok()); + ASSERT_EQ(external_counter.status().code(), StatusCode::KeyError); +} + TEST(TableScanTest, TestNoSnapshot) { std::string path = paimon::test::GetDataDir() + "/orc/append_table_with_nested_type.db/append_table_with_nested_type/"; @@ -61,6 +86,27 @@ TEST(TableScanTest, TestPkSchemaEvolutionScan) { ASSERT_OK_AND_ASSIGN(auto plan, table_scan->CreatePlan()); ASSERT_TRUE(plan->SnapshotId()); ASSERT_FALSE(plan->Splits().empty()); + + std::shared_ptr metrics = table_scan->GetMetrics(); + ASSERT_TRUE(metrics); + ASSERT_OK_AND_ASSIGN(uint64_t scanned_snapshot_id, + metrics->GetCounter(ScanMetrics::LAST_SCANNED_SNAPSHOT_ID)); + ASSERT_EQ(scanned_snapshot_id, static_cast(plan->SnapshotId().value())); + ASSERT_OK_AND_ASSIGN(uint64_t resulted_table_files, + metrics->GetCounter(ScanMetrics::LAST_SCAN_RESULTED_TABLE_FILES)); + ASSERT_GT(resulted_table_files, 0); + ASSERT_OK(metrics->GetCounter(ScanMetrics::LAST_MANIFEST_READ_DURATION)); + ASSERT_OK(metrics->GetHistogramStats(ScanMetrics::MANIFEST_READ_DURATION)); + ASSERT_OK_AND_ASSIGN(uint64_t lazy_decode_scanned_rows, + metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_SCANNED_ROWS)); + ASSERT_OK_AND_ASSIGN(uint64_t lazy_decode_materialized_rows, + metrics->GetCounter(ScanMetrics::LAST_LAZY_DECODE_MATERIALIZED_ROWS)); + ASSERT_GE(lazy_decode_scanned_rows, lazy_decode_materialized_rows); + + metrics->SetCounter(ScanMetrics::LAST_SCANNED_SNAPSHOT_ID, 0); + ASSERT_OK_AND_ASSIGN(uint64_t internal_snapshot_id, table_scan->GetMetrics()->GetCounter( + ScanMetrics::LAST_SCANNED_SNAPSHOT_ID)); + ASSERT_EQ(internal_snapshot_id, static_cast(plan->SnapshotId().value())); } TEST(TableScanTest, TestReadOptimizedPrimaryKeyStreamingScanUnsupported) { 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 e384146ed..7572f906e 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/global_system_tables.cpp b/src/paimon/core/table/system/global_system_tables.cpp index dd1b4e606..6523588e8 100644 --- a/src/paimon/core/table/system/global_system_tables.cpp +++ b/src/paimon/core/table/system/global_system_tables.cpp @@ -108,11 +108,12 @@ VariantType OptionalStringValue(const std::map& option Result OptionalLongValue(const std::map& options, const std::string& key) { - if (options.find(key) == options.end()) { + PAIMON_ASSIGN_OR_RAISE(std::optional value, + OptionsUtils::GetOptionalValueFromMap(options, key)); + if (!value) { return VariantType(NullType()); } - PAIMON_ASSIGN_OR_RAISE(int64_t value, OptionsUtils::GetValueFromMap(options, key)); - return VariantType(value); + return VariantType(value.value()); } Result IsEnabled(const GlobalSystemTableRegistryEntry& entry, 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 516861cb5..d7bfa0912 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/table/system/system_table_test.cpp b/src/paimon/core/table/system/system_table_test.cpp index 0c67ae3ac..7e5847dbf 100644 --- a/src/paimon/core/table/system/system_table_test.cpp +++ b/src/paimon/core/table/system/system_table_test.cpp @@ -33,10 +33,12 @@ #include "paimon/core/table/system/audit_log_system_table.h" #include "paimon/core/table/system/binlog_system_table.h" #include "paimon/core/table/system/read_optimized_system_table.h" +#include "paimon/core/table/system/system_table_scan.h" #include "paimon/defs.h" #include "paimon/fs/file_system.h" #include "paimon/fs/file_system_factory.h" #include "paimon/memory/memory_pool.h" +#include "paimon/metrics.h" #include "paimon/reader/batch_reader.h" #include "paimon/result.h" #include "paimon/status.h" @@ -199,4 +201,15 @@ TEST(SystemTableTest, TestGlobalSystemTableWithoutCatalogReturnsNotImplemented) "global system table requires catalog context: tables"); } +TEST(SystemTableTest, TestScanMetricsAreSnapshots) { + SystemTableScan scan("/tmp/table"); + std::shared_ptr metrics = scan.GetMetrics(); + ASSERT_TRUE(metrics); + metrics->SetCounter("external", 1); + + std::shared_ptr second_metrics = scan.GetMetrics(); + ASSERT_TRUE(second_metrics); + ASSERT_NOK_WITH_MSG(second_metrics->GetCounter("external"), "metric 'external' not found"); +} + } // namespace paimon::test diff --git a/src/paimon/core/utils/field_mapping.cpp b/src/paimon/core/utils/field_mapping.cpp index 447d8d589..be7287dd6 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 f8509fe23..43d782019 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,32 @@ 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(); + ColumnarRow row(fields, pool_, /*row_id=*/0); + for (int64_t i = 0; i < struct_array->length(); i++) { + row.SetRowId(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 +200,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 000000000..b576d674f --- /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 000000000..2ac62b92e --- /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.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index cf72da4ae..2ca444a7a 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -29,12 +29,14 @@ #include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/object_utils.h" #include "paimon/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" #include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/first_row_merge_function.h" #include "paimon/core/mergetree/compact/merge_function.h" #include "paimon/core/mergetree/compact/partial_update_merge_function.h" #include "paimon/core/options/merge_engine.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/status.h" namespace arrow { @@ -96,4 +98,52 @@ Result> PrimaryKeyTableUtils::CreateSequenceFi options.SequenceFieldSortOrderIsAscending()); } +Status PrimaryKeyTableUtils::ValidateRealtimeOptions(const CoreOptions& options, + const TableSchema& schema) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime does not support data evolution"); + } + if (options.IgnoreDelete()) { + return Status::NotImplemented("PK realtime requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime supports only ascending sequence.field.sort-order"); + } + if (options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime supports only the NONE changelog producer"); + } + if (options.DeletionVectorsEnabled()) { + return Status::NotImplemented("PK realtime does not support deletion vectors"); + } + if (options.NeedLookup()) { + return Status::NotImplemented("PK realtime does not support lookup"); + } + 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 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 does not support global indexes"); + } + } + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index 82a108ab7..114801cc1 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -36,6 +36,7 @@ class CoreOptions; class MemoryPool; class FieldsComparator; class DataField; +class TableSchema; class PrimaryKeyTableUtils { public: @@ -57,6 +58,8 @@ class PrimaryKeyTableUtils { static Result> CreateSequenceFieldsComparator( const std::vector& value_fields, const CoreOptions& options); + + static Status ValidateRealtimeOptions(const CoreOptions& options, const TableSchema& schema); }; } // namespace paimon 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 12713ca5b..796922336 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -22,7 +22,9 @@ #include #include #include +#include #include +#include #include "arrow/type.h" #include "gtest/gtest.h" @@ -32,6 +34,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/key_value.h" #include "paimon/core/mergetree/compact/merge_function.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/defs.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" @@ -40,6 +43,110 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +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(); +} + +} // namespace + +TEST(PrimaryKeyTableUtilsTest, TestSupportedRealtimeOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema())); +} + +TEST(PrimaryKeyTableUtilsTest, TestUnsupportedRealtimeOptions) { + 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"}, {Options::SEQUENCE_FIELD, "seq"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema())); + } +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeAcceptsInactiveMergeEngineOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::map option_map = { + {Options::BUCKET, "1"}, + {sequence_group, "seq"}, + {Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_DELETE, "true"}, + {Options::AGGREGATION_REMOVE_RECORD_ON_DELETE, "true"}, + {Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_SEQUENCE_GROUP, "seq"}, + }; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr schema, + TableSchema::Create(0, + arrow::schema({arrow::field("id", arrow::int64()), + arrow::field("value", arrow::utf8()), + arrow::field("seq", arrow::int64())}), + {}, {"id"}, option_map)); + ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *schema)); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsDeleteAndSequenceOrderingOptions) { + ASSERT_OK_AND_ASSIGN( + CoreOptions ignore_delete, + CoreOptions::FromMap({{Options::BUCKET, "1"}, {Options::IGNORE_DELETE, "true"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyTableUtils::ValidateRealtimeOptions(ignore_delete, *PkSchema()), + "requires default delete behavior"); + + ASSERT_OK_AND_ASSIGN( + CoreOptions descending, + CoreOptions::FromMap( + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD_SORT_ORDER, "descending"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyTableUtils::ValidateRealtimeOptions(descending, *PkSchema()), + "supports only ascending sequence.field.sort-order"); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeReportsSpecificLookupErrors) { + const std::vector, std::string>> cases = { + {{{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + "PK realtime does not support lookup"}, + {{{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + "PK realtime does not support deletion vectors"}, + {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + "PK realtime supports only the NONE changelog producer"}, + {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "lookup"}}, + "PK realtime supports only the NONE changelog producer"}, + }; + for (const auto& [option_map, expected_message] : cases) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + Status status = PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema()); + ASSERT_TRUE(status.IsNotImplemented()); + ASSERT_EQ(status.message(), expected_message); + } +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsFloatingPrimaryKeys) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_NOK_WITH_MSG( + PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema(arrow::float32())), + "FLOAT or DOUBLE primary keys"); + ASSERT_NOK_WITH_MSG( + PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema(arrow::float64())), + "FLOAT or DOUBLE primary keys"); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsEnabledGlobalIndex) { + 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(PrimaryKeyTableUtils::ValidateRealtimeOptions( + options, *PkSchema(arrow::int64(), option_map)), + "does not support global indexes"); +} TEST(PrimaryKeyTableUtilsTest, TestCreateSequenceFieldsComparator) { { diff --git a/src/paimon/format/avro/avro_direct_decoder.cpp b/src/paimon/format/avro/avro_direct_decoder.cpp index f837eed0d..9f9e36426 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. @@ -41,6 +57,20 @@ Status DecodeFieldToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder* decoder, arrow::ArrayBuilder* array_builder, AvroDirectDecoder::DecodeContext* ctx); +Status ReserveBuilderCapacityImpl(int64_t capacity, arrow::ArrayBuilder* array_builder) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(array_builder->Reserve(capacity)); + if (array_builder->type()->id() != arrow::Type::STRUCT) { + return Status::OK(); + } + + auto* struct_builder = checked_cast(array_builder); + for (int32_t i = 0; i < struct_builder->num_fields(); ++i) { + PAIMON_RETURN_NOT_OK( + ReserveBuilderCapacityImpl(capacity, struct_builder->field_builder(i))); + } + return Status::OK(); +} + /// \brief Skip an Avro value based on its schema without decoding Status SkipAvroValue(const ::avro::NodePtr& avro_node, ::avro::Decoder* decoder) { switch (avro_node->type()) { @@ -177,6 +207,7 @@ Status DecodeListToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder* de // Read array block count int64_t block_count = decoder->arrayStart(); while (block_count != 0) { + PAIMON_RETURN_NOT_OK(ReserveBuilderCapacityImpl(block_count, value_builder)); for (int64_t i = 0; i < block_count; ++i) { PAIMON_RETURN_NOT_OK(DecodeFieldToBuilder(element_node, /*projection=*/std::nullopt, decoder, value_builder, ctx)); @@ -205,6 +236,8 @@ Status DecodeMapToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder* dec // Read map block count int64_t block_count = decoder->mapStart(); while (block_count != 0) { + PAIMON_RETURN_NOT_OK(ReserveBuilderCapacityImpl(block_count, key_builder)); + PAIMON_RETURN_NOT_OK(ReserveBuilderCapacityImpl(block_count, item_builder)); for (int64_t i = 0; i < block_count; ++i) { PAIMON_RETURN_NOT_OK(DecodeFieldToBuilder(key_node, /*projection=*/std::nullopt, decoder, key_builder, ctx)); @@ -232,6 +265,8 @@ Status DecodeMapToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder* dec // Read array block count int64_t block_count = decoder->arrayStart(); while (block_count != 0) { + PAIMON_RETURN_NOT_OK(ReserveBuilderCapacityImpl(block_count, key_builder)); + PAIMON_RETURN_NOT_OK(ReserveBuilderCapacityImpl(block_count, item_builder)); for (int64_t i = 0; i < block_count; ++i) { PAIMON_RETURN_NOT_OK(DecodeFieldToBuilder(key_node, /*projection=*/std::nullopt, decoder, key_builder, ctx)); @@ -266,8 +301,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 +322,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 +331,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 +350,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)); @@ -431,4 +466,9 @@ Status AvroDirectDecoder::DecodeAvroToBuilder(const ::avro::NodePtr& avro_node, return DecodeFieldToBuilder(avro_node, projection, decoder, array_builder, ctx); } +Status AvroDirectDecoder::ReserveBuilderCapacity(int64_t capacity, + arrow::ArrayBuilder* array_builder) { + return ReserveBuilderCapacityImpl(capacity, array_builder); +} + } // namespace paimon::avro diff --git a/src/paimon/format/avro/avro_direct_decoder.h b/src/paimon/format/avro/avro_direct_decoder.h index c507091a7..3b316f540 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 @@ -61,6 +81,12 @@ class AvroDirectDecoder { const std::optional>& projection, ::avro::Decoder* decoder, arrow::ArrayBuilder* array_builder, DecodeContext* ctx); + + /// Reserve slots for a builder and any struct children with the same cardinality. + /// @param capacity Number of additional values to append. + /// @param array_builder Builder to reserve. + /// @return Status::OK if all reservations succeed. + static Status ReserveBuilderCapacity(int64_t capacity, arrow::ArrayBuilder* array_builder); }; } // namespace paimon::avro 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 f276d9469..8e65520d3 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 { @@ -338,6 +361,33 @@ TEST_F(AvroDirectEncoderDecoderTest, TestRecordType) { CheckResult(schema_json, input_array, &struct_builder); } +TEST_F(AvroDirectEncoderDecoderTest, TestReserveBuilderCapacity) { + std::shared_ptr nested_type = arrow::struct_({ + arrow::field("id", arrow::int32()), + arrow::field("nested", arrow::struct_({arrow::field("value", arrow::int64())})), + arrow::field("items", arrow::list(arrow::int32())), + }); + arrow::Result> builder_result = + arrow::MakeBuilder(nested_type); + ASSERT_TRUE(builder_result.ok()) << builder_result.status().ToString(); + std::unique_ptr builder = std::move(builder_result).ValueOrDie(); + + constexpr int64_t capacity = 1024; + ASSERT_OK(AvroDirectDecoder::ReserveBuilderCapacity(capacity, builder.get())); + + auto* root_builder = checked_cast(builder.get()); + ASSERT_GE(root_builder->capacity(), capacity); + ASSERT_GE(root_builder->field_builder(0)->capacity(), capacity); + + auto* nested_builder = checked_cast(root_builder->field_builder(1)); + ASSERT_GE(nested_builder->capacity(), capacity); + ASSERT_GE(nested_builder->field_builder(0)->capacity(), capacity); + + auto* list_builder = checked_cast(root_builder->field_builder(2)); + ASSERT_GE(list_builder->capacity(), capacity); + ASSERT_EQ(list_builder->value_builder()->capacity(), 0); +} + TEST_F(AvroDirectEncoderDecoderTest, TestDecodeWithProjection) { arrow::FieldVector fields = { arrow::field("f0", arrow::boolean()), diff --git a/src/paimon/format/avro/avro_file_batch_reader.cpp b/src/paimon/format/avro/avro_file_batch_reader.cpp index f48ec4cc4..f8b00c7c6 100644 --- a/src/paimon/format/avro/avro_file_batch_reader.cpp +++ b/src/paimon/format/avro/avro_file_batch_reader.cpp @@ -110,6 +110,10 @@ Result AvroFileBatchReader::NextBatch() { if (!reader_->hasMore()) { break; } + if (array_builder_->length() == 0) { + PAIMON_RETURN_NOT_OK( + AvroDirectDecoder::ReserveBuilderCapacity(batch_size_, array_builder_.get())); + } reader_->decr(); PAIMON_RETURN_NOT_OK(AvroDirectDecoder::DecodeAvroToBuilder( reader_->dataSchema().root(), read_fields_projection_, &reader_->decoder(), @@ -123,7 +127,10 @@ Result AvroFileBatchReader::NextBatch() { } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, array_builder_->Finish()); +#ifndef NDEBUG + // Keep structural validation in debug builds without adding its recursive cost to reads. PAIMON_RETURN_NOT_OK_FROM_ARROW(array->Validate()); +#endif std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get())); @@ -172,6 +179,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/mosaic/CMakeLists.txt b/src/paimon/format/mosaic/CMakeLists.txt new file mode 100644 index 000000000..ca1612f97 --- /dev/null +++ b/src/paimon/format/mosaic/CMakeLists.txt @@ -0,0 +1,60 @@ +# 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. + +if(PAIMON_ENABLE_MOSAIC) + set(PAIMON_MOSAIC_FILE_FORMAT + mosaic_file_batch_reader.cpp + mosaic_file_format.cpp + mosaic_file_format_factory.cpp + mosaic_format_writer.cpp + mosaic_stats.cpp + mosaic_stats_extractor.cpp + mosaic_stream.cpp + mosaic_writer_builder.cpp) + + add_paimon_lib(paimon_mosaic_file_format + SOURCES + ${PAIMON_MOSAIC_FILE_FORMAT} + DEPENDENCIES + paimon_shared + paimon_mosaic_ffi + STATIC_LINK_LIBS + arrow + fmt + paimon_mosaic_ffi + Threads::Threads + SHARED_LINK_LIBS + paimon_shared + paimon_mosaic_ffi + SHARED_LINK_FLAGS + ${PAIMON_VERSION_SCRIPT_FLAGS}) + + target_link_libraries(paimon_mosaic_file_format_objlib PUBLIC paimon_mosaic_ffi) + + if(PAIMON_BUILD_TESTS) + add_paimon_test(mosaic_format_test + SOURCES + mosaic_file_format_test.cpp + STATIC_LINK_LIBS + paimon_shared + test_utils_static + ${PAIMON_MOSAIC_FILE_FORMAT_STATIC_LINK_LIBS} + ${PAIMON_LOCAL_FILE_SYSTEM_STATIC_LINK_LIBS} + paimon_mosaic_ffi + ${GTEST_LINK_TOOLCHAIN}) + endif() +endif() diff --git a/src/paimon/format/mosaic/mosaic_ffi.h b/src/paimon/format/mosaic/mosaic_ffi.h new file mode 100644 index 000000000..5b7193eec --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_ffi.h @@ -0,0 +1,24 @@ +/* + * 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 + +extern "C" { +#include "mosaic.h" // NOLINT(build/include_subdir) +} diff --git a/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp new file mode 100644 index 000000000..b065c3632 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_batch_reader.cpp @@ -0,0 +1,286 @@ +/* + * 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/format/mosaic/mosaic_file_batch_reader.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/predicate/predicate_filter.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/math.h" +#include "paimon/core/stats/simple_stats.h" +#include "paimon/core/stats/simple_stats_converter.h" +#include "paimon/format/mosaic/mosaic_stats.h" +#include "paimon/fs/file_system.h" + +namespace paimon::mosaic { + +MosaicFileBatchReader::MosaicFileBatchReader( + const std::shared_ptr& input, int32_t batch_size, + std::unique_ptr input_context, MosaicReaderHandle* reader, + const std::shared_ptr& file_schema, uint32_t num_row_groups, uint64_t total_rows, + const std::shared_ptr& pool, const std::shared_ptr& arrow_pool) + : input_(input), + batch_size_(batch_size), + input_context_(std::move(input_context)), + reader_(reader), + file_schema_(file_schema), + num_row_groups_(num_row_groups), + total_rows_(total_rows), + pool_(pool), + arrow_pool_(arrow_pool), + metrics_(std::make_shared()) {} + +Result> MosaicFileBatchReader::Create( + const std::shared_ptr& input, int32_t batch_size, + const std::shared_ptr& pool) { + if (input == nullptr || pool == nullptr || batch_size <= 0) { + return Status::Invalid( + "Mosaic reader requires non-null input and memory pool, and positive batch size"); + } + std::shared_ptr arrow_pool = GetArrowPool(pool); + PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Mosaic input length")); + auto length = static_cast(signed_length); + auto input_context = std::make_unique(input, length); + MosaicInputFile input_file = {}; + input_file.ctx = input_context.get(); + input_file.read_at_fn = MosaicInputContext::ReadAt; + input_file.length_fn = MosaicInputContext::Length; + std::unique_ptr reader( + mosaic_reader_open(input_file), mosaic_reader_free); + if (reader == nullptr) { + return MosaicFfiError("open Mosaic reader", input_context->GetCallbackStatus()); + } + + ::ArrowSchema ffi_schema = {}; + if (mosaic_reader_export_schema(reader.get(), &ffi_schema) != 0) { + return MosaicFfiError("read Mosaic schema", input_context->GetCallbackStatus()); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_schema, + arrow::ImportSchema(&ffi_schema)); + + uint32_t num_row_groups = 0; + if (mosaic_reader_num_row_groups(reader.get(), &num_row_groups) != 0) { + return MosaicFfiError("read Mosaic row group count", input_context->GetCallbackStatus()); + } + uint64_t total_rows = 0; + for (uint32_t row_group = 0; row_group < num_row_groups; ++row_group) { + uint32_t row_count = 0; + if (mosaic_reader_row_group_num_rows(reader.get(), row_group, &row_count) != 0) { + return MosaicFfiError("read Mosaic row count", input_context->GetCallbackStatus()); + } + total_rows += row_count; + } + return std::unique_ptr( + new MosaicFileBatchReader(input, batch_size, std::move(input_context), reader.release(), + file_schema, num_row_groups, total_rows, pool, arrow_pool)); +} + +MosaicFileBatchReader::~MosaicFileBatchReader() { + CloseInternal(); +} + +Result> MosaicFileBatchReader::ReadNextRowGroup() { + while (next_row_group_ < num_row_groups_) { + uint32_t row_group = next_row_group_++; + uint32_t row_count = 0; + if (mosaic_reader_row_group_num_rows(reader_, row_group, &row_count) != 0) { + return MosaicFfiError("read Mosaic row count", input_context_->GetCallbackStatus()); + } + current_row_group_first_row_ = next_row_group_first_row_; + next_row_group_first_row_ += row_count; + PAIMON_ASSIGN_OR_RAISE(bool matches, MatchesRowGroup(row_group, row_count)); + if (!matches) { + continue; + } + + MosaicRowGroupReaderHandle* row_group_reader = + mosaic_reader_open_row_group(reader_, row_group); + if (row_group_reader == nullptr) { + return MosaicFfiError("open Mosaic row group", input_context_->GetCallbackStatus()); + } + MosaicRecordBatchHandle* record_batch = + mosaic_row_group_reader_read_columns(row_group_reader); + mosaic_row_group_reader_free(row_group_reader); + if (record_batch == nullptr) { + return MosaicFfiError("read Mosaic row group", input_context_->GetCallbackStatus()); + } + ::ArrowArray ffi_array = {}; + ::ArrowSchema ffi_schema = {}; + int32_t export_result = mosaic_record_batch_export(record_batch, &ffi_array, &ffi_schema); + mosaic_record_batch_free(record_batch); + if (export_result != 0) { + return MosaicFfiError("export Mosaic row group", input_context_->GetCallbackStatus()); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr batch, + arrow::ImportArray(&ffi_array, &ffi_schema)); + if (batch->length() != row_count) { + return Status::Invalid("Mosaic row group row count mismatch"); + } + if (batch->length() != 0) { + return batch; + } + } + return std::shared_ptr(); +} + +Result MosaicFileBatchReader::MatchesRowGroup(uint32_t row_group, uint32_t row_count) { + if (predicate_filter_ == nullptr) { + return true; + } + PAIMON_ASSIGN_OR_RAISE( + MosaicStatsUtils::RowGroupStatistics stats, + MosaicStatsUtils::ReadRowGroupStatistics(row_group, input_context_.get(), reader_)); + // This matches the Java Mosaic reader: a file without row-group statistics is always kept. + if (stats.empty()) { + return true; + } + std::vector row_group_stats; + row_group_stats.push_back(std::move(stats)); + PAIMON_ASSIGN_OR_RAISE( + ColumnStatsVector column_stats, + MosaicStatsUtils::ConvertColumnStatistics(file_schema_, row_group_stats, + /*missing_null_count_is_zero=*/false)); + PAIMON_ASSIGN_OR_RAISE(SimpleStats simple_stats, + SimpleStatsConverter::ToBinary(column_stats, pool_.get())); + return predicate_filter_->Test(file_schema_, row_count, simple_stats.MinValues(), + simple_stats.MaxValues(), simple_stats.NullCounts()); +} + +Result MosaicFileBatchReader::NextBatch() { + if (closed_) { + return Status::Invalid("Mosaic reader is closed"); + } + if (current_batch_ == nullptr || current_batch_offset_ == current_batch_->length()) { + PAIMON_ASSIGN_OR_RAISE(current_batch_, ReadNextRowGroup()); + current_batch_offset_ = 0; + } + if (current_batch_ == nullptr) { + previous_batch_row_count_ = 0; + return BatchReader::MakeEofBatch(); + } + + int64_t row_count = + std::min(batch_size_, current_batch_->length() - current_batch_offset_); + std::shared_ptr sliced_array = + current_batch_->Slice(current_batch_offset_, row_count); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_array, + ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); + + previous_first_row_ = current_row_group_first_row_ + current_batch_offset_; + previous_batch_row_count_ = row_count; + current_batch_offset_ += row_count; + auto ffi_array = std::make_unique<::ArrowArray>(); + auto ffi_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*normalized_array, ffi_array.get(), ffi_schema.get())); + return std::make_pair(std::move(ffi_array), std::move(ffi_schema)); +} + +Result> MosaicFileBatchReader::GetFileSchema() const { + auto schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*file_schema_, schema.get())); + return schema; +} + +Status MosaicFileBatchReader::SetReadSchema( + ::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + if (read_schema == nullptr) { + return Status::Invalid("Mosaic read schema is nullptr"); + } + (void)selection_bitmap; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, + arrow::ImportSchema(read_schema)); + std::vector names; + std::vector name_pointers; + names.reserve(schema->num_fields()); + name_pointers.reserve(schema->num_fields()); + for (const std::shared_ptr& field : schema->fields()) { + names.push_back(field->name()); + } + for (const std::string& name : names) { + name_pointers.push_back(name.c_str()); + } + if (mosaic_reader_set_projection(reader_, name_pointers.data(), name_pointers.size()) != 0) { + return MosaicFfiError("set Mosaic projection", input_context_->GetCallbackStatus()); + } + predicate_filter_ = std::dynamic_pointer_cast(predicate); + next_row_group_ = 0; + next_row_group_first_row_ = 0; + current_row_group_first_row_ = 0; + current_batch_.reset(); + current_batch_offset_ = 0; + previous_first_row_ = std::numeric_limits::max(); + previous_batch_row_count_ = 0; + return Status::OK(); +} + +Result MosaicFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch_row_id) const { + if (previous_batch_row_count_ == 0) { + return Status::Invalid(previous_first_row_ == std::numeric_limits::max() + ? "no Mosaic batch has been read yet" + : "last Mosaic batch was EOF"); + } + if (batch_row_id >= previous_batch_row_count_) { + return Status::Invalid(fmt::format("batch row id {} is out of range {}", batch_row_id, + previous_batch_row_count_)); + } + return previous_first_row_ + batch_row_id; +} + +Result MosaicFileBatchReader::GetNumberOfRows() const { + return total_rows_; +} + +std::shared_ptr MosaicFileBatchReader::GetReaderMetrics() const { + return metrics_; +} + +void MosaicFileBatchReader::Close() { + CloseInternal(); +} + +void MosaicFileBatchReader::CloseInternal() { + if (!closed_) { + if (reader_ != nullptr) { + mosaic_reader_free(reader_); + reader_ = nullptr; + } + if (input_ != nullptr) { + (void)input_->Close(); + } + input_context_.reset(); + input_.reset(); + current_batch_.reset(); + closed_ = true; + } +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_batch_reader.h b/src/paimon/format/mosaic/mosaic_file_batch_reader.h new file mode 100644 index 000000000..030e26c44 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_batch_reader.h @@ -0,0 +1,95 @@ +/* + * 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/type_fwd.h" +#include "paimon/format/mosaic/mosaic_ffi.h" +#include "paimon/format/mosaic/mosaic_stream.h" +#include "paimon/reader/file_batch_reader.h" +#include "paimon/result.h" + +namespace paimon { +class InputStream; +class MemoryPool; +class Metrics; +class PredicateFilter; +} // namespace paimon + +namespace paimon::mosaic { + +class MosaicFileBatchReader : public FileBatchReader { + public: + static Result> Create( + const std::shared_ptr& input, int32_t batch_size, + const std::shared_ptr& pool); + + ~MosaicFileBatchReader() override; + + Result NextBatch() override; + Result> GetFileSchema() const override; + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + Result GetNumberOfRows() const override; + std::shared_ptr GetReaderMetrics() const override; + void Close() override; + bool SupportPreciseBitmapSelection() const override { + return false; + } + + private: + MosaicFileBatchReader(const std::shared_ptr& input, int32_t batch_size, + std::unique_ptr input_context, + MosaicReaderHandle* reader, + const std::shared_ptr& file_schema, + uint32_t num_row_groups, uint64_t total_rows, + const std::shared_ptr& pool, + const std::shared_ptr& arrow_pool); + + Result> ReadNextRowGroup(); + Result MatchesRowGroup(uint32_t row_group, uint32_t row_count); + void CloseInternal(); + + std::shared_ptr input_; + int32_t batch_size_; + std::unique_ptr input_context_; + MosaicReaderHandle* reader_; + std::shared_ptr file_schema_; + uint32_t num_row_groups_; + uint64_t total_rows_; + uint32_t next_row_group_ = 0; + uint64_t next_row_group_first_row_ = 0; + uint64_t current_row_group_first_row_ = 0; + std::shared_ptr current_batch_; + int64_t current_batch_offset_ = 0; + uint64_t previous_first_row_ = std::numeric_limits::max(); + uint64_t previous_batch_row_count_ = 0; + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr predicate_filter_; + std::shared_ptr metrics_; + bool closed_ = false; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_format.cpp b/src/paimon/format/mosaic/mosaic_file_format.cpp new file mode 100644 index 000000000..937a72624 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_format.cpp @@ -0,0 +1,57 @@ +/* + * 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/format/mosaic/mosaic_file_format.h" + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/format/mosaic/mosaic_reader_builder.h" +#include "paimon/format/mosaic/mosaic_stats_extractor.h" +#include "paimon/format/mosaic/mosaic_writer_builder.h" + +namespace paimon::mosaic { + +Result> MosaicFileFormat::CreateReaderBuilder( + int32_t batch_size) const { + return std::make_unique(batch_size); +} + +Result> MosaicFileFormat::CreateWriterBuilder( + ::ArrowSchema* schema, int32_t batch_size) const { + (void)batch_size; + if (schema == nullptr) { + return Status::Invalid("Mosaic writer schema is nullptr"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr typed_schema, + arrow::ImportSchema(schema)); + return std::make_unique(typed_schema, options_); +} + +Result> MosaicFileFormat::CreateStatsExtractor( + ::ArrowSchema* schema) const { + if (schema == nullptr) { + return Status::Invalid("Mosaic stats schema is nullptr"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr typed_schema, + arrow::ImportSchema(schema)); + return std::make_unique(typed_schema); +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_format.h b/src/paimon/format/mosaic/mosaic_file_format.h new file mode 100644 index 000000000..90fb9fc9d --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_format.h @@ -0,0 +1,49 @@ +/* + * 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/format/file_format.h" + +namespace paimon::mosaic { + +class MosaicFileFormat : public FileFormat { + public: + explicit MosaicFileFormat(const std::map& options) + : identifier_("mosaic"), options_(options) {} + + const std::string& Identifier() const override { + return identifier_; + } + Result> CreateReaderBuilder(int32_t batch_size) const override; + Result> CreateWriterBuilder(::ArrowSchema* schema, + int32_t batch_size) const override; + Result> CreateStatsExtractor( + ::ArrowSchema* schema) const override; + + private: + std::string identifier_; + std::map options_; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_format_factory.cpp b/src/paimon/format/mosaic/mosaic_file_format_factory.cpp new file mode 100644 index 000000000..0b6b6b397 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_format_factory.cpp @@ -0,0 +1,36 @@ +/* + * 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/format/mosaic/mosaic_file_format_factory.h" + +#include "paimon/factories/factory.h" +#include "paimon/format/mosaic/mosaic_file_format.h" + +namespace paimon::mosaic { + +const char MosaicFileFormatFactory::IDENTIFIER[] = "mosaic"; + +Result> MosaicFileFormatFactory::Create( + const std::map& options) const { + return std::make_unique(options); +} + +REGISTER_PAIMON_FACTORY(MosaicFileFormatFactory); + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_format_factory.h b/src/paimon/format/mosaic/mosaic_file_format_factory.h new file mode 100644 index 000000000..e0326765d --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_format_factory.h @@ -0,0 +1,41 @@ +/* + * 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/format/file_format_factory.h" + +namespace paimon::mosaic { + +class MosaicFileFormatFactory : public FileFormatFactory { + public: + static const char IDENTIFIER[]; + + const char* Identifier() const override { + return IDENTIFIER; + } + Result> Create( + const std::map& options) const override; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_file_format_test.cpp b/src/paimon/format/mosaic/mosaic_file_format_test.cpp new file mode 100644 index 000000000..5dc37b5af --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_file_format_test.cpp @@ -0,0 +1,554 @@ +/* + * 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 +#include + +#include "arrow/api.h" +#include "arrow/array/concatenate.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/format/column_stats.h" +#include "paimon/format/file_format.h" +#include "paimon/format/file_format_factory.h" +#include "paimon/format/format_writer.h" +#include "paimon/format/mosaic/mosaic_format_defs.h" +#include "paimon/format/reader_builder.h" +#include "paimon/format/writer_builder.h" +#include "paimon/fs/file_system.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/reader/batch_reader.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/timezone_guard.h" + +namespace paimon::mosaic::test { + +class MosaicFileFormatTest : public ::testing::Test { + public: + // The footer layout stores the number of buckets followed by the number of row groups. + using FooterLayout = std::pair; + + void SetUp() override { + ASSERT_OK_AND_ASSIGN(format_, + FileFormatFactory::Get("mosaic", {{"file.format", "mosaic"}})); + file_system_ = std::make_shared(); + directory_ = paimon::test::UniqueTestDirectory::Create(); + ASSERT_NE(directory_, nullptr); + pool_ = GetDefaultPool(); + arrow_pool_ = GetArrowPool(pool_); + } + + Status WriteFile(const std::string& path, const std::shared_ptr& schema, + const std::shared_ptr& array, int32_t batch_size, + const std::shared_ptr& format = nullptr) const { + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + const std::shared_ptr& writer_format = format == nullptr ? format_ : format; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer_builder, + writer_format->CreateWriterBuilder(&ffi_schema, batch_size)); + writer_builder->WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, + file_system_->Create(path, /*overwrite=*/false)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + writer_builder->Build(output, "zstd")); + for (int64_t offset = 0; offset < array->length(); offset += batch_size) { + std::shared_ptr slice = array->Slice(offset, batch_size); + ::ArrowArray ffi_array = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*slice, &ffi_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&ffi_array)); + } + PAIMON_RETURN_NOT_OK(writer->Finish()); + return output->Close(); + } + + Result ReadFooterLayout(const std::string& path) const { + constexpr int64_t kFooterSize = 32; + constexpr int64_t kLayoutSize = 8; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, file_system_->Open(path)); + PAIMON_ASSIGN_OR_RAISE(int64_t file_size, input->Length()); + if (file_size < kFooterSize) { + return Status::Invalid("Mosaic file is shorter than its footer"); + } + std::array layout = {}; + PAIMON_ASSIGN_OR_RAISE(int64_t bytes_read, + input->Read(reinterpret_cast(layout.data()), kLayoutSize, + file_size - kFooterSize + /*layout offset=*/16)); + if (bytes_read != kLayoutSize) { + return Status::IOError("short read while reading Mosaic footer"); + } + auto decode_uint32 = [&layout](size_t offset) -> uint32_t { + return (uint32_t{layout[offset]} << 24) | (uint32_t{layout[offset + 1]} << 16) | + (uint32_t{layout[offset + 2]} << 8) | uint32_t{layout[offset + 3]}; + }; + return std::make_pair(decode_uint32(0), decode_uint32(4)); + } + + Result> ReadFile(const std::string& path, + const std::shared_ptr& schema, + int32_t batch_size, + uint64_t expected_row_count) const { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(batch_size)); + reader_builder->WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, file_system_->Open(path)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + reader_builder->Build(input)); + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + PAIMON_RETURN_NOT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + PAIMON_ASSIGN_OR_RAISE(uint64_t total_rows, reader->GetNumberOfRows()); + if (total_rows != expected_row_count) { + return Status::Invalid("unexpected Mosaic row count"); + } + + std::vector> batches; + uint64_t expected_first_row = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_RETURN_NOT_OK(paimon::test::ReadResultCollector::CheckBatchOffset(batch)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr arrow_batch, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (arrow_batch->length() > batch_size) { + return Status::Invalid("Mosaic read batch exceeds configured batch size"); + } + PAIMON_ASSIGN_OR_RAISE(uint64_t first_row, + reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); + if (first_row != expected_first_row) { + return Status::Invalid("unexpected Mosaic batch first row"); + } + expected_first_row += arrow_batch->length(); + batches.push_back(std::move(arrow_batch)); + } + if (expected_first_row != expected_row_count) { + return Status::Invalid("Mosaic batches do not contain all rows"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches, arrow_pool_.get())); + return result; + } + + void AssertReadWithBatchSizes(const std::string& path, + const std::shared_ptr& schema, + const std::shared_ptr& expected, + std::initializer_list batch_sizes) const { + for (int32_t batch_size : batch_sizes) { + SCOPED_TRACE("batch_size=" + std::to_string(batch_size)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadFile(path, schema, batch_size, + /*expected_row_count=*/static_cast(expected->length()))); + ASSERT_TRUE(actual->Equals(expected)) << actual->ToString() << "\nvs\n" + << expected->ToString(); + } + } + + protected: + std::shared_ptr format_; + std::shared_ptr file_system_; + std::unique_ptr directory_; + std::shared_ptr pool_; + std::unique_ptr arrow_pool_; +}; + +TEST_F(MosaicFileFormatTest, WriteThenRead) { + std::string path = PathUtil::JoinPath(directory_->Str(), "data.mosaic"); + arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), + arrow::field("name", arrow::utf8())}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data_type = arrow::struct_(fields); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON( + data_type, R"([[1,"one"],[2,null],[3,"three"],[4,"four"],[5,"five"]])") + .ValueOrDie(); + + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2)); + AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5, 8}); +} + +TEST_F(MosaicFileFormatTest, EmptyProjectionPreservesRowCount) { + std::string path = PathUtil::JoinPath(directory_->Str(), "empty-projection.mosaic"); + arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), + arrow::field("name", arrow::utf8())}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[1,"one"],[2,"two"],[3,"three"]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(/*batch_size=*/2)); + reader_builder->WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(path)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, reader_builder->Build(input)); + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema({}), &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + + for (int64_t expected_rows : {2, 1}) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + std::shared_ptr record_batch = + arrow::ImportRecordBatch(batch.first.get(), batch.second.get()).ValueOrDie(); + ASSERT_EQ(record_batch->num_columns(), 0); + ASSERT_EQ(record_batch->num_rows(), expected_rows); + } + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof_batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof_batch)); +} + +TEST_F(MosaicFileFormatTest, SetReadSchemaResetsReaderToFirstRow) { + std::string path = PathUtil::JoinPath(directory_->Str(), "reset-reader.mosaic"); + arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), + arrow::field("name", arrow::utf8())}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[1,"one"],[2,"two"],[3,"three"]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(/*batch_size=*/2)); + reader_builder->WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(path)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, reader_builder->Build(input)); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch first_batch, reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(uint64_t first_row, reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); + ASSERT_EQ(first_row, 0); + std::shared_ptr first_array = + arrow::ImportArray(first_batch.first.get(), first_batch.second.get()).ValueOrDie(); + ASSERT_TRUE(first_array->Equals(data->Slice(/*offset=*/0, /*length=*/2))); + + std::shared_ptr projected_schema = arrow::schema({fields[1]}); + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*projected_schema, &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_NOK(reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch projected_batch, reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(first_row, reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); + ASSERT_EQ(first_row, 0); + std::shared_ptr projected_array = + arrow::ImportArray(projected_batch.first.get(), projected_batch.second.get()).ValueOrDie(); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields[1]}), + R"([["one"],["two"]])") + .ValueOrDie(); + ASSERT_TRUE(projected_array->Equals(expected)) << projected_array->ToString(); +} + +TEST_F(MosaicFileFormatTest, WriterOptions) { + std::map options = { + {"file.format", "mosaic"}, {Options::FILE_BLOCK_SIZE, "1 B"}, + {MOSAIC_NUM_BUCKETS, "2"}, {MOSAIC_MAX_DICT_TOTAL_BYTES, "1 KB"}, + {MOSAIC_MAX_DICT_ENTRIES, "2"}, {MOSAIC_PAGE_SIZE_THRESHOLD, "1 B"}, + }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr configured_format, + FileFormatFactory::Get("mosaic", options)); + std::string path = PathUtil::JoinPath(directory_->Str(), "writer-options.mosaic"); + arrow::FieldVector fields = { + arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::int32()), + arrow::field("f2", arrow::int32()), arrow::field("f3", arrow::int32()), + arrow::field("f4", arrow::int32()), arrow::field("f5", arrow::int32()), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[1,2,3,4,5,6], + [7,8,9,10,11,12], + [13,14,15,16,17,18], + [19,20,21,22,23,24], + [25,26,27,28,29,30]])") + .ValueOrDie(); + + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2, configured_format)); + ASSERT_OK_AND_ASSIGN(FooterLayout footer_layout, ReadFooterLayout(path)); + ASSERT_EQ(footer_layout.first, 2); + ASSERT_EQ(footer_layout.second, 3); + AssertReadWithBatchSizes(path, schema, expected, {10}); +} + +TEST_F(MosaicFileFormatTest, ExtractStatistics) { + std::map options = { + {"file.format", "mosaic"}, + {Options::FILE_BLOCK_SIZE, "1 B"}, + {MOSAIC_STATS_COLUMNS, " id, name, ts, amount, all_null "}, + }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr configured_format, + FileFormatFactory::Get("mosaic", options)); + std::string path = PathUtil::JoinPath(directory_->Str(), "statistics.mosaic"); + arrow::FieldVector fields = { + arrow::field("id", arrow::int32()), + arrow::field("name", arrow::utf8()), + arrow::field("ts", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("amount", arrow::decimal128(10, 2)), + arrow::field("untracked", arrow::int64()), + arrow::field("all_null", arrow::int32()), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + [3,"three","1970-01-01 00:00:00.000000003","3.30",30,null], + [1,"one","1970-01-01 00:00:00.000000001","1.10",10,null], + [null,null,null,null,null,null], + [5,"five","1970-01-01 00:00:00.000000005","5.50",50,null], + [2,"","1970-01-01 00:00:00.000000002","2.20",20,null] + ])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2, configured_format)); + + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr extractor, + configured_format->CreateStatsExtractor(&ffi_schema)); + ASSERT_OK_AND_ASSIGN(auto result, extractor->ExtractWithFileInfo(file_system_, path, pool_)); + ASSERT_EQ(result.second.GetRowCount(), 5); + std::vector expected_stats = { + "min 1, max 5, null count 1", + "min , max three, null count 1", + "min 1970-01-01 00:00:00.000000001, max 1970-01-01 00:00:00.000000005, null count 1", + "min 1.10, max 5.50, null count 1", + "min null, max null, null count null", + "min null, max null, null count 5", + }; + ASSERT_EQ(result.first.size(), expected_stats.size()); + for (size_t i = 0; i < expected_stats.size(); ++i) { + ASSERT_EQ(result.first[i]->ToString(), expected_stats[i]); + } +} + +TEST_F(MosaicFileFormatTest, RowGroupPredicateFiltering) { + std::map options = { + {"file.format", "mosaic"}, + {Options::FILE_BLOCK_SIZE, "1 B"}, + {MOSAIC_STATS_COLUMNS, "id"}, + }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr configured_format, + FileFormatFactory::Get("mosaic", options)); + std::string path = PathUtil::JoinPath(directory_->Str(), "predicate.mosaic"); + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("untracked", arrow::int32())}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields), R"([[1,null],[2,2],[10,10],[11,11],[20,20],[21,21]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2, configured_format)); + ASSERT_OK_AND_ASSIGN(FooterLayout footer_layout, ReadFooterLayout(path)); + ASSERT_EQ(footer_layout.second, 3); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + configured_format->CreateReaderBuilder(/*batch_size=*/10)); + reader_builder->WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(path)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, reader_builder->Build(input)); + + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(15)); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + std::shared_ptr actual = + arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + ASSERT_TRUE(actual->Equals(data->Slice(/*offset=*/4, /*length=*/2))) << actual->ToString(); + ASSERT_OK_AND_ASSIGN(uint64_t first_row, reader->GetPreviousBatchFileRowId(/*batch_row_id=*/0)); + ASSERT_EQ(first_row, 4); + ASSERT_OK_AND_ASSIGN(batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); + + ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", + FieldType::INT, Literal(100)); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); + + ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + predicate = PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"untracked", + FieldType::INT, Literal(100)); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_without_stats, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + ASSERT_TRUE(actual_without_stats->Equals(arrow::ChunkedArray(data))) + << actual_without_stats->ToString(); + + ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + predicate = + PredicateBuilder::IsNull(/*field_index=*/1, /*field_name=*/"untracked", FieldType::INT); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_is_null_without_stats, + paimon::test::ReadResultCollector::CollectResult(reader.get())); + ASSERT_TRUE(actual_is_null_without_stats->Equals(arrow::ChunkedArray(data))) + << actual_is_null_without_stats->ToString(); +} + +TEST_F(MosaicFileFormatTest, WriteThenReadSupportedTypes) { + std::string path = PathUtil::JoinPath(directory_->Str(), "supported-types.mosaic"); + arrow::FieldVector fields = { + arrow::field("f0", arrow::boolean()), + arrow::field("f1", arrow::int8()), + arrow::field("f2", arrow::int16()), + arrow::field("f3", arrow::int32()), + arrow::field("f4", arrow::int64()), + arrow::field("f5", arrow::float32()), + arrow::field("f6", arrow::float64()), + arrow::field("f7", arrow::utf8()), + arrow::field("f8", arrow::binary()), + arrow::field("f9", arrow::map(arrow::int8(), arrow::int16())), + arrow::field("f10", arrow::list(arrow::float32())), + arrow::field("f12", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("f13", arrow::date32()), + arrow::field("f14", arrow::decimal128(2, 2)), + arrow::field("f15", arrow::decimal128(30, 2)), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + [true,-128,-32768,-2147483648,-4294967298,0.5,1.141592659,"mosaic","binary", + [[-1,-2],[1,2]],[1.5,null],"1970-01-01 00:00:00.000000001",-1,"-0.99","-123456789987654321.45"], + [false,127,32767,2147483647,4294967296,2.0,3.141592657,"", "", + [],[],"2030-12-31 23:59:59.999999999",12345,"0.78","123456789987654321.45"], + [null,null,null,null,null,null,null,null,null,null,null,null,null,null,null] + ])") + .ValueOrDie(); + + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2)); + AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5}); +} + +TEST_F(MosaicFileFormatTest, WriteThenReadNestedTypes) { + std::string path = PathUtil::JoinPath(directory_->Str(), "nested.mosaic"); + std::shared_ptr list_type = arrow::list(arrow::int32()); + std::shared_ptr map_type = arrow::map(arrow::utf8(), arrow::int64()); + std::shared_ptr nested_list_type = arrow::list(list_type); + std::shared_ptr list_of_maps_type = + arrow::list(arrow::map(arrow::utf8(), arrow::int32())); + std::shared_ptr map_of_lists_type = + arrow::map(arrow::utf8(), arrow::list(arrow::int32())); + arrow::FieldVector fields = { + arrow::field("id", arrow::int32(), false), + arrow::field("list_col", list_type), + arrow::field("map_col", map_type), + arrow::field("nested_list_col", nested_list_type), + arrow::field("list_of_maps_col", list_of_maps_type), + arrow::field("map_of_lists_col", map_of_lists_type), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([ + [1,[1,null,3],[["a",10],["b",null]],[[1,2],[3]], + [[["a",1]],[["b",2],["c",3]]],[["x",[1,2]],["y",[]]]], + [2,[],[],[[4]],[],[]], + [3,null,null,null,null,null], + [4,[4,5],[["c",30]],[[],[5,null]],[[]],[["z",[null,5]]]], + [5,[6],[["d",40],["e",50]],[],[[["d",4]]],[["w",[]]]] + ])") + .ValueOrDie(); + + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2)); + AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5, 8}); +} + +TEST_F(MosaicFileFormatTest, WriteThenReadTimestampTypes) { + const std::string timezone = "Asia/Tokyo"; + paimon::test::TimezoneGuard timezone_guard(timezone); + std::string path = PathUtil::JoinPath(directory_->Str(), "timestamp-types.mosaic"); + arrow::FieldVector fields = { + arrow::field("ts_milli", arrow::timestamp(arrow::TimeUnit::MILLI)), + arrow::field("ts_micro", arrow::timestamp(arrow::TimeUnit::MICRO)), + arrow::field("ts_nano", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("ts_tz_milli", arrow::timestamp(arrow::TimeUnit::MILLI, timezone)), + arrow::field("ts_tz_micro", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)), + arrow::field("ts_tz_nano", arrow::timestamp(arrow::TimeUnit::NANO, timezone)), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + ["1970-01-01 00:00:00.001","1970-01-01 00:00:00.000001","1970-01-01 00:00:00.000000001", + "1970-01-01 00:00:00.002","1970-01-01 00:00:00.000002","1970-01-01 00:00:00.000000002"], + ["2030-12-31 23:59:59.999","2030-12-31 23:59:59.999999","2030-12-31 23:59:59.999999999", + "2031-01-01 00:00:00.001","2031-01-01 00:00:00.000001","2031-01-01 00:00:00.000000001"], + [null,null,null,null,null,null] + ])") + .ValueOrDie(); + + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2)); + AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5}); +} + +TEST_F(MosaicFileFormatTest, RejectUnsupportedTimestampSecond) { + std::string path = PathUtil::JoinPath(directory_->Str(), "timestamp-second.mosaic"); + arrow::FieldVector fields = { + arrow::field("ts_second", arrow::timestamp(arrow::TimeUnit::SECOND)), + }; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([["1970-01-01 00:00:01"]])") + .ValueOrDie(); + + ASSERT_NOK_WITH_MSG(WriteFile(path, schema, array, /*batch_size=*/1), + "unsupported Timestamp unit: Second"); +} + +TEST_F(MosaicFileFormatTest, RejectUnsupportedStructType) { + std::string path = PathUtil::JoinPath(directory_->Str(), "struct.mosaic"); + std::shared_ptr struct_type = arrow::struct_( + {arrow::field("value", arrow::int32()), arrow::field("label", arrow::utf8())}); + arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), + arrow::field("struct_col", struct_type)}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[1,[10,"ten"]],[2,null]])") + .ValueOrDie(); + + ASSERT_NOK_WITH_MSG(WriteFile(path, schema, array, /*batch_size=*/2), + "unsupported DataType: Struct"); +} + +} // namespace paimon::mosaic::test diff --git a/src/paimon/format/mosaic/mosaic_format_defs.h b/src/paimon/format/mosaic/mosaic_format_defs.h new file mode 100644 index 000000000..4b5265edf --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_format_defs.h @@ -0,0 +1,40 @@ +/* + * 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::mosaic { + +/// Number of column buckets for parallel IO. +static inline const char MOSAIC_NUM_BUCKETS[] = "mosaic.num-buckets"; + +/// Max dict size per column. +static inline const char MOSAIC_MAX_DICT_TOTAL_BYTES[] = "mosaic.max-dict-total-bytes"; + +/// Max dict entries per column. +static inline const char MOSAIC_MAX_DICT_ENTRIES[] = "mosaic.max-dict-entries"; + +/// Min avg column page size to enable paged mode. +static inline const char MOSAIC_PAGE_SIZE_THRESHOLD[] = "mosaic.page-size-threshold"; + +/// Comma-separated list of column names to collect statistics for. Empty means no statistics +/// collection. +static inline const char MOSAIC_STATS_COLUMNS[] = "mosaic.stats-columns"; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_format_writer.cpp b/src/paimon/format/mosaic/mosaic_format_writer.cpp new file mode 100644 index 000000000..828ec565b --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_format_writer.cpp @@ -0,0 +1,131 @@ +/* + * 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/format/mosaic/mosaic_format_writer.h" + +#include + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/fs/file_system.h" + +namespace paimon::mosaic { + +MosaicFormatWriter::MosaicFormatWriter(const std::shared_ptr& output, + const std::shared_ptr& schema, + std::unique_ptr output_context, + MosaicWriterHandle* writer) + : output_(output), + schema_(schema), + output_context_(std::move(output_context)), + writer_(writer), + metrics_(std::make_shared()) {} + +Result> MosaicFormatWriter::Create( + const std::shared_ptr& output, const std::shared_ptr& schema, + const MosaicWriterOptions& options) { + if (output == nullptr || schema == nullptr) { + return Status::Invalid("Mosaic writer requires non-null output and schema"); + } + auto output_context = std::make_unique(output); + MosaicOutputFile output_file = {}; + output_file.ctx = output_context.get(); + output_file.write_fn = MosaicOutputContext::Write; + output_file.flush_fn = MosaicOutputContext::Flush; + output_file.get_pos_fn = MosaicOutputContext::GetPos; + + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + MosaicWriterHandle* writer = mosaic_writer_open(output_file, &ffi_schema, options); + if (writer == nullptr) { + Status status = MosaicFfiError("open Mosaic writer", output_context->GetCallbackStatus()); + ArrowSchemaRelease(&ffi_schema); + return status; + } + return std::unique_ptr( + new MosaicFormatWriter(output, schema, std::move(output_context), writer)); +} + +MosaicFormatWriter::~MosaicFormatWriter() { + if (writer_ != nullptr) { + if (!finished_) { + mosaic_writer_close(writer_); + } + mosaic_writer_free(writer_); + } +} + +Status MosaicFormatWriter::AddBatch(::ArrowArray* batch) { + if (batch == nullptr) { + return Status::Invalid("Mosaic writer batch is nullptr"); + } + if (finished_) { + return Status::Invalid("cannot add a batch after Mosaic writer is finished"); + } + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, &ffi_schema)); + if (mosaic_writer_write_batch(writer_, batch, &ffi_schema) != 0) { + Status status = MosaicFfiError("write Mosaic batch", output_context_->GetCallbackStatus()); + ArrowSchemaRelease(&ffi_schema); + return status; + } + return Status::OK(); +} + +Status MosaicFormatWriter::Flush() { + if (finished_) { + return Status::OK(); + } + return output_->Flush(); +} + +Status MosaicFormatWriter::Finish() { + if (finished_) { + return Status::OK(); + } + if (mosaic_writer_close(writer_) != 0) { + return MosaicFfiError("finish Mosaic writer", output_context_->GetCallbackStatus()); + } + finished_ = true; + return Status::OK(); +} + +Result MosaicFormatWriter::ReachTargetSize(bool suggested_check, int64_t target_size) const { + if (!suggested_check) { + return false; + } + int64_t estimated_size = 0; + if (mosaic_writer_estimated_file_size(writer_, &estimated_size) != 0) { + return MosaicFfiError("estimate Mosaic file size", output_context_->GetCallbackStatus()); + } + return estimated_size >= target_size; +} + +std::shared_ptr MosaicFormatWriter::GetWriterMetrics() const { + return metrics_; +} + +Status MosaicFormatWriter::AddMetadata(const std::map& metadata) { + (void)metadata; + return Status::NotImplemented("Mosaic writer metadata is not supported"); +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_format_writer.h b/src/paimon/format/mosaic/mosaic_format_writer.h new file mode 100644 index 000000000..b96aba887 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_format_writer.h @@ -0,0 +1,70 @@ +/* + * 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/format/format_writer.h" +#include "paimon/format/mosaic/mosaic_ffi.h" +#include "paimon/format/mosaic/mosaic_stream.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow +namespace paimon { +class Metrics; +class OutputStream; +} // namespace paimon + +namespace paimon::mosaic { + +class MosaicFormatWriter : public FormatWriter { + public: + static Result> Create( + const std::shared_ptr& output, const std::shared_ptr& schema, + const MosaicWriterOptions& options); + + ~MosaicFormatWriter() override; + + Status AddBatch(::ArrowArray* batch) override; + Status Flush() override; + Status Finish() override; + Result ReachTargetSize(bool suggested_check, int64_t target_size) const override; + std::shared_ptr GetWriterMetrics() const override; + Status AddMetadata(const std::map& metadata) override; + + private: + MosaicFormatWriter(const std::shared_ptr& output, + const std::shared_ptr& schema, + std::unique_ptr output_context, + MosaicWriterHandle* writer); + + std::shared_ptr output_; + std::shared_ptr schema_; + std::unique_ptr output_context_; + MosaicWriterHandle* writer_; + std::shared_ptr metrics_; + bool finished_ = false; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_reader_builder.h b/src/paimon/format/mosaic/mosaic_reader_builder.h new file mode 100644 index 000000000..ac1afba92 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_reader_builder.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 "paimon/format/mosaic/mosaic_file_batch_reader.h" +#include "paimon/format/reader_builder.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon::mosaic { + +class MosaicReaderBuilder : public ReaderBuilder { + public: + explicit MosaicReaderBuilder(int32_t batch_size) + : batch_size_(batch_size), pool_(GetDefaultPool()) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + pool_ = pool; + return this; + } + + Result> Build( + const std::shared_ptr& input) const override { + if (pool_ == nullptr) { + return Status::Invalid("Mosaic reader memory pool is nullptr"); + } + return MosaicFileBatchReader::Create(input, batch_size_, pool_); + } + + private: + int32_t batch_size_; + std::shared_ptr pool_; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stats.cpp b/src/paimon/format/mosaic/mosaic_stats.cpp new file mode 100644 index 000000000..17b788200 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stats.cpp @@ -0,0 +1,311 @@ +/* + * 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/format/mosaic/mosaic_stats.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/math.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/format/column_stats.h" +#include "paimon/format/mosaic/mosaic_stream.h" +#include "paimon/io/byte_order.h" + +namespace paimon::mosaic { + +namespace { + +template +using OptionalMinMax = std::pair, std::optional>; + +template +Result DecodeBigEndian(const std::vector& bytes) { + if (bytes.size() != sizeof(T)) { + return Status::Invalid( + fmt::format("invalid Mosaic statistic size {}, expected {}", bytes.size(), sizeof(T))); + } + T value; + std::memcpy(&value, bytes.data(), sizeof(T)); + if constexpr (SystemByteOrder() == ByteOrder::PAIMON_LITTLE_ENDIAN) { + value = EndianSwapValue(value); + } + return value; +} + +template > +Result, std::optional>> CollectMinMax( + const std::vector& stats, Decoder decoder, + Less less = Less()) { + std::optional min; + std::optional max; + for (const MosaicStatsUtils::ColumnStatistics* stat : stats) { + if (stat->min.has_value() != stat->max.has_value()) { + return Status::Invalid("Mosaic statistics contain incomplete min/max values"); + } + if (!stat->min.has_value()) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(T row_group_min, decoder(stat->min.value())); + PAIMON_ASSIGN_OR_RAISE(T row_group_max, decoder(stat->max.value())); + if (!min.has_value() || less(row_group_min, min.value())) { + min = std::move(row_group_min); + } + if (!max.has_value() || less(max.value(), row_group_max)) { + max = std::move(row_group_max); + } + } + return std::make_pair(std::move(min), std::move(max)); +} + +Result> CollectNullCount( + const std::vector& stats, + bool missing_null_count_is_zero) { + if (stats.empty()) { + return missing_null_count_is_zero ? std::optional(0) : std::nullopt; + } + uint64_t result = 0; + for (const MosaicStatsUtils::ColumnStatistics* stat : stats) { + if (stat->null_count > + static_cast(std::numeric_limits::max()) - result) { + return Status::Invalid("Mosaic null count exceeds int64 range"); + } + result += stat->null_count; + } + return std::optional(static_cast(result)); +} + +Result DecodeTimestamp(const std::vector& bytes, + const std::shared_ptr& type) { + if (type->unit() == arrow::TimeUnit::NANO) { + if (bytes.size() != 12) { + return Status::Invalid( + fmt::format("invalid Mosaic nanosecond timestamp statistic size {}", bytes.size())); + } + std::vector millis_bytes(bytes.begin(), bytes.begin() + 8); + std::vector nanos_bytes(bytes.begin() + 8, bytes.end()); + PAIMON_ASSIGN_OR_RAISE(int64_t millis, DecodeBigEndian(millis_bytes)); + PAIMON_ASSIGN_OR_RAISE(int32_t nanos, DecodeBigEndian(nanos_bytes)); + if (nanos < 0 || nanos > 999999) { + return Status::Invalid("invalid Mosaic nanosecond timestamp statistic"); + } + return Timestamp(millis, nanos); + } + PAIMON_ASSIGN_OR_RAISE(int64_t value, DecodeBigEndian(bytes)); + auto [millis, nanos] = DateTimeUtils::TimestampConverter( + value, DateTimeUtils::GetTimeTypeFromArrowType(type), DateTimeUtils::TimeType::MILLISECOND, + DateTimeUtils::TimeType::NANOSECOND); + return Timestamp(millis, static_cast(nanos)); +} + +Result> ConvertFieldStatistics( + const std::shared_ptr& type, + const std::vector& stats, + bool missing_null_count_is_zero) { + PAIMON_ASSIGN_OR_RAISE(std::optional null_count, + CollectNullCount(stats, missing_null_count_is_zero)); + switch (type->id()) { + case arrow::Type::BOOL: { + auto decoder = [](const std::vector& bytes) -> Result { + if (bytes.size() != 1 || bytes[0] > 1) { + return Status::Invalid("invalid Mosaic boolean statistic"); + } + return bytes[0] != 0; + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, decoder)); + return ColumnStats::CreateBooleanColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::INT8: { + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian)); + return ColumnStats::CreateTinyIntColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::INT16: { + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian)); + return ColumnStats::CreateSmallIntColumnStats(min_max.first, min_max.second, + null_count); + } + case arrow::Type::INT32: { + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian)); + return ColumnStats::CreateIntColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::INT64: { + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian)); + return ColumnStats::CreateBigIntColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::FLOAT: { + auto less = [](float lhs, float rhs) { + return FieldsComparator::CompareFloatingPoint(lhs, rhs) < 0; + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian, less)); + return ColumnStats::CreateFloatColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::DOUBLE: { + auto less = [](double lhs, double rhs) { + return FieldsComparator::CompareFloatingPoint(lhs, rhs) < 0; + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian, less)); + return ColumnStats::CreateDoubleColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::STRING: { + auto decoder = [](const std::vector& bytes) -> Result { + if (bytes.empty()) { + return std::string(); + } + return std::string(reinterpret_cast(bytes.data()), bytes.size()); + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, decoder)); + return ColumnStats::CreateStringColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::BINARY: + return ColumnStats::CreateStringColumnStats(std::nullopt, std::nullopt, null_count); + case arrow::Type::DATE32: { + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, DecodeBigEndian)); + return ColumnStats::CreateDateColumnStats(min_max.first, min_max.second, null_count); + } + case arrow::Type::TIMESTAMP: { + auto timestamp_type = checked_pointer_cast(type); + auto decoder = [×tamp_type](const std::vector& bytes) { + return DecodeTimestamp(bytes, timestamp_type); + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, decoder)); + return ColumnStats::CreateTimestampColumnStats( + min_max.first, min_max.second, null_count, + DateTimeUtils::GetPrecisionFromType(timestamp_type)); + } + case arrow::Type::DECIMAL128: { + auto decimal_type = checked_pointer_cast(type); + if (decimal_type->precision() > 18) { + return ColumnStats::CreateDecimalColumnStats(std::nullopt, std::nullopt, null_count, + decimal_type->precision(), + decimal_type->scale()); + } + auto decoder = [&decimal_type](const std::vector& bytes) -> Result { + PAIMON_ASSIGN_OR_RAISE(int64_t value, DecodeBigEndian(bytes)); + return Decimal::FromUnscaledLong(value, decimal_type->precision(), + decimal_type->scale()); + }; + PAIMON_ASSIGN_OR_RAISE(OptionalMinMax min_max, + CollectMinMax(stats, decoder)); + return ColumnStats::CreateDecimalColumnStats(min_max.first, min_max.second, null_count, + decimal_type->precision(), + decimal_type->scale()); + } + case arrow::Type::LIST: + return ColumnStats::CreateNestedColumnStats(FieldType::ARRAY, null_count); + case arrow::Type::MAP: + return ColumnStats::CreateNestedColumnStats(FieldType::MAP, null_count); + case arrow::Type::STRUCT: + return ColumnStats::CreateNestedColumnStats(FieldType::STRUCT, null_count); + default: + return Status::Invalid( + fmt::format("cannot fetch Mosaic statistics for type {}", type->ToString())); + } +} + +} // namespace + +Result MosaicStatsUtils::ReadRowGroupStatistics( + uint32_t row_group, const MosaicInputContext* input_context, MosaicReaderHandle* reader) { + uint32_t num_stats = 0; + if (mosaic_reader_row_group_num_stats(reader, row_group, &num_stats) != 0) { + return MosaicFfiError("read Mosaic row group statistic count", + input_context->GetCallbackStatus()); + } + if (num_stats == 0) { + return RowGroupStatistics(); + } + std::vector names(num_stats); + std::vector null_counts(num_stats); + std::vector min_ptrs(num_stats); + std::vector min_lens(num_stats); + std::vector max_ptrs(num_stats); + std::vector max_lens(num_stats); + if (mosaic_reader_row_group_stats(reader, row_group, names.data(), null_counts.data(), + min_ptrs.data(), min_lens.data(), max_ptrs.data(), + max_lens.data()) != 0) { + return MosaicFfiError("read Mosaic row group statistics", + input_context->GetCallbackStatus()); + } + RowGroupStatistics result; + result.reserve(num_stats); + for (uint32_t i = 0; i < num_stats; ++i) { + if (names[i] == nullptr || (min_ptrs[i] == nullptr) != (max_ptrs[i] == nullptr) || + (min_ptrs[i] == nullptr && (min_lens[i] != 0 || max_lens[i] != 0))) { + return Status::Invalid("invalid Mosaic row group statistics"); + } + ColumnStatistics stats = {null_counts[i], std::nullopt, std::nullopt}; + if (min_ptrs[i] != nullptr) { + stats.min = min_lens[i] == 0 + ? std::vector() + : std::vector(min_ptrs[i], min_ptrs[i] + min_lens[i]); + stats.max = max_lens[i] == 0 + ? std::vector() + : std::vector(max_ptrs[i], max_ptrs[i] + max_lens[i]); + } + auto [iter, inserted] = result.emplace(names[i], std::move(stats)); + if (!inserted) { + return Status::Invalid( + fmt::format("duplicate Mosaic statistics for column {}", iter->first)); + } + } + return result; +} + +Result MosaicStatsUtils::ConvertColumnStatistics( + const std::shared_ptr& schema, + const std::vector& row_group_stats, bool missing_null_count_is_zero) { + ColumnStatsVector result; + result.reserve(schema->num_fields()); + for (const std::shared_ptr& field : schema->fields()) { + std::vector field_stats; + field_stats.reserve(row_group_stats.size()); + for (const RowGroupStatistics& stats : row_group_stats) { + auto iter = stats.find(field->name()); + if (iter != stats.end()) { + field_stats.push_back(&iter->second); + } + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr column_stats, + ConvertFieldStatistics(field->type(), field_stats, missing_null_count_is_zero)); + result.push_back(std::move(column_stats)); + } + return result; +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stats.h b/src/paimon/format/mosaic/mosaic_stats.h new file mode 100644 index 000000000..d5350edd7 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stats.h @@ -0,0 +1,62 @@ +/* + * 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 + +#include "paimon/format/mosaic/mosaic_ffi.h" +#include "paimon/result.h" +#include "paimon/type_fwd.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon::mosaic { + +class MosaicInputContext; + +class MosaicStatsUtils { + public: + struct ColumnStatistics { + uint64_t null_count; + std::optional> min; + std::optional> max; + }; + + using RowGroupStatistics = std::unordered_map; + + MosaicStatsUtils() = delete; + ~MosaicStatsUtils() = delete; + + static Result ReadRowGroupStatistics( + uint32_t row_group, const MosaicInputContext* input_context, MosaicReaderHandle* reader); + + static Result ConvertColumnStatistics( + const std::shared_ptr& schema, + const std::vector& row_group_stats, bool missing_null_count_is_zero); +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stats_extractor.cpp b/src/paimon/format/mosaic/mosaic_stats_extractor.cpp new file mode 100644 index 000000000..bef972848 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stats_extractor.cpp @@ -0,0 +1,84 @@ +/* + * 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/format/mosaic/mosaic_stats_extractor.h" + +#include +#include +#include + +#include "paimon/common/utils/math.h" +#include "paimon/format/mosaic/mosaic_ffi.h" +#include "paimon/format/mosaic/mosaic_stats.h" +#include "paimon/format/mosaic/mosaic_stream.h" +#include "paimon/fs/file_system.h" + +namespace paimon::mosaic { + +Result> +MosaicStatsExtractor::ExtractWithFileInfo(const std::shared_ptr& file_system, + const std::string& path, + const std::shared_ptr& pool) { + if (file_system == nullptr || pool == nullptr) { + return Status::Invalid("Mosaic stats extractor requires file system and memory pool"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, file_system->Open(path)); + PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Mosaic input length")); + auto input_context = + std::make_unique(input, static_cast(signed_length)); + MosaicInputFile input_file = {}; + input_file.ctx = input_context.get(); + input_file.read_at_fn = MosaicInputContext::ReadAt; + input_file.length_fn = MosaicInputContext::Length; + std::unique_ptr reader( + mosaic_reader_open(input_file), mosaic_reader_free); + if (reader == nullptr) { + return MosaicFfiError("open Mosaic reader", input_context->GetCallbackStatus()); + } + + uint32_t num_row_groups = 0; + if (mosaic_reader_num_row_groups(reader.get(), &num_row_groups) != 0) { + return MosaicFfiError("read Mosaic row group count", input_context->GetCallbackStatus()); + } + int64_t row_count = 0; + std::vector row_group_stats; + row_group_stats.reserve(num_row_groups); + for (uint32_t row_group = 0; row_group < num_row_groups; ++row_group) { + uint32_t row_group_row_count = 0; + if (mosaic_reader_row_group_num_rows(reader.get(), row_group, &row_group_row_count) != 0) { + return MosaicFfiError("read Mosaic row count", input_context->GetCallbackStatus()); + } + if (row_group_row_count > + static_cast(std::numeric_limits::max() - row_count)) { + return Status::Invalid("Mosaic row count exceeds int64 range"); + } + row_count += row_group_row_count; + PAIMON_ASSIGN_OR_RAISE( + MosaicStatsUtils::RowGroupStatistics stats, + MosaicStatsUtils::ReadRowGroupStatistics(row_group, input_context.get(), reader.get())); + row_group_stats.push_back(std::move(stats)); + } + PAIMON_ASSIGN_OR_RAISE(ColumnStatsVector stats, + MosaicStatsUtils::ConvertColumnStatistics( + schema_, row_group_stats, /*missing_null_count_is_zero=*/false)); + return std::make_pair(std::move(stats), FileInfo(row_count)); +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stats_extractor.h b/src/paimon/format/mosaic/mosaic_stats_extractor.h new file mode 100644 index 000000000..cfa064490 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stats_extractor.h @@ -0,0 +1,55 @@ +/* + * 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/format/format_stats_extractor.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon::mosaic { + +class MosaicStatsExtractor : public FormatStatsExtractor { + public: + explicit MosaicStatsExtractor(const std::shared_ptr& schema) : schema_(schema) {} + + Result Extract(const std::shared_ptr& file_system, + const std::string& path, + const std::shared_ptr& pool) override { + using ExtractResult = std::pair; + PAIMON_ASSIGN_OR_RAISE(ExtractResult result, ExtractWithFileInfo(file_system, path, pool)); + return result.first; + } + + Result> ExtractWithFileInfo( + const std::shared_ptr& file_system, const std::string& path, + const std::shared_ptr& pool) override; + + private: + std::shared_ptr schema_; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stream.cpp b/src/paimon/format/mosaic/mosaic_stream.cpp new file mode 100644 index 000000000..87225fe99 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stream.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/format/mosaic/mosaic_stream.h" + +#include + +#include "paimon/common/utils/math.h" +#include "paimon/format/mosaic/mosaic_ffi.h" +#include "paimon/fs/file_system.h" +#include "paimon/result.h" + +namespace paimon::mosaic { + +void MosaicInputContext::SetCallbackStatus(const Status& status) { + std::lock_guard lock(mutex_); + if (callback_status_.ok()) { + callback_status_ = status; + } +} + +Status MosaicInputContext::GetCallbackStatus() const { + std::lock_guard lock(mutex_); + return callback_status_; +} + +int32_t MosaicInputContext::ReadAt(void* context, uint64_t offset, uint8_t* buffer, + size_t length) noexcept { + auto* input_context = static_cast(context); + if (input_context == nullptr || buffer == nullptr) { + if (input_context != nullptr) { + input_context->SetCallbackStatus(Status::Invalid("invalid Mosaic read request")); + } + return -1; + } + Status status = ValidateValueInRange(offset, "Mosaic read offset"); + if (status.ok()) { + status = ValidateValueInRange(length, "Mosaic read length"); + } + if (!status.ok()) { + input_context->SetCallbackStatus(status); + return -1; + } + auto read_length = static_cast(length); + auto read_offset = static_cast(offset); + Result result = + input_context->input_->Read(reinterpret_cast(buffer), read_length, read_offset); + if (!result.ok()) { + input_context->SetCallbackStatus(result.status()); + return -1; + } + int64_t bytes_read = std::move(result).value(); + if (bytes_read != read_length) { + input_context->SetCallbackStatus(Status::IOError("short read while reading Mosaic file")); + return -1; + } + return 0; +} + +uint64_t MosaicInputContext::Length(void* context) noexcept { + auto* input_context = static_cast(context); + return input_context == nullptr ? 0 : input_context->length_; +} + +void MosaicOutputContext::SetCallbackStatus(const Status& status) { + std::lock_guard lock(mutex_); + if (callback_status_.ok()) { + callback_status_ = status; + } +} + +Status MosaicOutputContext::GetCallbackStatus() const { + std::lock_guard lock(mutex_); + return callback_status_; +} + +int32_t MosaicOutputContext::Write(void* context, const uint8_t* data, size_t length) noexcept { + auto* output_context = static_cast(context); + if (output_context == nullptr || data == nullptr) { + if (output_context != nullptr) { + output_context->SetCallbackStatus(Status::Invalid("invalid Mosaic write request")); + } + return -1; + } + Status status = ValidateValueInRange(length, "Mosaic write length"); + if (!status.ok()) { + output_context->SetCallbackStatus(status); + return -1; + } + auto write_length = static_cast(length); + Result result = + output_context->output_->Write(reinterpret_cast(data), write_length); + if (!result.ok()) { + output_context->SetCallbackStatus(result.status()); + return -1; + } + int64_t bytes_written = std::move(result).value(); + if (bytes_written != write_length) { + output_context->SetCallbackStatus(Status::IOError("short write while writing Mosaic file")); + return -1; + } + return 0; +} + +int32_t MosaicOutputContext::Flush(void* context) noexcept { + auto* output_context = static_cast(context); + if (output_context == nullptr) { + return -1; + } + Status status = output_context->output_->Flush(); + if (!status.ok()) { + output_context->SetCallbackStatus(status); + return -1; + } + return 0; +} + +int64_t MosaicOutputContext::GetPos(void* context) noexcept { + auto* output_context = static_cast(context); + if (output_context == nullptr) { + return -1; + } + Result result = output_context->output_->GetPos(); + if (!result.ok()) { + output_context->SetCallbackStatus(result.status()); + return -1; + } + return std::move(result).value(); +} + +Status MosaicFfiError(const std::string& operation, const Status& callback_status) { + if (!callback_status.ok()) { + return callback_status.WithMessage(operation, ": ", callback_status.message()); + } + const char* error = mosaic_last_error(); + return Status::Invalid(operation, ": ", error == nullptr ? "unknown Mosaic error" : error); +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_stream.h b/src/paimon/format/mosaic/mosaic_stream.h new file mode 100644 index 000000000..93f27a4f1 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_stream.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 +#include +#include + +#include "paimon/status.h" + +namespace paimon { +class InputStream; +class OutputStream; +} // namespace paimon + +namespace paimon::mosaic { + +class MosaicInputContext { + public: + MosaicInputContext(const std::shared_ptr& input, uint64_t length) + : input_(input), length_(length) {} + + static int32_t ReadAt(void* context, uint64_t offset, uint8_t* buffer, size_t length) noexcept; + static uint64_t Length(void* context) noexcept; + + Status GetCallbackStatus() const; + + private: + void SetCallbackStatus(const Status& status); + + std::shared_ptr input_; + uint64_t length_; + mutable std::mutex mutex_; + Status callback_status_; +}; + +class MosaicOutputContext { + public: + explicit MosaicOutputContext(const std::shared_ptr& output) : output_(output) {} + + static int32_t Write(void* context, const uint8_t* data, size_t length) noexcept; + static int32_t Flush(void* context) noexcept; + static int64_t GetPos(void* context) noexcept; + + Status GetCallbackStatus() const; + + private: + void SetCallbackStatus(const Status& status); + + std::shared_ptr output_; + mutable std::mutex mutex_; + Status callback_status_; +}; + +Status MosaicFfiError(const std::string& operation, const Status& callback_status); + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_writer_builder.cpp b/src/paimon/format/mosaic/mosaic_writer_builder.cpp new file mode 100644 index 000000000..20b9407fd --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_writer_builder.cpp @@ -0,0 +1,104 @@ +/* + * 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/format/mosaic/mosaic_writer_builder.h" + +#include +#include +#include + +#include "arrow/type.h" +#include "paimon/common/options/memory_size.h" +#include "paimon/common/utils/math.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/defs.h" +#include "paimon/format/mosaic/mosaic_format_defs.h" + +namespace paimon::mosaic { + +Result> MosaicWriterBuilder::Build( + const std::shared_ptr& output, const std::string& compression) { + if (pool_ == nullptr) { + return Status::Invalid("Mosaic writer memory pool is nullptr"); + } + std::string normalized = StringUtils::ToLowerCase(compression); + uint8_t compression_id = 0; + if (normalized == "zstd" || normalized == "zstandard") { + compression_id = 1; + } else if (normalized != "none" && normalized != "null" && normalized != "uncompressed") { + return Status::Invalid("unknown Mosaic compression ", compression); + } + MosaicWriterOptions writer_options = mosaic_writer_options_default(); + writer_options.compression = compression_id; + if (compression_id == 1) { + PAIMON_ASSIGN_OR_RAISE( + writer_options.zstd_level, + OptionsUtils::GetValueFromMap(options_, Options::FILE_COMPRESSION_ZSTD_LEVEL, + writer_options.zstd_level)); + } + PAIMON_ASSIGN_OR_RAISE(writer_options.num_buckets, + OptionsUtils::GetValueFromMap(options_, MOSAIC_NUM_BUCKETS, + writer_options.num_buckets)); + auto max_dict_total_bytes = options_.find(MOSAIC_MAX_DICT_TOTAL_BYTES); + if (max_dict_total_bytes != options_.end()) { + PAIMON_ASSIGN_OR_RAISE(int64_t value, MemorySize::ParseBytes(max_dict_total_bytes->second)); + PAIMON_RETURN_NOT_OK(ValidateValueInRange(value, "Mosaic max dict total bytes")); + writer_options.max_dict_total_bytes = static_cast(value); + } + PAIMON_ASSIGN_OR_RAISE(writer_options.max_dict_entries, + OptionsUtils::GetValueFromMap( + options_, MOSAIC_MAX_DICT_ENTRIES, writer_options.max_dict_entries)); + auto page_size_threshold = options_.find(MOSAIC_PAGE_SIZE_THRESHOLD); + if (page_size_threshold != options_.end()) { + PAIMON_ASSIGN_OR_RAISE(int64_t value, MemorySize::ParseBytes(page_size_threshold->second)); + PAIMON_RETURN_NOT_OK(ValidateValueInRange(value, "Mosaic page size threshold")); + writer_options.page_size_threshold = static_cast(value); + } + auto block_size = options_.find(Options::FILE_BLOCK_SIZE); + if (block_size != options_.end()) { + PAIMON_ASSIGN_OR_RAISE(int64_t row_group_max_size, + MemorySize::ParseBytes(block_size->second)); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(row_group_max_size, "Mosaic row group max size")); + writer_options.row_group_max_size = static_cast(row_group_max_size); + } + std::vector stats_columns; + auto stats_columns_iter = options_.find(MOSAIC_STATS_COLUMNS); + if (stats_columns_iter != options_.end()) { + for (std::string column : StringUtils::Split(stats_columns_iter->second, ",")) { + StringUtils::Trim(&column); + if (!column.empty() && schema_->GetFieldByName(column) != nullptr) { + stats_columns.push_back(std::move(column)); + } + } + } + std::vector stats_column_pointers; + stats_column_pointers.reserve(stats_columns.size()); + for (const std::string& column : stats_columns) { + stats_column_pointers.push_back(column.c_str()); + } + writer_options.stats_columns = stats_column_pointers.data(); + PAIMON_RETURN_NOT_OK( + ValidateValueInRange(stats_column_pointers.size(), "Mosaic stats column count")); + writer_options.num_stats_columns = static_cast(stats_column_pointers.size()); + return MosaicFormatWriter::Create(output, schema_, writer_options); +} + +} // namespace paimon::mosaic diff --git a/src/paimon/format/mosaic/mosaic_writer_builder.h b/src/paimon/format/mosaic/mosaic_writer_builder.h new file mode 100644 index 000000000..d3d40ff57 --- /dev/null +++ b/src/paimon/format/mosaic/mosaic_writer_builder.h @@ -0,0 +1,56 @@ +/* + * 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/format/mosaic/mosaic_format_writer.h" +#include "paimon/format/writer_builder.h" +#include "paimon/memory/memory_pool.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon::mosaic { + +class MosaicWriterBuilder : public WriterBuilder { + public: + MosaicWriterBuilder(const std::shared_ptr& schema, + const std::map& options) + : schema_(schema), options_(options), pool_(GetDefaultPool()) {} + + WriterBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + pool_ = pool; + return this; + } + + Result> Build(const std::shared_ptr& output, + const std::string& compression) override; + + private: + std::shared_ptr schema_; + std::map options_; + std::shared_ptr pool_; +}; + +} // namespace paimon::mosaic diff --git a/src/paimon/format/orc/orc_adapter.cpp b/src/paimon/format/orc/orc_adapter.cpp index 4a03d7cda..dc99c5e68 100644 --- a/src/paimon/format/orc/orc_adapter.cpp +++ b/src/paimon/format/orc/orc_adapter.cpp @@ -940,6 +940,8 @@ Result> OrcAdapter::AppendBatch( MakeArrowBuilder(type, batch, pool)); std::shared_ptr array; PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Finish(&array)); + // Keep this check in release builds so malformed nested arrays return a Status before they + // reach Arrow constructors that enforce their invariants with a process-terminating check. PAIMON_RETURN_NOT_OK_FROM_ARROW(array->Validate()); return array; } diff --git a/src/paimon/format/orc/orc_file_batch_reader.cpp b/src/paimon/format/orc/orc_file_batch_reader.cpp index cd627eb5b..a201839b3 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.h b/src/paimon/format/orc/orc_file_batch_reader.h index 85673a93e..b48f7cdb0 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -85,7 +85,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { return reader_->GetNumberOfRows(); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { return reader_->GetNextRowToRead(); } 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 038c93e2e..c39ba5166 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/orc/orc_format_writer.cpp b/src/paimon/format/orc/orc_format_writer.cpp index 1a394ca94..cc723b9d8 100644 --- a/src/paimon/format/orc/orc_format_writer.cpp +++ b/src/paimon/format/orc/orc_format_writer.cpp @@ -45,6 +45,7 @@ #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/schema/arrow_schema_validator.h" +#include "paimon/defs.h" #include "paimon/format/orc/orc_adapter.h" #include "paimon/format/orc/orc_format_defs.h" #include "paimon/format/orc/orc_memory_pool.h" @@ -236,18 +237,6 @@ Status OrcFormatWriter::AddMetadata(const std::map& me return Status::OK(); } -namespace { - -Result GetMemorySizeOption(const std::map& options, - const std::string& key, uint64_t default_value) { - PAIMON_ASSIGN_OR_RAISE(std::string value, OptionsUtils::GetValueFromMap( - options, key, std::to_string(default_value))); - PAIMON_ASSIGN_OR_RAISE(int64_t bytes, MemorySize::ParseBytes(value)); - return static_cast(bytes); -} - -} // namespace - Result<::orc::WriterOptions> OrcFormatWriter::PrepareWriterOptions( const std::map& options, const std::string& file_compression, const std::shared_ptr& data_type) { @@ -261,15 +250,21 @@ Result<::orc::WriterOptions> OrcFormatWriter::PrepareWriterOptions( } } ::orc::WriterOptions writer_options; - PAIMON_ASSIGN_OR_RAISE(uint64_t stripe_size, - GetMemorySizeOption(options, ORC_STRIPE_SIZE, DEFAULT_STRIPE_SIZE)); + int64_t stripe_size = DEFAULT_STRIPE_SIZE; + auto file_block_size = options.find(Options::FILE_BLOCK_SIZE); + if (file_block_size != options.end()) { + PAIMON_ASSIGN_OR_RAISE(stripe_size, MemorySize::ParseBytes(file_block_size->second)); + } else { + PAIMON_ASSIGN_OR_RAISE(stripe_size, OptionsUtils::GetValueFromMap( + options, ORC_STRIPE_SIZE, DEFAULT_STRIPE_SIZE)); + } writer_options.setStripeSize(stripe_size); PAIMON_ASSIGN_OR_RAISE(::orc::CompressionKind compression, ToOrcCompressionKind(StringUtils::ToLowerCase(file_compression))); writer_options.setCompression(compression); - PAIMON_ASSIGN_OR_RAISE( - uint64_t compression_block_size, - GetMemorySizeOption(options, ORC_COMPRESSION_BLOCK_SIZE, DEFAULT_COMPRESSION_BLOCK_SIZE)); + PAIMON_ASSIGN_OR_RAISE(uint64_t compression_block_size, OptionsUtils::GetValueFromMap( + options, ORC_COMPRESSION_BLOCK_SIZE, + DEFAULT_COMPRESSION_BLOCK_SIZE)); writer_options.setCompressionBlockSize(compression_block_size); PAIMON_ASSIGN_OR_RAISE( double dictionary_key_threshold, diff --git a/src/paimon/format/orc/orc_format_writer_test.cpp b/src/paimon/format/orc/orc_format_writer_test.cpp index fff6816c0..70719eb51 100644 --- a/src/paimon/format/orc/orc_format_writer_test.cpp +++ b/src/paimon/format/orc/orc_format_writer_test.cpp @@ -288,6 +288,13 @@ TEST_F(OrcFormatWriterTest, TestPrepareWriterOptions) { OrcFormatWriter::PrepareWriterOptions(options, "zstd", data_type)); ASSERT_FALSE(writer_options.getEnableDictionary()); } + { + std::map options = {{ORC_STRIPE_SIZE, "4096"}, + {Options::FILE_BLOCK_SIZE, "8 KB"}}; + ASSERT_OK_AND_ASSIGN(::orc::WriterOptions writer_options, + OrcFormatWriter::PrepareWriterOptions(options, "zstd", data_type)); + ASSERT_EQ(writer_options.getStripeSize(), 8 * 1024); + } { // test disable config for timestamp with timezone arrow::FieldVector invalid_fields = { diff --git a/src/paimon/format/parquet/CMakeLists.txt b/src/paimon/format/parquet/CMakeLists.txt index a1a566c08..968546581 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 48a4430ad..fb5b65741 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 78a07c794..02e6ae5b1 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 aaef711e0..41dcde193 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 1b4bdd30d..f20f224fe 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 683bde71e..a143ae5a2 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 67b24ac50..daacaceee 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 56d36d36e..adb271734 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 7f7c114ed..d2053f120 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 c0cd40e19..7605c4242 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( @@ -134,6 +154,7 @@ ParquetFileBatchReader::ParquetFileBatchReader( arrow_pool_(arrow_pool), input_stream_(std::move(input_stream)), reader_(std::move(reader)), + read_ranges_(reader_->GetAllRowGroupRanges()), metrics_(std::make_shared()), storage_read_bytes_(std::move(storage_read_bytes)), logger_(Logger::GetLogger("ParquetFileBatchReader")) {} @@ -143,14 +164,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( @@ -289,6 +310,7 @@ Status ParquetFileBatchReader::SetReadSchema( PAIMON_RETURN_NOT_OK(UpdateAllTargetRowRanges(target_row_groups)); PAIMON_RETURN_NOT_OK(reader_->PrepareForReadingLazy(target_row_groups, column_indices)); + PAIMON_RETURN_NOT_OK(reader_->ApplyReadRanges(read_ranges_)); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetFileBatchReader::SetReadSchema") return Status::OK(); @@ -635,14 +657,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 +677,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 +686,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 +763,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 +794,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 7e5e9afab..2b1097fb4 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, @@ -125,12 +128,13 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { return reader_->GetNumberOfRows(); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { assert(reader_); return reader_->GetNextRowToRead(); } Status SetReadRanges(const std::vector>& read_ranges) override { + read_ranges_ = read_ranges; return reader_->ApplyReadRanges(read_ranges); } @@ -162,7 +166,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) { @@ -257,6 +262,8 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr read_data_type_; + std::vector> read_ranges_; + std::shared_ptr metrics_; // storageReadBytes counter shared with the underlying ArrowInputStreamAdapter. std::shared_ptr> storage_read_bytes_; 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 bcd4af210..1409f24bf 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_format_defs.h b/src/paimon/format/parquet/parquet_format_defs.h index 433103e54..8b205a09a 100644 --- a/src/paimon/format/parquet/parquet_format_defs.h +++ b/src/paimon/format/parquet/parquet_format_defs.h @@ -99,7 +99,7 @@ static inline const char PARQUET_READ_PREDICATE_NODE_COUNT_LIMIT[] = static inline const char PARQUET_READ_ENABLE_PAGE_INDEX_FILTER[] = "parquet.read.enable-page-index-filter"; -// Default is true. Compaction will set to false to reduce memory consumption. +// Default is true. static inline const char PARQUET_READ_ENABLE_PRE_BUFFER[] = "parquet.read.enable-pre-buffer"; static constexpr uint32_t DEFAULT_PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT = 0; diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index 519763724..112167bf3 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 4b8f97f0f..8dfe7faac 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 4dc7fbe58..65998426d 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 000000000..ad60caef3 --- /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/parquet_writer_builder.cpp b/src/paimon/format/parquet/parquet_writer_builder.cpp index 5946dd704..19d3e7fd4 100644 --- a/src/paimon/format/parquet/parquet_writer_builder.cpp +++ b/src/paimon/format/parquet/parquet_writer_builder.cpp @@ -24,10 +24,12 @@ #include "arrow/util/compression.h" #include "arrow/util/type_fwd.h" #include "fmt/format.h" +#include "paimon/common/options/memory_size.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/options_utils.h" #include "paimon/core/core_options.h" +#include "paimon/defs.h" #include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/format/parquet/parquet_format_writer.h" #include "paimon/status.h" @@ -77,9 +79,14 @@ Result> ParquetWriterBuilder::Prepa builder.compression_level(file_compression_level); } - PAIMON_ASSIGN_OR_RAISE(int64_t row_group_size, OptionsUtils::GetValueFromMap( - options_, PARQUET_BLOCK_SIZE, - ::parquet::DEFAULT_MAX_ROW_GROUP_SIZE)); + int64_t row_group_size = ::parquet::DEFAULT_MAX_ROW_GROUP_SIZE; + auto file_block_size = options_.find(Options::FILE_BLOCK_SIZE); + if (file_block_size != options_.end()) { + PAIMON_ASSIGN_OR_RAISE(row_group_size, MemorySize::ParseBytes(file_block_size->second)); + } else { + PAIMON_ASSIGN_OR_RAISE(row_group_size, OptionsUtils::GetValueFromMap( + options_, PARQUET_BLOCK_SIZE, row_group_size)); + } builder.max_row_group_size(row_group_size); PAIMON_ASSIGN_OR_RAISE(int64_t page_size, diff --git a/src/paimon/format/parquet/parquet_writer_builder_test.cpp b/src/paimon/format/parquet/parquet_writer_builder_test.cpp index 2c5de2cf3..3a95a8de3 100644 --- a/src/paimon/format/parquet/parquet_writer_builder_test.cpp +++ b/src/paimon/format/parquet/parquet_writer_builder_test.cpp @@ -77,6 +77,21 @@ TEST(ParquetWriterBuilderTest, PrepareWriterProperties) { ASSERT_EQ(3, properties->default_column_properties().compression_level()); } +TEST(ParquetWriterBuilderTest, PrepareWriterPropertiesWithFileBlockSize) { + arrow::FieldVector fields; + std::shared_ptr schema = arrow::schema(fields); + std::map options = { + {Options::FILE_FORMAT, "parquet"}, + {Options::MANIFEST_FORMAT, "parquet"}, + {PARQUET_BLOCK_SIZE, "2048"}, + {Options::FILE_BLOCK_SIZE, "8 KB"}, + }; + ParquetWriterBuilder builder(schema, /*batch_size=*/1024, options); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<::parquet::WriterProperties> properties, + builder.PrepareWriterProperties("zstd")); + ASSERT_EQ(properties->max_row_group_size(), 8 * 1024); +} + TEST(ParquetWriterBuilderTest, PrepareWriterPropertiesWithZstdLevelPriority) { arrow::FieldVector fields; std::shared_ptr schema = arrow::schema(fields); diff --git a/src/paimon/format/parquet/predicate_pushdown_test.cpp b/src/paimon/format/parquet/predicate_pushdown_test.cpp index 9723794eb..26175e1df 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 fd60e7b2f..404b0521d 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/fs/local/local_file.cpp b/src/paimon/fs/local/local_file.cpp index 645306676..a3f960af8 100644 --- a/src/paimon/fs/local/local_file.cpp +++ b/src/paimon/fs/local/local_file.cpp @@ -49,7 +49,7 @@ Result> LocalFile::Create(const std::string& path_str // local file system does not support path_string with scheme, e.g., "file:/tmp" will be // rewritten to "/tmp" PAIMON_ASSIGN_OR_RAISE(Path path, PathUtil::ToPath(path_string)); - if (!path.scheme.empty() && StringUtils::ToLowerCase(path.scheme) != "file") { + if (!path.scheme.empty() && !StringUtils::EqualsIgnoreCase(path.scheme, "file")) { return Status::Invalid(fmt::format("invalid scheme {} for local file system", path.scheme)); } if (path.path.empty() || path.path[0] != '/') { diff --git a/src/paimon/fs/local/local_file_test.cpp b/src/paimon/fs/local/local_file_test.cpp index f22095c90..25b9f6db9 100644 --- a/src/paimon/fs/local/local_file_test.cpp +++ b/src/paimon/fs/local/local_file_test.cpp @@ -27,6 +27,11 @@ namespace paimon::test { +TEST(LocalFileTest, TestSchemeCaseInsensitive) { + ASSERT_OK(LocalFile::Create("FILE:/tmp")); + ASSERT_NOK(LocalFile::Create("s3:/tmp")); +} + TEST(LocalFileTest, TestReadWriteEmptyContent) { auto test_root_dir = UniqueTestDirectory::Create(); ASSERT_TRUE(test_root_dir); diff --git a/src/paimon/fs/oss/CMakeLists.txt b/src/paimon/fs/oss/CMakeLists.txt new file mode 100644 index 000000000..fb0ea6744 --- /dev/null +++ b/src/paimon/fs/oss/CMakeLists.txt @@ -0,0 +1,50 @@ +# 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. + +if(PAIMON_ENABLE_OSS) + add_paimon_lib(paimon_oss_file_system + SOURCES + oss_file_system.cpp + oss_file_system_factory.cpp + EXTRA_INCLUDES + ${OSS_SDK_V2_INCLUDE_DIR} + DEPENDENCIES + paimon_shared + CURL::libcurl + STATIC_LINK_LIBS + alibabacloud_oss_v2::oss + fmt + SHARED_LINK_LIBS + paimon_shared + SHARED_LINK_FLAGS + ${PAIMON_VERSION_SCRIPT_FLAGS}) + + add_dependencies(paimon_oss_file_system_objlib oss_sdk_v2_ep) + + if(PAIMON_BUILD_TESTS) + add_paimon_test(oss_file_system_test + SOURCES + oss_file_system_test.cpp + EXTRA_INCLUDES + ${OSS_SDK_V2_INCLUDE_DIR} + STATIC_LINK_LIBS + paimon_shared + test_utils_static + ${PAIMON_OSS_FILE_SYSTEM_STATIC_LINK_LIBS} + ${GTEST_LINK_TOOLCHAIN}) + endif() +endif() diff --git a/src/paimon/fs/oss/oss_file_system.cpp b/src/paimon/fs/oss/oss_file_system.cpp new file mode 100644 index 000000000..bc283cc8b --- /dev/null +++ b/src/paimon/fs/oss/oss_file_system.cpp @@ -0,0 +1,192 @@ +/* + * 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/fs/oss/oss_file_system.h" + +#include +#include +#include +#include + +#include "alibabacloud/oss2/OSSClient.h" +#include "alibabacloud/oss2/Operation.h" +#include "alibabacloud/oss2/Types.h" +#include "alibabacloud/oss2/io/ByteWriter.h" +#include "alibabacloud/oss2/models/BucketBasic.h" +#include "alibabacloud/oss2/models/ObjectBasic.h" +#include "fmt/format.h" +#include "paimon/executor.h" + +namespace paimon::oss { +namespace { + +namespace oss2 = alibabacloud::oss2; + +constexpr int32_t kMaxKeysPerRequest = 1000; + +bool IsNotFoundError(const oss2::OperationError& error) { + return error.getStatusCode() == 404 || error.getCode() == "NoSuchKey" || + error.getCode() == "NoSuchBucket" || error.getCode() == "NotFound"; +} + +Status ToPaimonStatus(const oss2::OperationError& error, const std::string& operation, + const ObjectStorePath& path) { + std::string message = fmt::format("OSS {} 'oss://{}/{}' failed: code={}, status={}, message={}", + operation, path.bucket, path.key, error.getCode(), + error.getStatusCode(), error.getMessage()); + if (!error.getRequestId().empty()) { + message += fmt::format(", request_id={}", error.getRequestId()); + } + if (IsNotFoundError(error)) { + return Status::NotExist(message); + } + if (error.getCode() == "RequestCanceled") { + return Status::Cancelled(message); + } + return Status::IOError(message); +} + +class OssObjectStoreClient : public ObjectStoreClient, + public std::enable_shared_from_this { + public: + OssObjectStoreClient(std::string bucket, std::shared_ptr client, + std::unique_ptr executor) + : bucket_(std::move(bucket)), client_(std::move(client)), executor_(std::move(executor)) {} + + Result HeadObject(const ObjectStorePath& path) const override { + PAIMON_RETURN_NOT_OK(ValidateBucket(path)); + oss2::models::HeadObjectRequest request; + request.setBucket(path.bucket).setKey(path.key); + oss2::HeadObjectOutcome outcome = client_->headObject(request); + if (!outcome.has_value()) { + return ToPaimonStatus(outcome.error(), "HeadObject", path); + } + const oss2::models::HeadObjectResult& result = outcome.value(); + if (result.getContentLength() < 0) { + return Status::IOError("OSS HeadObject response is missing Content-Length"); + } + return ObjectMetadata{ + path.key, result.getContentLength(), + ObjectStoreFileSystemUtils::ParseModificationTime(result.getLastModified())}; + } + + Result ListObjects(const ObjectStorePath& path, + const std::string& continuation_token, + int32_t max_keys) const override { + PAIMON_RETURN_NOT_OK(ValidateBucket(path)); + oss2::models::ListObjectsV2Request request; + request.setBucket(path.bucket).setPrefix(path.key).setDelimiter("/"); + if (!continuation_token.empty()) { + request.setContinuationToken(continuation_token); + } + if (max_keys > 0) { + // OSS limits ListObjectsV2 requests to 1000 keys. + request.setMaxKeys(std::min(max_keys, kMaxKeysPerRequest)); + } + oss2::ListObjectsV2Outcome outcome = client_->listObjectsV2(request); + if (!outcome.has_value()) { + return ToPaimonStatus(outcome.error(), "ListObjectsV2", path); + } + const oss2::models::ListObjectsV2Result& value = outcome.value(); + ListObjectsResult result; + result.objects.reserve(value.getContents().size()); + for (const oss2::models::ObjectSummary& object : value.getContents()) { + result.objects.push_back(ObjectMetadata{ + object.key, object.size, + ObjectStoreFileSystemUtils::ParseModificationTime(object.lastModified)}); + } + result.common_prefixes.reserve(value.getCommonPrefixes().size()); + for (const oss2::models::CommonPrefix& prefix : value.getCommonPrefixes()) { + result.common_prefixes.push_back(prefix.prefix); + } + result.is_truncated = value.getIsTruncated(); + result.continuation_token = value.getNextContinuationToken(); + return result; + } + + Result GetObjectRange(const ObjectStorePath& path, int64_t offset, int64_t size, + char* buffer) const override { + PAIMON_RETURN_NOT_OK(ValidateBucket(path)); + if (size == 0) { + return 0; + } + auto writer = std::make_shared>(); + oss2::SinkFactory sink; + sink.isOneShot = false; + sink.supplier = [buffer, size, writer](int64_t, const oss2::HeaderCollection&) { + auto memory_writer = std::make_shared( + reinterpret_cast(buffer), static_cast(size)); + *writer = memory_writer; + return memory_writer; + }; + oss2::models::GetObjectRequest request; + request.setBucket(path.bucket) + .setKey(path.key) + .setRange(fmt::format("bytes={}-{}", offset, offset + size - 1)) + .setRangeBehavior("standard") + .setSinkFactory(std::move(sink)); + oss2::GetObjectOutcome outcome = client_->getObject(request); + if (!outcome.has_value()) { + return ToPaimonStatus(outcome.error(), "GetObject", path); + } + int64_t written = + *writer ? static_cast((*writer)->written()) : static_cast(0); + if (written != size) { + return Status::IOError( + fmt::format("OSS GetObject read {} bytes for oss://{}/{}, expected {}", written, + path.bucket, path.key, size)); + } + return written; + } + + void GetObjectRangeAsync(const ObjectStorePath& path, int64_t offset, int64_t size, + char* buffer, std::function&& callback) const override { + std::shared_ptr self = shared_from_this(); + executor_->Add([self = std::move(self), path, offset, size, buffer, + callback = std::move(callback)]() mutable { + Result result = self->GetObjectRange(path, offset, size, buffer); + callback(result.ok() ? Status::OK() : result.status()); + }); + } + + private: + Status ValidateBucket(const ObjectStorePath& path) const { + if (path.bucket != bucket_) { + return Status::Invalid( + fmt::format("OSS file system for bucket '{}' cannot access " + "'oss://{}/{}'", + bucket_, path.bucket, path.key)); + } + return Status::OK(); + } + + std::string bucket_; + std::shared_ptr client_; + std::unique_ptr executor_; +}; + +} // namespace + +OssFileSystem::OssFileSystem(std::string bucket, std::shared_ptr client, + std::unique_ptr executor) + : ObjectStoreFileSystem("oss", std::make_shared( + std::move(bucket), std::move(client), std::move(executor))) { +} + +} // namespace paimon::oss diff --git a/src/paimon/fs/oss/oss_file_system.h b/src/paimon/fs/oss/oss_file_system.h new file mode 100644 index 000000000..0749718e9 --- /dev/null +++ b/src/paimon/fs/oss/oss_file_system.h @@ -0,0 +1,51 @@ +/* + * 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/fs/object_store_file_system.h" +#include "paimon/executor.h" + +namespace alibabacloud::oss2 { +class OSSClient; +} + +namespace paimon::oss { + +inline constexpr char kOssAccessKeyIdOption[] = "fs.oss.accessKeyId"; +inline constexpr char kOssAccessKeySecretOption[] = "fs.oss.accessKeySecret"; +inline constexpr char kOssEndpointOption[] = "fs.oss.endpoint"; +inline constexpr char kOssRegionOption[] = "fs.oss.region"; +inline constexpr char kOssSignatureVersionOption[] = "fs.oss.signatureVersion"; +inline constexpr char kOssSecurityTokenOption[] = "fs.oss.securityToken"; +inline constexpr char kOssSessionTokenOption[] = "fs.oss.sessionToken"; +inline constexpr char kOssUsePathStyleOption[] = "fs.oss.usePathStyle"; +inline constexpr char kOssExecutorThreadCountOption[] = "fs.oss.executor.thread-count"; + +class OssFileSystem : public ObjectStoreFileSystem { + public: + OssFileSystem(std::string bucket, std::shared_ptr client, + std::unique_ptr executor); + ~OssFileSystem() override = default; +}; + +} // namespace paimon::oss diff --git a/src/paimon/fs/oss/oss_file_system_factory.cpp b/src/paimon/fs/oss/oss_file_system_factory.cpp new file mode 100644 index 000000000..327586d0e --- /dev/null +++ b/src/paimon/fs/oss/oss_file_system_factory.cpp @@ -0,0 +1,234 @@ +/* + * 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/fs/oss/oss_file_system_factory.h" + +#include +#include +#include +#include +#include +#include + +#include "alibabacloud/oss2/ClientConfiguration.h" +#include "alibabacloud/oss2/OSSClient.h" +#include "alibabacloud/oss2/credentials/CredentialsProvider.h" +#include "fmt/format.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/factories/factory.h" +#include "paimon/fs/oss/oss_file_system.h" + +namespace paimon::oss { +namespace { + +namespace oss2 = alibabacloud::oss2; + +constexpr std::string_view kOssOptionPrefix = "fs.oss."; +constexpr char kEndpointPrefix[] = "oss-"; +constexpr char kEndpointSuffix[] = ".aliyuncs.com"; +constexpr char kDualStackEndpointSuffix[] = ".oss.aliyuncs.com"; +constexpr char kInternalEndpointSuffix[] = "-internal"; + +bool IsValidRegion(const std::string& region) { + size_t first_separator = region.find('-'); + if (first_separator < 2 || first_separator > 3 || first_separator + 1 == region.size()) { + return false; + } + for (size_t i = 0; i < first_separator; ++i) { + if (!std::islower(static_cast(region[i]))) { + return false; + } + } + bool previous_was_separator = true; + for (size_t i = first_separator + 1; i < region.size(); ++i) { + char value = region[i]; + if (value == '-') { + if (previous_was_separator || i + 1 == region.size()) { + return false; + } + previous_was_separator = true; + } else if (std::islower(static_cast(value)) || + std::isdigit(static_cast(value))) { + previous_was_separator = false; + } else { + return false; + } + } + return !StringUtils::EndsWith(region, "-dualstack") && !StringUtils::EndsWith(region, "-pub"); +} + +std::string GetBucketOptionKey(const std::string& bucket, std::string_view option) { + return fmt::format("fs.oss.bucket.{}.{}", bucket, option.substr(kOssOptionPrefix.size())); +} + +const std::string* FindOption(const std::map& options, + const std::string& bucket, std::string_view option, + std::string* option_key) { + std::string bucket_option_key = GetBucketOptionKey(bucket, option); + auto bucket_option = options.find(bucket_option_key); + if (bucket_option != options.end()) { + *option_key = std::move(bucket_option_key); + return &bucket_option->second; + } + option_key->assign(option); + auto global_option = options.find(*option_key); + return global_option == options.end() ? nullptr : &global_option->second; +} + +std::string GetOption(const std::map& options, const std::string& bucket, + std::string_view option) { + std::string option_key; + const std::string* value = FindOption(options, bucket, option, &option_key); + return value == nullptr ? "" : *value; +} + +Result GetRequiredOption(const std::map& options, + const std::string& bucket, std::string_view option) { + std::string option_key; + const std::string* value = FindOption(options, bucket, option, &option_key); + if (value == nullptr || value->empty()) { + return Status::Invalid(fmt::format("OSS option '{}' must not be empty", option_key)); + } + return *value; +} + +Result> CreateExecutor(const std::map& options, + const std::string& bucket) { + std::string option_key; + const std::string* value = + FindOption(options, bucket, kOssExecutorThreadCountOption, &option_key); + if (value == nullptr) { + return CreateDefaultExecutor(); + } + std::optional thread_count = StringUtils::StringToValue(*value); + if (!thread_count.has_value() || *thread_count == 0) { + return Status::Invalid(fmt::format( + "OSS executor thread count for option '{}' must be greater than 0", option_key)); + } + return CreateDefaultExecutor(*thread_count); +} + +std::string NormalizeEndpoint(std::string endpoint) { + if (!endpoint.empty() && endpoint.find("://") == std::string::npos) { + endpoint = "https://" + endpoint; + } + return endpoint; +} + +std::string InferRegion(std::string endpoint) { + size_t scheme = endpoint.find("://"); + if (scheme != std::string::npos) { + endpoint.erase(0, scheme + 3); + } + size_t slash = endpoint.find('/'); + if (slash != std::string::npos) { + endpoint.erase(slash); + } + size_t port_separator = endpoint.rfind(':'); + if (port_separator != std::string::npos && endpoint.find(':') == port_separator) { + std::optional port = + StringUtils::StringToValue(endpoint.substr(port_separator + 1)); + if (port.has_value()) { + endpoint.erase(port_separator); + } + } + + std::string region; + if (StringUtils::StartsWith(endpoint, kEndpointPrefix) && + StringUtils::EndsWith(endpoint, kEndpointSuffix)) { + region = endpoint.substr( + sizeof(kEndpointPrefix) - 1, + endpoint.size() - (sizeof(kEndpointPrefix) - 1) - (sizeof(kEndpointSuffix) - 1)); + if (StringUtils::EndsWith(region, kInternalEndpointSuffix)) { + region.erase(region.size() - (sizeof(kInternalEndpointSuffix) - 1)); + } + } else if (StringUtils::EndsWith(endpoint, kDualStackEndpointSuffix)) { + region = endpoint.substr(0, endpoint.size() - (sizeof(kDualStackEndpointSuffix) - 1)); + } + return IsValidRegion(region) ? region : ""; +} + +} // namespace + +const char OssFileSystemFactory::IDENTIFIER[] = "oss"; + +Result> OssFileSystemFactory::Create( + const std::string& path, const std::map& options) const { + PAIMON_ASSIGN_OR_RAISE(Path parsed_path, PathUtil::ToPath(path)); + if (parsed_path.scheme != "oss" || parsed_path.authority.empty()) { + return Status::Invalid(fmt::format("invalid OSS path '{}'", path)); + } + const std::string& bucket = parsed_path.authority; + PAIMON_ASSIGN_OR_RAISE(std::string access_key_id, + GetRequiredOption(options, bucket, kOssAccessKeyIdOption)); + PAIMON_ASSIGN_OR_RAISE(std::string access_key_secret, + GetRequiredOption(options, bucket, kOssAccessKeySecretOption)); + std::string endpoint = GetOption(options, bucket, kOssEndpointOption); + std::string region = GetOption(options, bucket, kOssRegionOption); + if (region.empty()) { + region = InferRegion(endpoint); + } + if (endpoint.empty() && region.empty()) { + return Status::Invalid("OSS endpoint or region must be configured"); + } + std::string signature_version = GetOption(options, bucket, kOssSignatureVersionOption); + if (!signature_version.empty() && signature_version != "v1" && signature_version != "v4") { + return Status::Invalid( + fmt::format("invalid OSS signature version '{}'", signature_version)); + } + if (region.empty() && signature_version != "v1") { + return Status::Invalid( + "OSS region must be configured when the endpoint does not identify a region"); + } + std::string security_token = GetOption(options, bucket, kOssSecurityTokenOption); + if (security_token.empty()) { + security_token = GetOption(options, bucket, kOssSessionTokenOption); + } + + oss2::ClientConfiguration config = oss2::ClientConfiguration::loadDefault(); + if (!endpoint.empty()) { + config.endpoint = NormalizeEndpoint(endpoint); + } + if (!region.empty()) { + config.region = region; + } + if (!signature_version.empty()) { + config.signatureVersion = signature_version; + } + config.userAgent = "paimon-cpp"; + config.credentialsProvider = std::make_shared( + std::move(access_key_id), std::move(access_key_secret), std::move(security_token)); + std::string path_style = GetOption(options, bucket, kOssUsePathStyleOption); + if (!path_style.empty()) { + std::optional value = StringUtils::StringToValue(path_style); + if (!value.has_value()) { + return Status::Invalid(fmt::format("invalid boolean value '{}' for OSS option '{}'", + path_style, kOssUsePathStyleOption)); + } + config.usePathStyle = *value; + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr executor, CreateExecutor(options, bucket)); + return std::make_unique(bucket, std::make_shared(config), + std::move(executor)); +} + +REGISTER_PAIMON_FACTORY(OssFileSystemFactory); + +} // namespace paimon::oss diff --git a/src/paimon/fs/oss/oss_file_system_factory.h b/src/paimon/fs/oss/oss_file_system_factory.h new file mode 100644 index 000000000..fe92715e5 --- /dev/null +++ b/src/paimon/fs/oss/oss_file_system_factory.h @@ -0,0 +1,38 @@ +/* + * 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/fs/file_system_factory.h" + +namespace paimon::oss { + +class OssFileSystemFactory : public FileSystemFactory { + public: + static const char IDENTIFIER[]; + + const char* Identifier() const override { + return IDENTIFIER; + } + + Result> Create( + const std::string& path, const std::map& options) const override; +}; + +} // namespace paimon::oss diff --git a/src/paimon/fs/oss/oss_file_system_test.cpp b/src/paimon/fs/oss/oss_file_system_test.cpp new file mode 100644 index 000000000..fbd847f80 --- /dev/null +++ b/src/paimon/fs/oss/oss_file_system_test.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/fs/oss/oss_file_system.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "alibabacloud/oss2/ClientConfiguration.h" +#include "alibabacloud/oss2/OSSClient.h" +#include "alibabacloud/oss2/credentials/CredentialsProvider.h" +#include "alibabacloud/oss2/io/ByteWriter.h" +#include "alibabacloud/oss2/transport/HttpTransport.h" +#include "gtest/gtest.h" +#include "paimon/fs/oss/oss_file_system_factory.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::oss { +namespace { + +namespace oss2 = alibabacloud::oss2; + +class MockHttpTransport : public oss2::HttpTransport { + public: + oss2::ResponseResult send(std::unique_ptr& request, + const oss2::RequestOptions& options) override { + requests_.emplace_back(std::make_unique(*request)); + if (responses_.empty()) { + return oss2::TransportError{std::make_error_code(std::errc::no_message_available), "", + ""}; + } + std::unique_ptr response = std::move(responses_.front()); + responses_.erase(responses_.begin()); + if (response->statusCode / 100 == 2 && options.sinkFactory.has_value() && + response->body != nullptr) { + int64_t content_length = -1; + auto content_length_header = response->headers.find("Content-Length"); + if (content_length_header != response->headers.end()) { + content_length = std::stoll(content_length_header->second); + } + std::shared_ptr sink = + options.sinkFactory.value()(content_length, response->headers); + std::ostringstream body; + body << response->body->rdbuf(); + const std::string data = body.str(); + sink->write(reinterpret_cast(data.data()), data.size()); + response->body.reset(); + } + return response; + } + + std::string getName() const override { + return "MockHttpTransport"; + } + + void AddResponse(int status_code, oss2::HeaderCollection headers, std::string body = "") { + std::shared_ptr response_body; + if (!body.empty()) { + response_body = std::make_shared(std::move(body)); + } + responses_.emplace_back(std::make_unique(oss2::ResponseMessage{ + status_code, "", std::move(headers), std::move(response_body), nullptr})); + } + + std::vector> responses_; + std::vector> requests_; +}; + +std::unique_ptr CreateFileSystem( + const std::shared_ptr& transport) { + oss2::ClientConfiguration config = oss2::ClientConfiguration::loadDefault(); + config.region = "cn-hangzhou"; + config.credentialsProvider = + std::make_shared("access-key", "secret-key"); + config.httpTransport = transport; + return std::make_unique("bucket", std::make_shared(config), + CreateDefaultExecutor()); +} + +} // namespace + +TEST(OssFileSystemFactoryTest, TestOptionValidation) { + OssFileSystemFactory factory; + std::map options; + ASSERT_NOK(factory.Create("s3://bucket/key", options)); + ASSERT_NOK(factory.Create("oss://bucket/key", options)); + + options[kOssAccessKeyIdOption] = "access-key"; + options[kOssAccessKeySecretOption] = "secret-key"; + options[kOssEndpointOption] = "oss-cn-hangzhou.aliyuncs.com"; + options[kOssUsePathStyleOption] = "treu"; + ASSERT_NOK(factory.Create("oss://bucket/key", options)); + + options[kOssUsePathStyleOption] = "false"; + options[kOssSignatureVersionOption] = "v2"; + ASSERT_NOK(factory.Create("oss://bucket/key", options)); + + options[kOssSignatureVersionOption] = "v4"; + options[kOssExecutorThreadCountOption] = "0"; + ASSERT_NOK(factory.Create("oss://bucket/key", options)); + + options[kOssExecutorThreadCountOption] = "4"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options[kOssEndpointOption] = ""; + options[kOssRegionOption] = "cn-hangzhou"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); +} + +TEST(OssFileSystemFactoryTest, TestBucketOptionsOverrideGlobalOptions) { + OssFileSystemFactory factory; + std::map options = { + {kOssAccessKeyIdOption, ""}, + {kOssAccessKeySecretOption, ""}, + {kOssEndpointOption, ""}, + {"fs.oss.bucket.bucket.accessKeyId", "access-key"}, + {"fs.oss.bucket.bucket.accessKeySecret", "secret-key"}, + {"fs.oss.bucket.bucket.endpoint", "oss-cn-hangzhou.aliyuncs.com"}, + }; + ASSERT_OK(factory.Create("oss://bucket/key", options)); +} + +TEST(OssFileSystemFactoryTest, TestBucketOptionErrorReportsBucketKey) { + OssFileSystemFactory factory; + std::map options = { + {kOssAccessKeyIdOption, "access-key"}, + {kOssAccessKeySecretOption, "secret-key"}, + {kOssEndpointOption, "oss-cn-hangzhou.aliyuncs.com"}, + {"fs.oss.bucket.bucket.accessKeyId", ""}, + }; + ASSERT_NOK_WITH_MSG(factory.Create("oss://bucket/key", options), + "fs.oss.bucket.bucket.accessKeyId"); +} + +TEST(OssFileSystemFactoryTest, TestEndpointRegionValidation) { + OssFileSystemFactory factory; + std::map options = { + {kOssAccessKeyIdOption, "access-key"}, + {kOssAccessKeySecretOption, "secret-key"}, + }; + + options[kOssEndpointOption] = "oss-cn-hangzhou.aliyuncs.com"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options[kOssEndpointOption] = "oss-cn-hangzhou-internal.aliyuncs.com"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options[kOssEndpointOption] = "oss-ap-southeast-1.aliyuncs.com:443"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options[kOssEndpointOption] = "cn-hangzhou.oss.aliyuncs.com"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options[kOssEndpointOption] = "oss-accelerate.aliyuncs.com"; + ASSERT_NOK_WITH_MSG(factory.Create("oss://bucket/key", options), "OSS region must be"); + + options[kOssEndpointOption] = "oss-accelerate-overseas.aliyuncs.com"; + ASSERT_NOK_WITH_MSG(factory.Create("oss://bucket/key", options), "OSS region must be"); + + options[kOssEndpointOption] = "oss.example.com"; + ASSERT_NOK_WITH_MSG(factory.Create("oss://bucket/key", options), "OSS region must be"); + + options[kOssRegionOption] = "cn-hangzhou"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); + + options.erase(kOssRegionOption); + options[kOssSignatureVersionOption] = "v1"; + ASSERT_OK(factory.Create("oss://bucket/key", options)); +} + +TEST(OssFileSystemTest, TestHeadObjectParsesMetadata) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(200, {{"Content-Length", "3"}, + {"Last-Modified", "not-a-timestamp"}, + {"x-oss-request-id", "request-id"}}); + std::unique_ptr file_system = CreateFileSystem(transport); + + ASSERT_OK_AND_ASSIGN(FileStatus status, file_system->GetFileStatus("oss://bucket/key")); + ASSERT_EQ(3, status.GetLen()); + ASSERT_EQ(FileStatus::kUnknownModificationTime, status.GetModificationTime()); + ASSERT_EQ(1U, transport->requests_.size()); + ASSERT_EQ("HEAD", transport->requests_[0]->method); +} + +TEST(OssFileSystemTest, TestHeadObjectNotFound) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(404, {{"x-oss-request-id", "request-id"}}, + "NoSuchKeymissing"); + transport->AddResponse(200, {}, + "false"); + std::unique_ptr file_system = CreateFileSystem(transport); + + Result status = file_system->GetFileStatus("oss://bucket/missing"); + ASSERT_TRUE(status.status().IsNotExist()) << status.status().ToString(); + ASSERT_NOK_WITH_MSG(status, "does not exist"); +} + +TEST(OssFileSystemTest, TestHeadObjectErrorMapping) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(403, {{"x-oss-request-id", "request-id"}}, + "AccessDenieddenied"); + std::unique_ptr file_system = CreateFileSystem(transport); + + Result status = file_system->GetFileStatus("oss://bucket/key"); + ASSERT_TRUE(status.status().IsIOError()) << status.status().ToString(); + ASSERT_NOK_WITH_MSG(status, "code=AccessDenied, status=403"); +} + +TEST(OssFileSystemTest, TestListObjects) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(404, {{"x-oss-request-id", "request-id"}}, + "NoSuchKeymissing"); + transport->AddResponse(200, {{"x-oss-request-id", "request-id"}}, R"( + + false + + prefix/file + 3 + 2024-01-01T00:00:00.000Z + + prefix/sub/ +)"); + std::unique_ptr file_system = CreateFileSystem(transport); + std::vector statuses; + + ASSERT_OK(file_system->ListFileStatus("oss://bucket/prefix", &statuses)); + ASSERT_EQ(2U, statuses.size()); + ASSERT_EQ(1704067200000, statuses[0].GetModificationTime()); + ASSERT_EQ(2U, transport->requests_.size()); + ASSERT_EQ("HEAD", transport->requests_[0]->method); + ASSERT_EQ("GET", transport->requests_[1]->method); + ASSERT_NE(std::string::npos, transport->requests_[1]->uri.find("list-type=2")); +} + +TEST(OssFileSystemTest, TestListObjectsErrorMapping) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(403, {{"x-oss-request-id", "request-id"}}, + "AccessDenieddenied"); + std::unique_ptr file_system = CreateFileSystem(transport); + std::vector statuses; + + Status status = file_system->ListDir("oss://bucket/prefix/", &statuses); + ASSERT_TRUE(status.IsIOError()) << status.ToString(); + ASSERT_NOK_WITH_MSG(status, "code=AccessDenied, status=403"); +} + +TEST(OssFileSystemTest, TestGetObjectRangeAndShortRead) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(206, {{"Content-Length", "3"}}, "abc"); + std::unique_ptr file_system = CreateFileSystem(transport); + FileStatus file_status("oss://bucket/key", 3); + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, file_system->Open(file_status)); + char buffer[3]; + + ASSERT_OK_AND_ASSIGN(int64_t bytes_read, stream->Read(buffer, 3, 0)); + ASSERT_EQ(3, bytes_read); + ASSERT_EQ("abc", std::string(buffer, sizeof(buffer))); + ASSERT_EQ(1U, transport->requests_.size()); + ASSERT_EQ("bytes=0-2", transport->requests_[0]->headers.at("range")); + + std::shared_ptr short_transport = std::make_shared(); + short_transport->AddResponse(206, {{"Content-Length", "2"}}, "ab"); + std::unique_ptr short_file_system = CreateFileSystem(short_transport); + ASSERT_OK_AND_ASSIGN(std::unique_ptr short_stream, + short_file_system->Open(file_status)); + ASSERT_NOK_WITH_MSG(short_stream->Read(buffer, 3, 0), "expected 3"); +} + +TEST(OssFileSystemTest, TestGetObjectRangeErrorMapping) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(403, {{"x-oss-request-id", "request-id"}}, + "AccessDenieddenied"); + std::unique_ptr file_system = CreateFileSystem(transport); + FileStatus file_status("oss://bucket/key", 3); + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, file_system->Open(file_status)); + char buffer[3]; + + Result result = stream->Read(buffer, 3, 0); + ASSERT_TRUE(result.status().IsIOError()) << result.status().ToString(); + ASSERT_NOK_WITH_MSG(result, "code=AccessDenied, status=403"); +} + +TEST(OssFileSystemTest, TestGetObjectRangeAsync) { + std::shared_ptr transport = std::make_shared(); + transport->AddResponse(206, {{"Content-Length", "3"}}, "abc"); + std::unique_ptr file_system = CreateFileSystem(transport); + FileStatus file_status("oss://bucket/key", 3); + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, file_system->Open(file_status)); + char buffer[3]; + std::promise promise; + std::future future = promise.get_future(); + + stream->ReadAsync(buffer, 3, 0, + [&promise](Status status) { promise.set_value(std::move(status)); }); + + ASSERT_EQ(std::future_status::ready, future.wait_for(std::chrono::seconds(5))); + ASSERT_OK(future.get()); + ASSERT_EQ("abc", std::string(buffer, sizeof(buffer))); + ASSERT_EQ(1U, transport->requests_.size()); + ASSERT_EQ("bytes=0-2", transport->requests_[0]->headers.at("range")); +} + +} // namespace paimon::oss diff --git a/src/paimon/fs/s3/s3_file_system.cpp b/src/paimon/fs/s3/s3_file_system.cpp index 49668b70f..b9b98f1e5 100644 --- a/src/paimon/fs/s3/s3_file_system.cpp +++ b/src/paimon/fs/s3/s3_file_system.cpp @@ -71,11 +71,6 @@ Result ParseNonNegativeInt64(const std::string& value, const std::strin return *result; } -int64_t ParseModificationTime(const std::string& value) { - time_t seconds = curl_getdate(value.c_str(), nullptr); - return seconds == static_cast(-1) ? 0 : static_cast(seconds) * 1000; -} - std::string XmlUnescape(const std::string& value) { const std::pair entities[] = { {"&", "&"}, {"<", "<"}, {">", ">"}, {""", "\""}, {"'", "'"}}; @@ -580,22 +575,22 @@ bool IsIpAddressAuthority(const std::string& authority) { } const char* AwsDnsSuffixForRegion(const std::string& region) { - if (region.rfind("cn-", 0) == 0) { + if (StringUtils::StartsWith(region, "cn-")) { return "amazonaws.com.cn"; } - if (region.rfind("eusc-de-", 0) == 0) { + if (StringUtils::StartsWith(region, "eusc-de-")) { return "amazonaws.eu"; } - if (region.rfind("us-iso-", 0) == 0) { + if (StringUtils::StartsWith(region, "us-iso-")) { return "c2s.ic.gov"; } - if (region.rfind("us-isob-", 0) == 0) { + if (StringUtils::StartsWith(region, "us-isob-")) { return "sc2s.sgov.gov"; } - if (region.rfind("eu-isoe-", 0) == 0) { + if (StringUtils::StartsWith(region, "eu-isoe-")) { return "cloud.adc-e.uk"; } - if (region.rfind("us-isof-", 0) == 0) { + if (StringUtils::StartsWith(region, "us-isof-")) { return "csp.hci.ic.gov"; } return "amazonaws.com"; @@ -662,10 +657,10 @@ class S3ObjectStoreClient : public ObjectStoreClient, if (length == response.headers.end()) { return Status::IOError("HeadObject response is missing Content-Length"); } - int64_t modification_time = 0; + int64_t modification_time = FileStatus::kUnknownModificationTime; auto modified = response.headers.find("last-modified"); if (modified != response.headers.end()) { - modification_time = ParseModificationTime(modified->second); + modification_time = ObjectStoreFileSystemUtils::ParseModificationTime(modified->second); } PAIMON_ASSIGN_OR_RAISE(int64_t object_size, ParseNonNegativeInt64(length->second, "Content-Length")); @@ -712,10 +707,10 @@ class S3ObjectStoreClient : public ObjectStoreClient, PAIMON_ASSIGN_OR_RAISE(std::string decoded_key, PercentDecode(*key, "Key")); PAIMON_ASSIGN_OR_RAISE(int64_t object_size, ParseNonNegativeInt64(*size, "ListObjectsV2 Size")); - int64_t modified = 0; + int64_t modified = FileStatus::kUnknownModificationTime; auto last_modified = TagValue(block, "LastModified"); if (last_modified) { - modified = ParseModificationTime(*last_modified); + modified = ObjectStoreFileSystemUtils::ParseModificationTime(*last_modified); } result.objects.push_back(ObjectMetadata{decoded_key, object_size, modified}); } diff --git a/src/paimon/fs/s3/s3_file_system_test.cpp b/src/paimon/fs/s3/s3_file_system_test.cpp index 09b1bb4ba..822e1aede 100644 --- a/src/paimon/fs/s3/s3_file_system_test.cpp +++ b/src/paimon/fs/s3/s3_file_system_test.cpp @@ -21,11 +21,17 @@ #include +#include +#include #include #include #include #include +#include +#include +#include #include +#include #include #include @@ -40,6 +46,9 @@ class MockHttpClient : public HttpClient { public: Result Execute(const HttpRequest& request, const HttpBodyConsumer& consumer) const override { + if (before_execute_) { + before_execute_(); + } request_ = request; HttpResponse response; response.status_code = status_code_; @@ -55,6 +64,7 @@ class MockHttpClient : public HttpClient { int32_t status_code_ = 200; HttpHeaders response_headers_; std::string body_; + std::function before_execute_; }; class ScopedEnvironmentVariable { @@ -292,13 +302,13 @@ TEST(S3ObjectStoreClientTest, TestInvalidModificationTime) { ASSERT_OK_AND_ASSIGN(std::shared_ptr client, MakeS3ObjectStoreClient(StaticOptions(), http)); ASSERT_OK_AND_ASSIGN(auto metadata, client->HeadObject({"bucket", "file"})); - ASSERT_EQ(metadata.modification_time, 0); + ASSERT_EQ(metadata.modification_time, FileStatus::kUnknownModificationTime); http->body_ = "falsefile" "invalid12"; ASSERT_OK_AND_ASSIGN(auto result, client->ListObjects({"bucket", ""}, "", 0)); - ASSERT_EQ(result.objects[0].modification_time, 0); + ASSERT_EQ(result.objects[0].modification_time, FileStatus::kUnknownModificationTime); } TEST(S3ObjectStoreClientTest, TestRegionFromEnvironment) { @@ -428,13 +438,55 @@ TEST(S3ObjectStoreClientTest, TestRangeAndListObjects) { ASSERT_TRUE(result.is_truncated); ASSERT_EQ(result.continuation_token, "next token"); ASSERT_EQ(result.objects[0].key, "dir/a&b"); + ASSERT_EQ(result.objects[0].modification_time, 1767225600000); ASSERT_EQ(result.objects[1].key, "dir/a<b"); + ASSERT_EQ(result.objects[1].modification_time, FileStatus::kUnknownModificationTime); ASSERT_EQ(result.common_prefixes[0], "dir/sub/"); ASSERT_NE(http->request_.url.find("amazonaws.com/?list-type=2"), std::string::npos); ASSERT_NE(http->request_.url.find("encoding-type=url"), std::string::npos); ASSERT_NE(http->request_.url.find("continuation-token=old%20token"), std::string::npos); } +TEST(S3ObjectStoreClientTest, TestGetObjectRangeAsyncClientLifetime) { + auto http = std::make_shared(); + http->body_ = "data"; + std::mutex mutex; + std::condition_variable condition; + bool request_started = false; + bool release_request = false; + http->before_execute_ = [&] { + std::unique_lock lock(mutex); + request_started = true; + condition.notify_one(); + condition.wait(lock, [&] { return release_request; }); + }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr client, + MakeS3ObjectStoreClient(StaticOptions(), http)); + char buffer[4]; + auto promise = std::make_shared>(); + std::future future = promise->get_future(); + + client->GetObjectRangeAsync({"bucket", "key"}, 0, 4, buffer, [promise](Status status) { + promise->set_value(std::move(status)); + }); + { + std::unique_lock lock(mutex); + ASSERT_TRUE( + condition.wait_for(lock, std::chrono::seconds(5), [&] { return request_started; })); + } + std::thread destruction_thread([client = std::move(client)]() mutable { client.reset(); }); + { + std::lock_guard lock(mutex); + release_request = true; + } + condition.notify_one(); + + destruction_thread.join(); + ASSERT_EQ(std::future_status::ready, future.wait_for(std::chrono::seconds(5))); + ASSERT_OK(future.get()); + ASSERT_EQ("data", std::string(buffer, sizeof(buffer))); +} + TEST(S3ObjectStoreClientTest, TestUrlEncodedListObjects) { auto http = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr client, diff --git a/src/paimon/global_index/lucene/jieba_analyzer.cpp b/src/paimon/global_index/lucene/jieba_analyzer.cpp index 39cecec2f..e0a71a81a 100644 --- a/src/paimon/global_index/lucene/jieba_analyzer.cpp +++ b/src/paimon/global_index/lucene/jieba_analyzer.cpp @@ -17,6 +17,8 @@ */ #include "paimon/global_index/lucene/jieba_analyzer.h" +#include + #include "paimon/common/utils/string_utils.h" #include "paimon/global_index/lucene/lucene_utils.h" @@ -94,9 +96,7 @@ void JiebaTokenizer::NormalizeCase(std::string* term) { } } if (is_alphanumeric && !term->empty()) { - std::transform(term->begin(), term->end(), term->begin(), [](char ch) { - return static_cast(std::tolower(static_cast(ch))); - }); + *term = StringUtils::ToLowerCase(*term); } } diff --git a/src/paimon/rest/dlf_auth.cpp b/src/paimon/rest/dlf_auth.cpp new file mode 100644 index 000000000..6a4906279 --- /dev/null +++ b/src/paimon/rest/dlf_auth.cpp @@ -0,0 +1,792 @@ +/* + * 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/rest/dlf_auth.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/catalog_options.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/common/utils/url_utils.h" +#include "paimon/common/utils/uuid.h" +#include "paimon/rest/rest_http_client.h" +#include "rapidjson/document.h" + +namespace paimon { + +namespace { + +constexpr int32_t kEcsMetadataRequestTimeoutMillis = 3 * 60 * 1000; + +constexpr int64_t kTokenExpirationSafeTimeMillis = 60 * 60 * 1000; +constexpr size_t kMaxTokenResponseBytes = 1024 * 1024; +constexpr const char kDefaultEcsMetadataUrl[] = + "http://100.100.100.200/latest/meta-data/Ram/security-credentials/"; + +constexpr const char kAuthorizationHeader[] = "Authorization"; +constexpr const char kContentMd5Header[] = "Content-MD5"; +constexpr const char kContentTypeHeader[] = "Content-Type"; +constexpr const char kDlfDateHeader[] = "x-dlf-date"; +constexpr const char kDlfSecurityTokenHeader[] = "x-dlf-security-token"; +constexpr const char kDlfVersionHeader[] = "x-dlf-version"; +constexpr const char kDlfContentSha256Header[] = "x-dlf-content-sha256"; +constexpr const char kUnsignedPayload[] = "UNSIGNED-PAYLOAD"; +constexpr const char kJsonMediaType[] = "application/json"; + +constexpr const char kOpenApiDateHeader[] = "Date"; +constexpr const char kOpenApiAcceptHeader[] = "Accept"; +constexpr const char kOpenApiHostHeader[] = "Host"; +constexpr const char kAcsSignatureMethodHeader[] = "x-acs-signature-method"; +constexpr const char kAcsSignatureNonceHeader[] = "x-acs-signature-nonce"; +constexpr const char kAcsSignatureVersionHeader[] = "x-acs-signature-version"; +constexpr const char kAcsVersionHeader[] = "x-acs-version"; +constexpr const char kAcsSecurityTokenHeader[] = "x-acs-security-token"; + +Result RequiredNonEmptyOption(const std::map& options, + const std::string& key) { + Result value = OptionsUtils::GetNonEmptyValueFromMap(options, key); + if (!value.ok()) { + return Status::Invalid(fmt::format("option '{}' must be configured for DLF auth", key)); + } + return value.value(); +} + +Result RequiredJsonString(const rapidjson::Value& object, const char* key) { + if (!object.HasMember(key) || !object[key].IsString() || object[key].GetStringLength() == 0) { + return Status::Invalid(fmt::format("DLF token field '{}' must be a non-empty string", key)); + } + return std::string(object[key].GetString(), object[key].GetStringLength()); +} + +Result> OptionalJsonString(const rapidjson::Value& object, + const char* key) { + if (!object.HasMember(key) || object[key].IsNull()) { + return std::optional(); + } + if (!object[key].IsString()) { + return Status::Invalid(fmt::format("DLF token field '{}' must be a string", key)); + } + return std::optional( + std::string(object[key].GetString(), object[key].GetStringLength())); +} + +Result ToUtc(std::chrono::system_clock::time_point time) { + std::time_t seconds = std::chrono::system_clock::to_time_t(time); + std::tm utc{}; + if (gmtime_r(&seconds, &utc) == nullptr) { + return Status::Invalid("failed to convert DLF signing time to UTC"); + } + return utc; +} + +Result FormatDlfTime(std::chrono::system_clock::time_point time) { + PAIMON_ASSIGN_OR_RAISE(std::tm utc, ToUtc(time)); + std::array buffer{}; + if (std::strftime(buffer.data(), buffer.size(), "%Y%m%dT%H%M%SZ", &utc) == 0) { + return Status::Invalid("failed to format DLF signing time"); + } + return std::string(buffer.data()); +} + +Result FormatRfc1123Time(std::chrono::system_clock::time_point time) { + PAIMON_ASSIGN_OR_RAISE(std::tm utc, ToUtc(time)); + static constexpr std::array kWeekdays = {"Sun", "Mon", "Tue", "Wed", + "Thu", "Fri", "Sat"}; + static constexpr std::array kMonths = { + "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"}; + if (utc.tm_wday < 0 || utc.tm_wday >= static_cast(kWeekdays.size()) || + utc.tm_mon < 0 || utc.tm_mon >= static_cast(kMonths.size())) { + return Status::Invalid("failed to format DLF OpenAPI signing time"); + } + return fmt::format("{}, {:02d} {} {:04d} {:02d}:{:02d}:{:02d} GMT", kWeekdays[utc.tm_wday], + utc.tm_mday, kMonths[utc.tm_mon], utc.tm_year + 1900, utc.tm_hour, + utc.tm_min, utc.tm_sec); +} + +Result ParseExpiration(const std::string& expiration) { + static const std::regex kExpirationPattern( + "^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z$"); + if (!std::regex_match(expiration, kExpirationPattern)) { + return Status::Invalid("invalid DLF token expiration"); + } + std::tm utc{}; + std::istringstream stream(expiration); + stream >> std::get_time(&utc, "%Y-%m-%dT%H:%M:%SZ"); + if (stream.fail() || stream.peek() != std::char_traits::eof()) { + return Status::Invalid("invalid DLF token expiration"); + } + int32_t year = utc.tm_year; + int32_t month = utc.tm_mon; + int32_t day = utc.tm_mday; + int32_t hour = utc.tm_hour; + int32_t minute = utc.tm_min; + int32_t second = utc.tm_sec; + std::time_t timestamp = timegm(&utc); + std::tm verified{}; + if (timestamp == static_cast(-1) || gmtime_r(×tamp, &verified) == nullptr) { + return Status::Invalid("invalid DLF token expiration"); + } + if (verified.tm_year != year || verified.tm_mon != month || verified.tm_mday != day || + verified.tm_hour != hour || verified.tm_min != minute || verified.tm_sec != second) { + return Status::Invalid("invalid DLF token expiration"); + } + if (timestamp > std::numeric_limits::max() / 1000) { + return Status::Invalid("DLF token expiration is out of range"); + } + return static_cast(timestamp) * 1000; +} + +using Bytes = std::vector; +using EvpMdContext = std::unique_ptr; +using EvpPkey = std::unique_ptr; + +Result Digest(const EVP_MD* digest, std::string_view data) { + EvpMdContext context(EVP_MD_CTX_new(), EVP_MD_CTX_free); + if (!context || EVP_DigestInit_ex(context.get(), digest, nullptr) != 1 || + EVP_DigestUpdate(context.get(), data.data(), data.size()) != 1) { + return Status::IOError("failed to calculate DLF request digest"); + } + Bytes output(EVP_MAX_MD_SIZE); + unsigned int output_size = 0; + if (EVP_DigestFinal_ex(context.get(), output.data(), &output_size) != 1) { + return Status::IOError("failed to calculate DLF request digest"); + } + output.resize(output_size); + return output; +} + +Result Hmac(const EVP_MD* digest, const Bytes& key, std::string_view data) { + if (key.size() > static_cast(INT_MAX)) { + return Status::Invalid("DLF signing key is too large"); + } + EvpPkey signing_key( + EVP_PKEY_new_mac_key(EVP_PKEY_HMAC, nullptr, key.data(), static_cast(key.size())), + EVP_PKEY_free); + EvpMdContext context(EVP_MD_CTX_new(), EVP_MD_CTX_free); + if (!signing_key || !context || + EVP_DigestSignInit(context.get(), nullptr, digest, nullptr, signing_key.get()) != 1 || + EVP_DigestSignUpdate(context.get(), data.data(), data.size()) != 1) { + return Status::IOError("failed to calculate DLF request signature"); + } + size_t output_size = 0; + if (EVP_DigestSignFinal(context.get(), nullptr, &output_size) != 1) { + return Status::IOError("failed to calculate DLF request signature"); + } + Bytes output(output_size); + if (EVP_DigestSignFinal(context.get(), output.data(), &output_size) != 1) { + return Status::IOError("failed to calculate DLF request signature"); + } + output.resize(output_size); + return output; +} + +Bytes ToBytes(const std::string& value) { + return Bytes(value.begin(), value.end()); +} + +std::string HexEncode(const Bytes& value) { + static constexpr char kHex[] = "0123456789abcdef"; + std::string encoded; + encoded.reserve(value.size() * 2); + for (uint8_t byte : value) { + encoded.push_back(kHex[byte >> 4]); + encoded.push_back(kHex[byte & 0x0f]); + } + return encoded; +} + +Result Base64Encode(const Bytes& value) { + if (value.size() > static_cast(INT_MAX)) { + return Status::Invalid("DLF digest is too large to encode"); + } + size_t capacity = 4 * ((value.size() + 2) / 3) + 1; + std::string encoded(capacity, '\0'); + int32_t size = EVP_EncodeBlock(reinterpret_cast(encoded.data()), value.data(), + static_cast(value.size())); + if (size < 0) { + return Status::IOError("failed to encode DLF request digest"); + } + encoded.resize(static_cast(size)); + return encoded; +} + +Result Md5Base64(const std::string& value) { + PAIMON_ASSIGN_OR_RAISE(Bytes digest, Digest(EVP_md5(), value)); + return Base64Encode(digest); +} + +std::string Trimmed(const std::string& value) { + std::string trimmed = value; + StringUtils::Trim(&trimmed); + return trimmed; +} + +std::string DefaultCanonicalRequest(const RestAuthParameter& parameter, + const DlfRequestSigner::Headers& headers) { + std::string canonical = parameter.method + "\n" + parameter.resource_path + "\n"; + bool first = true; + for (const auto& [key, value] : parameter.parameters) { + if (!first) { + canonical += "&"; + } + canonical += Trimmed(key); + if (!value.empty()) { + canonical += "=" + Trimmed(value); + } + first = false; + } + + static const std::set kSignedHeaders = { + "content-md5", "content-type", "x-dlf-content-sha256", + "x-dlf-date", "x-dlf-version", "x-dlf-security-token"}; + std::map sorted_headers; + for (const auto& [key, value] : headers) { + std::string lower_key = StringUtils::ToLowerCase(key); + if (kSignedHeaders.count(lower_key) > 0) { + sorted_headers[lower_key] = Trimmed(value); + } + } + for (const auto& [key, value] : sorted_headers) { + canonical += "\n" + key + ":" + value; + } + auto content_iter = headers.find(kDlfContentSha256Header); + std::string content_sha = + content_iter == headers.end() ? std::string(kUnsignedPayload) : content_iter->second; + return canonical + "\n" + content_sha; +} + +Result RequiredHeader(const DlfRequestSigner::Headers& headers, + const std::string& name) { + auto iter = headers.find(name); + if (iter == headers.end() || iter->second.empty()) { + return Status::Invalid(fmt::format("DLF signing header '{}' is missing", name)); + } + return iter->second; +} + +std::string OpenApiCanonicalizedHeaders(const DlfRequestSigner::Headers& headers) { + std::map sorted; + for (const auto& [key, value] : headers) { + std::string lower_key = StringUtils::ToLowerCase(key); + if (StringUtils::StartsWith(lower_key, "x-acs-")) { + sorted[lower_key] = Trimmed(value); + } + } + std::string canonical; + for (const auto& [key, value] : sorted) { + canonical += key + ":" + value + "\n"; + } + return canonical; +} + +std::string OpenApiCanonicalizedResource(const RestAuthParameter& parameter) { + std::string resource = UrlUtils::DecodeString(parameter.resource_path); + if (parameter.parameters.empty()) { + return resource; + } + resource += "?"; + bool first = true; + for (const auto& [key, value] : parameter.parameters) { + if (!first) { + resource += "&"; + } + resource += key; + std::string decoded = UrlUtils::DecodeString(value); + if (!decoded.empty()) { + resource += "=" + decoded; + } + first = false; + } + return resource; +} + +Result GenerateNonce(std::chrono::system_clock::time_point now) { + std::string uuid; + if (!UUID::Generate(&uuid)) { + return Status::IOError("failed to generate DLF OpenAPI signing nonce"); + } + int64_t millis = + std::chrono::duration_cast(now.time_since_epoch()).count(); + std::ostringstream thread_id; + thread_id << std::this_thread::get_id(); + return fmt::format("{}{}{}", uuid, millis, thread_id.str()); +} + +Result> CreateSigner(const std::string& algorithm, + const std::string& region) { + if (algorithm == DlfDefaultSigner::kIdentifier) { + return std::make_unique(region); + } + if (algorithm == DlfOpenApiSigner::kIdentifier) { + return std::make_unique(); + } + return Status::Invalid(fmt::format( + "unsupported DLF signing algorithm '{}', supported values are 'default' and 'openapi'", + algorithm)); +} + +} // namespace + +DlfToken::DlfToken(const std::string& access_key_id, const std::string& access_key_secret, + const std::optional& security_token, + const std::optional& expiration_at_millis) + : access_key_id_(access_key_id), + access_key_secret_(access_key_secret), + security_token_(security_token), + expiration_at_millis_(expiration_at_millis) {} + +Result DlfToken::FromJson(const std::string& json) { + rapidjson::Document document; + document.Parse(json.data(), json.size()); + if (document.HasParseError() || !document.IsObject()) { + return Status::Invalid("failed to parse DLF token JSON"); + } + PAIMON_ASSIGN_OR_RAISE(std::string access_key_id, RequiredJsonString(document, "AccessKeyId")); + PAIMON_ASSIGN_OR_RAISE(std::string access_key_secret, + RequiredJsonString(document, "AccessKeySecret")); + PAIMON_ASSIGN_OR_RAISE(std::optional security_token, + OptionalJsonString(document, "SecurityToken")); + PAIMON_ASSIGN_OR_RAISE(std::optional expiration, + OptionalJsonString(document, "Expiration")); + std::optional expiration_at_millis; + if (expiration) { + PAIMON_ASSIGN_OR_RAISE(int64_t parsed_expiration, ParseExpiration(expiration.value())); + expiration_at_millis = parsed_expiration; + } + return DlfToken(access_key_id, access_key_secret, security_token, expiration_at_millis); +} + +bool DlfToken::ShouldRefresh(std::chrono::system_clock::time_point now) const { + if (!expiration_at_millis_) { + return false; + } + int64_t now_millis = + std::chrono::duration_cast(now.time_since_epoch()).count(); + return expiration_at_millis_.value() - now_millis < kTokenExpirationSafeTimeMillis; +} + +DlfLocalFileTokenLoader::DlfLocalFileTokenLoader(const std::string& token_file_path, + int32_t max_attempts, + std::chrono::milliseconds retry_delay) + : token_file_path_(token_file_path), max_attempts_(max_attempts), retry_delay_(retry_delay) {} + +Result DlfLocalFileTokenLoader::LoadToken() { + if (token_file_path_.empty()) { + return Status::Invalid("DLF token file path is empty"); + } + if (max_attempts_ <= 0 || retry_delay_.count() < 0) { + return Status::Invalid("invalid DLF token file retry configuration"); + } + Status last_status = Status::Invalid("failed to load DLF token file"); + for (int32_t attempt = 1; attempt <= max_attempts_; ++attempt) { + std::ifstream file(token_file_path_, std::ios::binary); + if (!file.is_open()) { + last_status = Status::IOError( + fmt::format("failed to read DLF token file '{}'", token_file_path_)); + } else { + std::string contents(kMaxTokenResponseBytes + 1, '\0'); + file.read(contents.data(), static_cast(contents.size())); + std::streamsize size = file.gcount(); + if (file.bad()) { + last_status = Status::IOError( + fmt::format("failed to read DLF token file '{}'", token_file_path_)); + } else if (size > static_cast(kMaxTokenResponseBytes)) { + last_status = Status::Invalid("DLF token file is too large"); + } else { + contents.resize(static_cast(size)); + Result token = DlfToken::FromJson(contents); + if (token.ok()) { + return token; + } + last_status = Status::Invalid("failed to parse DLF token file"); + } + } + if (attempt < max_attempts_) { + std::this_thread::sleep_for(retry_delay_ * attempt); + } + } + return last_status; +} + +std::string DlfLocalFileTokenLoader::Description() const { + return token_file_path_; +} + +DlfEcsTokenLoader::DlfEcsTokenLoader(const std::string& metadata_url, + const std::optional& role_name, + std::unique_ptr http_client) + : metadata_url_(metadata_url), role_name_(role_name), http_client_(std::move(http_client)) {} + +std::unique_ptr DlfEcsTokenLoader::Create( + const std::string& metadata_url, const std::optional& role_name) { + return std::make_unique(metadata_url, role_name, + std::make_unique()); +} + +Result DlfEcsTokenLoader::Get(const std::string& url) const { + if (!http_client_) { + return Status::Invalid("DLF ECS metadata HTTP client is not configured"); + } + HttpRequest request; + request.url = url; + request.request_timeout_ms = kEcsMetadataRequestTimeoutMillis; + std::string body; + Result response = + http_client_->Execute(request, [&body](const char* data, int64_t size) { + if (size < 0 || body.size() + static_cast(size) > kMaxTokenResponseBytes) { + return Status::Invalid("DLF ECS metadata response is too large"); + } + body.append(data, static_cast(size)); + return Status::OK(); + }); + if (!response.ok()) { + return Status::IOError("failed to request DLF credentials from ECS metadata service: ", + response.status().message()); + } + HttpResponse http_response = std::move(response).value(); + if (http_response.status_code < 200 || http_response.status_code >= 300) { + return Status::IOError(fmt::format("DLF ECS metadata service returned HTTP status {}", + http_response.status_code)); + } + if (StringUtils::IsNullOrWhitespaceOnly(body)) { + return Status::Invalid("DLF ECS metadata service returned an empty response"); + } + return body; +} + +Result DlfEcsTokenLoader::LoadToken() { + if (metadata_url_.empty()) { + return Status::Invalid("DLF ECS metadata URL is empty"); + } + if (!role_name_) { + PAIMON_ASSIGN_OR_RAISE(std::string role, Get(metadata_url_)); + StringUtils::Trim(&role); + if (role.empty()) { + return Status::Invalid("DLF ECS metadata service returned an empty role name"); + } + role_name_ = role; + } + PAIMON_ASSIGN_OR_RAISE(std::string token_json, Get(metadata_url_ + role_name_.value())); + Result token = DlfToken::FromJson(token_json); + if (!token.ok()) { + return Status::Invalid("failed to parse DLF ECS token response"); + } + return token; +} + +std::string DlfEcsTokenLoader::Description() const { + return metadata_url_; +} + +DlfDefaultSigner::DlfDefaultSigner(const std::string& region) : region_(region) {} + +Result DlfDefaultSigner::SignHeaders( + const std::string& body, std::chrono::system_clock::time_point now, + const std::optional& security_token, const std::string& host) const { + PAIMON_ASSIGN_OR_RAISE(std::string date_time, FormatDlfTime(now)); + Headers headers = {{kDlfDateHeader, date_time}, + {kDlfContentSha256Header, kUnsignedPayload}, + {kDlfVersionHeader, "v1"}}; + if (!body.empty()) { + PAIMON_ASSIGN_OR_RAISE(std::string content_md5, Md5Base64(body)); + headers[kContentTypeHeader] = kJsonMediaType; + headers[kContentMd5Header] = content_md5; + } + if (security_token) { + headers[kDlfSecurityTokenHeader] = security_token.value(); + } + return headers; +} + +Result DlfDefaultSigner::Authorization(const RestAuthParameter& parameter, + const DlfToken& token, const std::string& host, + const Headers& sign_headers) const { + PAIMON_ASSIGN_OR_RAISE(std::string date_time, RequiredHeader(sign_headers, kDlfDateHeader)); + if (date_time.size() < 8) { + return Status::Invalid("DLF signing date is invalid"); + } + std::string date = date_time.substr(0, 8); + std::string scope = fmt::format("{}/{}/DlfNext/aliyun_v4_request", date, region_); + std::string canonical_request = DefaultCanonicalRequest(parameter, sign_headers); + PAIMON_ASSIGN_OR_RAISE(Bytes canonical_hash, Digest(EVP_sha256(), canonical_request)); + std::string string_to_sign = + fmt::format("DLF4-HMAC-SHA256\n{}\n{}\n{}", date_time, scope, HexEncode(canonical_hash)); + + PAIMON_ASSIGN_OR_RAISE( + Bytes date_key, + Hmac(EVP_sha256(), ToBytes("aliyun_v4" + token.GetAccessKeySecret()), date)); + PAIMON_ASSIGN_OR_RAISE(Bytes region_key, Hmac(EVP_sha256(), date_key, region_)); + PAIMON_ASSIGN_OR_RAISE(Bytes service_key, Hmac(EVP_sha256(), region_key, "DlfNext")); + PAIMON_ASSIGN_OR_RAISE(Bytes signing_key, Hmac(EVP_sha256(), service_key, "aliyun_v4_request")); + PAIMON_ASSIGN_OR_RAISE(Bytes signature, Hmac(EVP_sha256(), signing_key, string_to_sign)); + return fmt::format("DLF4-HMAC-SHA256 Credential={}/{},Signature={}", token.GetAccessKeyId(), + scope, HexEncode(signature)); +} + +Result DlfOpenApiSigner::SignHeaders( + const std::string& body, std::chrono::system_clock::time_point now, + const std::optional& security_token, const std::string& host) const { + if (host.empty()) { + return Status::Invalid("DLF OpenAPI signing host is empty"); + } + PAIMON_ASSIGN_OR_RAISE(std::string date, FormatRfc1123Time(now)); + PAIMON_ASSIGN_OR_RAISE(std::string nonce, GenerateNonce(now)); + Headers headers = {{kOpenApiDateHeader, date}, {kOpenApiAcceptHeader, kJsonMediaType}, + {kOpenApiHostHeader, host}, {kAcsSignatureMethodHeader, "HMAC-SHA1"}, + {kAcsSignatureNonceHeader, nonce}, {kAcsSignatureVersionHeader, "1.0"}, + {kAcsVersionHeader, "2026-01-18"}}; + if (!body.empty()) { + PAIMON_ASSIGN_OR_RAISE(std::string content_md5, Md5Base64(body)); + headers[kContentMd5Header] = content_md5; + headers[kContentTypeHeader] = kJsonMediaType; + } + if (security_token) { + headers[kAcsSecurityTokenHeader] = security_token.value(); + } + return headers; +} + +Result DlfOpenApiSigner::Authorization(const RestAuthParameter& parameter, + const DlfToken& token, const std::string& host, + const Headers& sign_headers) const { + PAIMON_ASSIGN_OR_RAISE(std::string accept, RequiredHeader(sign_headers, kOpenApiAcceptHeader)); + PAIMON_ASSIGN_OR_RAISE(std::string date, RequiredHeader(sign_headers, kOpenApiDateHeader)); + std::string content_md5; + auto md5_iter = sign_headers.find(kContentMd5Header); + if (md5_iter != sign_headers.end()) { + content_md5 = md5_iter->second; + } + std::string content_type; + auto type_iter = sign_headers.find(kContentTypeHeader); + if (type_iter != sign_headers.end()) { + content_type = type_iter->second; + } + std::string string_to_sign = + parameter.method + "\n" + accept + "\n" + content_md5 + "\n" + content_type + "\n" + date + + "\n" + OpenApiCanonicalizedHeaders(sign_headers) + OpenApiCanonicalizedResource(parameter); + PAIMON_ASSIGN_OR_RAISE(Bytes signature, + Hmac(EVP_sha1(), ToBytes(token.GetAccessKeySecret()), string_to_sign)); + PAIMON_ASSIGN_OR_RAISE(std::string encoded_signature, Base64Encode(signature)); + return fmt::format("acs {}:{}", token.GetAccessKeyId(), encoded_signature); +} + +DlfAuthProvider::DlfAuthProvider(std::unique_ptr token_loader, + const std::optional& token, const std::string& host, + std::unique_ptr signer, Clock clock) + : token_loader_(std::move(token_loader)), + token_(token), + host_(host), + signer_(std::move(signer)), + clock_(std::move(clock)) {} + +Result> DlfAuthProvider::Create( + const std::map& options) { + PAIMON_ASSIGN_OR_RAISE(std::string uri, RequiredNonEmptyOption(options, CatalogOptions::URI)); + std::string region; + PAIMON_ASSIGN_OR_RAISE( + std::optional configured_region, + OptionsUtils::GetOptionalValueFromMap(options, CatalogOptions::DLF_REGION)); + if (configured_region) { + if (configured_region->empty()) { + return Status::Invalid("option 'dlf.region' must not be empty"); + } + region = configured_region.value(); + } else { + PAIMON_ASSIGN_OR_RAISE(region, ParseRegionFromUri(uri)); + } + + std::string algorithm; + PAIMON_ASSIGN_OR_RAISE(std::optional configured_algorithm, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_SIGNING_ALGORITHM)); + if (configured_algorithm) { + algorithm = configured_algorithm.value(); + } else { + algorithm = ParseSigningAlgorithmFromUri(uri); + } + + PAIMON_ASSIGN_OR_RAISE(std::optional loader_name, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_TOKEN_LOADER)); + PAIMON_ASSIGN_OR_RAISE(std::optional token_path, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_TOKEN_PATH)); + if (loader_name) { + if (loader_name.value() == "ecs") { + PAIMON_ASSIGN_OR_RAISE(std::optional configured_metadata_url, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_TOKEN_ECS_METADATA_URL)); + std::string metadata_url = configured_metadata_url.value_or(kDefaultEcsMetadataUrl); + PAIMON_ASSIGN_OR_RAISE(std::optional role_name, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_TOKEN_ECS_ROLE_NAME)); + return FromTokenLoader(DlfEcsTokenLoader::Create(metadata_url, role_name), uri, region, + algorithm, std::chrono::system_clock::now); + } + if (loader_name.value() == "local_file") { + PAIMON_ASSIGN_OR_RAISE(std::string path, + RequiredNonEmptyOption(options, CatalogOptions::DLF_TOKEN_PATH)); + return FromTokenLoader( + std::make_unique(path, 5, std::chrono::seconds(1)), uri, + region, algorithm, std::chrono::system_clock::now); + } + return Status::NotImplemented( + fmt::format("unsupported DLF token loader '{}', supported values are 'ecs' and " + "'local_file'", + loader_name.value())); + } + if (token_path) { + if (token_path->empty()) { + return Status::Invalid("option 'dlf.token-path' must not be empty"); + } + return FromTokenLoader(std::make_unique(token_path.value(), 5, + std::chrono::seconds(1)), + uri, region, algorithm, std::chrono::system_clock::now); + } + + PAIMON_ASSIGN_OR_RAISE(std::optional access_key_id, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_ACCESS_KEY_ID)); + PAIMON_ASSIGN_OR_RAISE(std::optional access_key_secret, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_ACCESS_KEY_SECRET)); + if (access_key_id && !access_key_id->empty() && access_key_secret && + !access_key_secret->empty()) { + PAIMON_ASSIGN_OR_RAISE(std::optional security_token, + OptionsUtils::GetOptionalValueFromMap( + options, CatalogOptions::DLF_SECURITY_TOKEN)); + DlfToken token(access_key_id.value(), access_key_secret.value(), security_token, + std::nullopt); + return FromAccessKey(token, uri, region, algorithm, std::chrono::system_clock::now); + } + return Status::Invalid("DLF token path or access key must be configured for DLF auth"); +} + +Result> DlfAuthProvider::FromAccessKey( + const DlfToken& token, const std::string& uri, const std::string& region, + const std::string& signing_algorithm, Clock clock) { + if (token.GetAccessKeyId().empty() || token.GetAccessKeySecret().empty()) { + return Status::Invalid("DLF access key id and secret must not be empty"); + } + PAIMON_ASSIGN_OR_RAISE(std::string host, ExtractHost(uri)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr signer, + CreateSigner(signing_algorithm, region)); + return std::unique_ptr( + new DlfAuthProvider(nullptr, token, host, std::move(signer), std::move(clock))); +} + +Result> DlfAuthProvider::FromTokenLoader( + std::unique_ptr token_loader, const std::string& uri, const std::string& region, + const std::string& signing_algorithm, Clock clock) { + if (!token_loader) { + return Status::Invalid("DLF token loader must not be null"); + } + PAIMON_ASSIGN_OR_RAISE(std::string host, ExtractHost(uri)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr signer, + CreateSigner(signing_algorithm, region)); + return std::unique_ptr(new DlfAuthProvider( + std::move(token_loader), std::nullopt, host, std::move(signer), std::move(clock))); +} + +Result DlfAuthProvider::GetFreshToken(std::chrono::system_clock::time_point now) const { + std::scoped_lock lock(token_mutex_); + if (token_ && !token_->ShouldRefresh(now)) { + return token_.value(); + } + if (!token_loader_) { + return Status::Invalid("DLF credentials expired and no token loader is configured"); + } + PAIMON_ASSIGN_OR_RAISE(DlfToken loaded, token_loader_->LoadToken()); + if (loaded.GetAccessKeyId().empty() || loaded.GetAccessKeySecret().empty()) { + return Status::Invalid("DLF token loader returned empty access key credentials"); + } + token_ = loaded; + return loaded; +} + +Result> DlfAuthProvider::MergeAuthHeader( + const std::map& base_header, + const RestAuthParameter& parameter) const { + PAIMON_ASSIGN_OR_RAISE(DlfToken token, GetFreshToken(clock_())); + std::chrono::system_clock::time_point signing_time = clock_(); + PAIMON_ASSIGN_OR_RAISE( + DlfRequestSigner::Headers sign_headers, + signer_->SignHeaders(parameter.data, signing_time, token.GetSecurityToken(), host_)); + PAIMON_ASSIGN_OR_RAISE(std::string authorization, + signer_->Authorization(parameter, token, host_, sign_headers)); + std::map headers = base_header; + for (const auto& [key, value] : sign_headers) { + headers[key] = value; + } + headers[kAuthorizationHeader] = authorization; + return headers; +} + +Result DlfAuthProvider::ParseRegionFromUri(const std::string& uri) { + static const std::regex kRegionPattern("(?:pre-)?([a-z]+-[a-z]+(?:-[0-9]+)?)"); + std::smatch match; + if (std::regex_search(uri, match, kRegionPattern) && match.size() > 1 && + !match.str(1).empty()) { + return match.str(1); + } + return Status::Invalid( + "could not determine DLF region from option 'dlf.region' or REST catalog URI"); +} + +std::string DlfAuthProvider::ParseSigningAlgorithmFromUri(const std::string& uri) { + std::string lower_uri = StringUtils::ToLowerCase(uri); + return lower_uri.find("dlfnext") == std::string::npos ? DlfDefaultSigner::kIdentifier + : DlfOpenApiSigner::kIdentifier; +} + +Result DlfAuthProvider::ExtractHost(const std::string& uri) { + std::string host = RestHttpClient::NormalizeUri(uri); + std::string lower_uri = StringUtils::ToLowerCase(host); + if (StringUtils::StartsWith(lower_uri, "http://")) { + host.erase(0, 7); + } else if (StringUtils::StartsWith(lower_uri, "https://")) { + host.erase(0, 8); + } + size_t path = host.find('/'); + if (path != std::string::npos) { + host.resize(path); + } + if (host.empty() || host.find('?') != std::string::npos || + host.find('#') != std::string::npos || host.find('@') != std::string::npos) { + return Status::Invalid("could not determine DLF signing host from REST catalog URI"); + } + return host; +} + +} // namespace paimon diff --git a/src/paimon/rest/dlf_auth.h b/src/paimon/rest/dlf_auth.h new file mode 100644 index 000000000..c10b655b7 --- /dev/null +++ b/src/paimon/rest/dlf_auth.h @@ -0,0 +1,208 @@ +/* + * 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 +#include +#include +#include + +#include "paimon/common/utils/http_client.h" +#include "paimon/rest/rest_auth.h" + +namespace paimon { + +/// Access key credentials used to sign DLF REST requests. +class DlfToken { + public: + DlfToken(const std::string& access_key_id, const std::string& access_key_secret, + const std::optional& security_token, + const std::optional& expiration_at_millis); + + static Result FromJson(const std::string& json); + + const std::string& GetAccessKeyId() const { + return access_key_id_; + } + + const std::string& GetAccessKeySecret() const { + return access_key_secret_; + } + + const std::optional& GetSecurityToken() const { + return security_token_; + } + + const std::optional& GetExpirationAtMillis() const { + return expiration_at_millis_; + } + + bool ShouldRefresh(std::chrono::system_clock::time_point now) const; + + private: + std::string access_key_id_; + std::string access_key_secret_; + std::optional security_token_; + std::optional expiration_at_millis_; +}; + +/// Loads refreshable DLF credentials. +class DlfTokenLoader { + public: + virtual ~DlfTokenLoader() = default; + + virtual Result LoadToken() = 0; + virtual std::string Description() const = 0; +}; + +/// Loads a DLF STS token from a JSON file. +class DlfLocalFileTokenLoader : public DlfTokenLoader { + public: + DlfLocalFileTokenLoader(const std::string& token_file_path, int32_t max_attempts, + std::chrono::milliseconds retry_delay); + + Result LoadToken() override; + std::string Description() const override; + + private: + std::string token_file_path_; + int32_t max_attempts_; + std::chrono::milliseconds retry_delay_; +}; + +/// Loads a DLF STS token from the Alibaba Cloud ECS metadata service. +class DlfEcsTokenLoader : public DlfTokenLoader { + public: + DlfEcsTokenLoader(const std::string& metadata_url, const std::optional& role_name, + std::unique_ptr http_client); + + static std::unique_ptr Create(const std::string& metadata_url, + const std::optional& role_name); + + Result LoadToken() override; + std::string Description() const override; + + private: + Result Get(const std::string& url) const; + + std::string metadata_url_; + std::optional role_name_; + std::unique_ptr http_client_; +}; + +/// Signs a DLF REST request using one of the endpoint-specific algorithms. +class DlfRequestSigner { + public: + using Headers = std::map; + + virtual ~DlfRequestSigner() = default; + + virtual Result SignHeaders(const std::string& body, + std::chrono::system_clock::time_point now, + const std::optional& security_token, + const std::string& host) const = 0; + + virtual Result Authorization(const RestAuthParameter& parameter, + const DlfToken& token, const std::string& host, + const Headers& sign_headers) const = 0; +}; + +/// DLF4-HMAC-SHA256 signer used by the default DLF VPC endpoint. +class DlfDefaultSigner : public DlfRequestSigner { + public: + static constexpr const char* kIdentifier = "default"; + + explicit DlfDefaultSigner(const std::string& region); + + Result SignHeaders(const std::string& body, std::chrono::system_clock::time_point now, + const std::optional& security_token, + const std::string& host) const override; + + Result Authorization(const RestAuthParameter& parameter, const DlfToken& token, + const std::string& host, + const Headers& sign_headers) const override; + + private: + std::string region_; +}; + +/// ROA HMAC-SHA1 signer used by DlfNext OpenAPI endpoints. +class DlfOpenApiSigner : public DlfRequestSigner { + public: + static constexpr const char* kIdentifier = "openapi"; + + Result SignHeaders(const std::string& body, std::chrono::system_clock::time_point now, + const std::optional& security_token, + const std::string& host) const override; + + Result Authorization(const RestAuthParameter& parameter, const DlfToken& token, + const std::string& host, + const Headers& sign_headers) const override; +}; + +/// Generates DLF authentication headers and refreshes expiring credentials. +class DlfAuthProvider : public AuthProvider { + public: + using Clock = std::function; + + static Result> Create( + const std::map& options); + + static Result> FromAccessKey( + const DlfToken& token, const std::string& uri, const std::string& region, + const std::string& signing_algorithm, Clock clock); + + static Result> FromTokenLoader( + std::unique_ptr token_loader, const std::string& uri, + const std::string& region, const std::string& signing_algorithm, Clock clock); + + Result> MergeAuthHeader( + const std::map& base_header, + const RestAuthParameter& parameter) const override; + + bool AllowsRedirects() const override { + return false; + } + + static Result ParseRegionFromUri(const std::string& uri); + static std::string ParseSigningAlgorithmFromUri(const std::string& uri); + static Result ExtractHost(const std::string& uri); + + private: + DlfAuthProvider(std::unique_ptr token_loader, + const std::optional& token, const std::string& host, + std::unique_ptr signer, Clock clock); + + Result GetFreshToken(std::chrono::system_clock::time_point now) const; + + std::unique_ptr token_loader_; + mutable std::optional token_; + std::string host_; + std::unique_ptr signer_; + Clock clock_; + mutable std::mutex token_mutex_; +}; + +} // namespace paimon diff --git a/src/paimon/rest/dlf_auth_test.cpp b/src/paimon/rest/dlf_auth_test.cpp new file mode 100644 index 000000000..afb9cec50 --- /dev/null +++ b/src/paimon/rest/dlf_auth_test.cpp @@ -0,0 +1,468 @@ +/* + * 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/rest/dlf_auth.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/catalog_options.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +namespace { + +using StringMap = std::map; + +std::chrono::system_clock::time_point FixedTime() { + return std::chrono::system_clock::from_time_t(1744775086); +} + +class SequenceTokenLoader : public DlfTokenLoader { + public: + explicit SequenceTokenLoader( + std::vector tokens, + std::chrono::milliseconds load_delay = std::chrono::milliseconds(0)) + : tokens_(std::move(tokens)), load_delay_(load_delay) {} + + Result LoadToken() override { + std::this_thread::sleep_for(load_delay_); + int32_t index = load_count_.fetch_add(1); + if (index >= static_cast(tokens_.size())) { + return Status::Invalid("test token loader exhausted"); + } + return tokens_[index]; + } + + std::string Description() const override { + return "test sequence"; + } + + int32_t GetLoadCount() const { + return load_count_.load(); + } + + private: + std::vector tokens_; + std::chrono::milliseconds load_delay_; + std::atomic load_count_{0}; +}; + +class MockEcsHttpClient : public HttpClient { + public: + explicit MockEcsHttpClient(const std::string& metadata_url, bool direct_token = false) + : metadata_url_(metadata_url), direct_token_(direct_token) {} + + Result Execute(const HttpRequest& request, + const HttpBodyConsumer& consumer) const override { + last_request_timeout_ms_.store(request.request_timeout_ms); + HttpResponse response; + response.status_code = 200; + std::string body; + if (request.url == metadata_url_) { + if (direct_token_) { + token_requests_.fetch_add(1); + body = R"({"AccessKeyId":"ecs-ak","AccessKeySecret":"ecs-sk",)" + R"("SecurityToken":"ecs-sts","Expiration":"2027-04-16T05:44:46Z"})"; + } else { + role_requests_.fetch_add(1); + body = " test-role\n"; + } + } else if (request.url == metadata_url_ + "test-role") { + token_requests_.fetch_add(1); + body = R"({"AccessKeyId":"ecs-ak","AccessKeySecret":"ecs-sk",)" + R"("SecurityToken":"ecs-sts","Expiration":"2027-04-16T05:44:46Z"})"; + } else { + response.status_code = 404; + } + if (!body.empty()) { + PAIMON_RETURN_NOT_OK(consumer(body.data(), static_cast(body.size()))); + response.body_size = static_cast(body.size()); + } + return response; + } + + int32_t GetRoleRequestCount() const { + return role_requests_.load(); + } + + int32_t GetTokenRequestCount() const { + return token_requests_.load(); + } + + int64_t GetLastRequestTimeoutMillis() const { + return last_request_timeout_ms_.load(); + } + + private: + std::string metadata_url_; + bool direct_token_; + mutable std::atomic role_requests_{0}; + mutable std::atomic token_requests_{0}; + mutable std::atomic last_request_timeout_ms_{-1}; +}; + +class FailingEcsHttpClient : public HttpClient { + public: + Result Execute(const HttpRequest&, const HttpBodyConsumer&) const override { + return Status::IOError("connection refused"); + } +}; + +} // namespace + +TEST(DlfDefaultSignerTest, SignsJavaCompatibleRequest) { + DlfDefaultSigner signer("cn-beijing"); + const std::string body = R"({"name":"t1"})"; + DlfToken token("YourAccessKeyId", "YourAccessKeySecret", "securityToken", std::nullopt); + RestAuthParameter parameter = + RestAuthParameter::Create("POST", "/v1/wh/databases/db/tables", + {{"warehouse", "my instance"}, {"branch", "main"}}, body); + + ASSERT_OK_AND_ASSIGN(DlfRequestSigner::Headers headers, + signer.SignHeaders(body, FixedTime(), token.GetSecurityToken(), "unused")); + ASSERT_EQ("20250416T034446Z", headers.at("x-dlf-date")); + ASSERT_EQ("Od9T1x3c2+JusJPFMpXe9Q==", headers.at("Content-MD5")); + ASSERT_EQ("application/json", headers.at("Content-Type")); + ASSERT_EQ("UNSIGNED-PAYLOAD", headers.at("x-dlf-content-sha256")); + ASSERT_EQ("v1", headers.at("x-dlf-version")); + ASSERT_EQ("securityToken", headers.at("x-dlf-security-token")); + + ASSERT_OK_AND_ASSIGN(std::string authorization, + signer.Authorization(parameter, token, "unused", headers)); + ASSERT_EQ( + "DLF4-HMAC-SHA256 Credential=YourAccessKeyId/20250416/cn-beijing/" + "DlfNext/aliyun_v4_request,Signature=" + "22594f8bbb8bb0ec296ced6003b7ffdf7022a8ca3815da5b53090daa11a06558", + authorization); +} + +TEST(DlfDefaultSignerTest, MatchesJavaGoldenAuthorization) { + // Mirrors Java DLFAuthSignatureTest#testGetAuthorization. + DlfDefaultSigner signer("cn-hangzhou"); + const std::string body = R"({"name":"database","options":{"a":"b"}})"; + DlfToken token("access-key-id", "access-key-secret", "securityToken", std::nullopt); + RestAuthParameter parameter = RestAuthParameter::Create("POST", "/v1/paimon/databases", + {{"k1", "v1"}, {"k2", "v2"}}, body); + const std::chrono::system_clock::time_point signing_time = + std::chrono::system_clock::from_time_t(1701605532); + + ASSERT_OK_AND_ASSIGN(DlfRequestSigner::Headers headers, + signer.SignHeaders(body, signing_time, token.GetSecurityToken(), "host")); + ASSERT_OK_AND_ASSIGN(std::string authorization, + signer.Authorization(parameter, token, "host", headers)); + ASSERT_EQ( + "DLF4-HMAC-SHA256 Credential=access-key-id/20231203/cn-hangzhou/" + "DlfNext/aliyun_v4_request,Signature=" + "c72caf1d40b55b1905d891ee3e3de48a2f8bebefa7e39e4f277acc93c269c5e3", + authorization); +} + +TEST(DlfDefaultSignerTest, OmitsBodyAndSecurityTokenHeadersWhenAbsent) { + DlfDefaultSigner signer("cn-hangzhou"); + ASSERT_OK_AND_ASSIGN(DlfRequestSigner::Headers headers, + signer.SignHeaders("", FixedTime(), std::nullopt, "unused")); + ASSERT_EQ(3, headers.size()); + ASSERT_EQ(0, headers.count("Content-MD5")); + ASSERT_EQ(0, headers.count("Content-Type")); + ASSERT_EQ(0, headers.count("x-dlf-security-token")); +} + +TEST(DlfOpenApiSignerTest, SignsJavaCompatibleRequest) { + DlfOpenApiSigner signer; + const std::string body = R"({"CategoryName":"test","CategoryType":"UNSTRUCTURED"})"; + const std::string host = "dlfnext.cn-beijing.aliyuncs.com"; + DlfToken token("YourAccessKeyId", "YourAccessKeySecret", "securityToken", std::nullopt); + RestAuthParameter parameter = + RestAuthParameter::Create("POST", "/llm-p2e4XXXXXXXXsvtn/datacenter/category", {}, body); + + ASSERT_OK_AND_ASSIGN(DlfRequestSigner::Headers headers, + signer.SignHeaders(body, FixedTime(), token.GetSecurityToken(), host)); + headers["x-acs-signature-nonce"] = "ef34aae7-7bd2-413d-a541-680cd2c48538"; + ASSERT_EQ("Wed, 16 Apr 2025 03:44:46 GMT", headers.at("Date")); + ASSERT_EQ("q2qaEcR4P47+Z7CUzHRTBw==", headers.at("Content-MD5")); + ASSERT_EQ("application/json", headers.at("Accept")); + ASSERT_EQ("application/json", headers.at("Content-Type")); + ASSERT_EQ(host, headers.at("Host")); + ASSERT_EQ("HMAC-SHA1", headers.at("x-acs-signature-method")); + ASSERT_EQ("1.0", headers.at("x-acs-signature-version")); + ASSERT_EQ("2026-01-18", headers.at("x-acs-version")); + ASSERT_EQ("securityToken", headers.at("x-acs-security-token")); + + ASSERT_OK_AND_ASSIGN(std::string authorization, + signer.Authorization(parameter, token, host, headers)); + ASSERT_EQ("acs YourAccessKeyId:wX4CDPSCtfgYkxdK9tJIO3ez5VI=", authorization); +} + +TEST(DlfOpenApiSignerTest, DecodesPathAndQueryValuesBeforeSigning) { + DlfOpenApiSigner signer; + DlfToken token("ak", "sk", std::nullopt, std::nullopt); + RestAuthParameter parameter = RestAuthParameter::Create( + "GET", "/v1/%24snapshots", {{"z", ""}, {"name", "hello world"}}, ""); + ASSERT_OK_AND_ASSIGN(DlfRequestSigner::Headers headers, + signer.SignHeaders("", FixedTime(), std::nullopt, "host")); + headers["x-acs-signature-nonce"] = "fixed-nonce"; + + ASSERT_OK_AND_ASSIGN(std::string authorization, + signer.Authorization(parameter, token, "host", headers)); + ASSERT_EQ("acs ak:vD+M7291KKoOTvt2gQuD6jPDlw8=", authorization); + ASSERT_EQ(0, headers.count("Content-MD5")); + ASSERT_EQ(0, headers.count("Content-Type")); +} + +TEST(DlfTokenTest, ParsesExpirationAndRefreshBoundary) { + ASSERT_OK_AND_ASSIGN( + DlfToken token, + DlfToken::FromJson(R"({"AccessKeyId":"ak","AccessKeySecret":"sk","SecurityToken":"sts",)" + R"("Expiration":"2025-04-16T05:44:46Z","Ignored":"value"})")); + ASSERT_EQ("ak", token.GetAccessKeyId()); + ASSERT_EQ("sk", token.GetAccessKeySecret()); + ASSERT_EQ(std::optional("sts"), token.GetSecurityToken()); + ASSERT_EQ(std::optional(1744782286000), token.GetExpirationAtMillis()); + ASSERT_FALSE(token.ShouldRefresh(FixedTime() + std::chrono::hours(1))); + ASSERT_TRUE( + token.ShouldRefresh(FixedTime() + std::chrono::hours(1) + std::chrono::milliseconds(1))); + + DlfToken permanent("ak", "sk", std::nullopt, std::nullopt); + ASSERT_FALSE(permanent.ShouldRefresh(FixedTime() + std::chrono::hours(100000))); +} + +TEST(DlfTokenTest, ParseFailureDoesNotLeakCredentials) { + const std::string secret = "STSSECRET_AKID_9999"; + Status status = + DlfToken::FromJson(R"({"AccessKeyId":"ak","AccessKeySecret":")" + secret).status(); + ASSERT_FALSE(status.ok()); + ASSERT_EQ(std::string::npos, status.ToString().find(secret)); +} + +TEST(DlfLocalFileTokenLoaderTest, LoadsTokenAndRedactsMalformedContent) { + std::unique_ptr test_dir = UniqueTestDirectory::Create(); + ASSERT_NE(nullptr, test_dir); + const std::string path = test_dir->Str() + "/dlf-token.json"; + { + std::ofstream file(path); + file << R"({"AccessKeyId":"file-ak","AccessKeySecret":"file-sk",)" + R"("SecurityToken":"file-sts","Expiration":"2027-04-16T05:44:46Z"})"; + } + DlfLocalFileTokenLoader loader(path, 1, std::chrono::milliseconds(0)); + ASSERT_OK_AND_ASSIGN(DlfToken token, loader.LoadToken()); + ASSERT_EQ("file-ak", token.GetAccessKeyId()); + ASSERT_EQ(std::optional("file-sts"), token.GetSecurityToken()); + + std::map options = { + {CatalogOptions::URI, "https://cn-hangzhou-vpc.dlf.aliyuncs.com"}, + {CatalogOptions::TOKEN_PROVIDER, "dlf"}, + {CatalogOptions::DLF_TOKEN_PATH, path}, + {CatalogOptions::DLF_ACCESS_KEY_ID, "ignored-ak"}, + {CatalogOptions::DLF_ACCESS_KEY_SECRET, "ignored-sk"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr provider, AuthProvider::Create(options)); + RestAuthParameter parameter = RestAuthParameter::Create("GET", "/v1/config", {}, ""); + ASSERT_OK_AND_ASSIGN(StringMap headers, provider->MergeAuthHeader({}, parameter)); + ASSERT_NE(std::string::npos, headers.at("Authorization").find("Credential=file-ak/")); + + options[CatalogOptions::DLF_TOKEN_LOADER] = "local_file"; + ASSERT_OK_AND_ASSIGN(provider, AuthProvider::Create(options)); + ASSERT_OK(provider->MergeAuthHeader({}, parameter)); + + const std::string secret = "FILE_SECRET_9999"; + { + std::ofstream file(path); + file << R"({"AccessKeyId":"ak","AccessKeySecret":")" << secret; + } + Status status = loader.LoadToken().status(); + ASSERT_FALSE(status.ok()); + ASSERT_EQ(std::string::npos, status.ToString().find(secret)); +} + +TEST(DlfEcsTokenLoaderTest, DiscoversRoleOnceAndRefreshesToken) { + const std::string metadata_url = "http://100.100.100.200/metadata/"; + auto http_client = std::make_unique(metadata_url); + MockEcsHttpClient* http_client_ptr = http_client.get(); + DlfEcsTokenLoader loader(metadata_url, std::nullopt, std::move(http_client)); + + ASSERT_OK_AND_ASSIGN(DlfToken first, loader.LoadToken()); + ASSERT_OK_AND_ASSIGN(DlfToken second, loader.LoadToken()); + ASSERT_EQ("ecs-ak", first.GetAccessKeyId()); + ASSERT_EQ("ecs-ak", second.GetAccessKeyId()); + ASSERT_EQ(1, http_client_ptr->GetRoleRequestCount()); + ASSERT_EQ(2, http_client_ptr->GetTokenRequestCount()); + ASSERT_EQ(180000, http_client_ptr->GetLastRequestTimeoutMillis()); +} + +TEST(DlfEcsTokenLoaderTest, ExplicitEmptyRoleUsesMetadataUrlAsTokenEndpoint) { + const std::string metadata_url = "http://100.100.100.200/metadata/token"; + auto http_client = std::make_unique(metadata_url, /*direct_token=*/true); + MockEcsHttpClient* http_client_ptr = http_client.get(); + DlfEcsTokenLoader loader(metadata_url, std::string(""), std::move(http_client)); + + ASSERT_OK_AND_ASSIGN(DlfToken token, loader.LoadToken()); + ASSERT_EQ("ecs-ak", token.GetAccessKeyId()); + ASSERT_EQ(0, http_client_ptr->GetRoleRequestCount()); + ASSERT_EQ(1, http_client_ptr->GetTokenRequestCount()); + ASSERT_EQ(180000, http_client_ptr->GetLastRequestTimeoutMillis()); +} + +TEST(DlfEcsTokenLoaderTest, PreservesTransportFailureDetail) { + DlfEcsTokenLoader loader("http://100.100.100.200/metadata/token", std::string(""), + std::make_unique()); + Status status = loader.LoadToken().status(); + ASSERT_NOK_WITH_MSG(status, "failed to request DLF credentials from ECS metadata service"); + ASSERT_NOK_WITH_MSG(status, "connection refused"); +} + +TEST(DlfAuthProviderTest, RefreshesWithinSafeWindow) { + std::atomic now_seconds{1744775086}; + std::vector tokens; + tokens.emplace_back("ak-1", "sk-1", std::nullopt, 1744782286000); + tokens.emplace_back("ak-2", "sk-2", std::nullopt, std::nullopt); + auto loader = std::make_unique(tokens); + SequenceTokenLoader* loader_ptr = loader.get(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr provider, + DlfAuthProvider::FromTokenLoader( + std::move(loader), "https://cn-beijing-vpc.dlf.aliyuncs.com", "cn-beijing", "default", + [&] { return std::chrono::system_clock::from_time_t(now_seconds.load()); })); + RestAuthParameter parameter = RestAuthParameter::Create("GET", "/v1/config", {}, ""); + + ASSERT_OK_AND_ASSIGN(StringMap first, provider->MergeAuthHeader({}, parameter)); + ASSERT_NE(std::string::npos, first.at("Authorization").find("Credential=ak-1/")); + ASSERT_OK(provider->MergeAuthHeader({}, parameter)); + ASSERT_EQ(1, loader_ptr->GetLoadCount()); + + now_seconds.fetch_add(3601); + ASSERT_OK_AND_ASSIGN(StringMap refreshed, provider->MergeAuthHeader({}, parameter)); + ASSERT_NE(std::string::npos, refreshed.at("Authorization").find("Credential=ak-2/")); + ASSERT_EQ(2, loader_ptr->GetLoadCount()); +} + +TEST(DlfAuthProviderTest, ConcurrentFirstUseLoadsTokenOnce) { + auto loader = std::make_unique( + std::vector{ + DlfToken("concurrent-ak", "concurrent-sk", std::nullopt, std::nullopt)}, + std::chrono::milliseconds(10)); + SequenceTokenLoader* loader_ptr = loader.get(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr provider, + DlfAuthProvider::FromTokenLoader(std::move(loader), + "https://cn-beijing-vpc.dlf.aliyuncs.com", + "cn-beijing", "default", FixedTime)); + RestAuthParameter parameter = RestAuthParameter::Create("GET", "/v1/config", {}, ""); + std::vector threads; + std::vector statuses; + std::mutex statuses_mutex; + for (int32_t i = 0; i < 16; ++i) { + threads.emplace_back([&] { + Status status = provider->MergeAuthHeader({}, parameter).status(); + std::scoped_lock lock(statuses_mutex); + statuses.push_back(status); + }); + } + for (std::thread& thread : threads) { + thread.join(); + } + ASSERT_EQ(16, statuses.size()); + for (const Status& status : statuses) { + ASSERT_OK(status); + } + ASSERT_EQ(1, loader_ptr->GetLoadCount()); +} + +TEST(DlfAuthProviderTest, SelectsEndpointSignerAndCredentialSource) { + std::map options = { + {CatalogOptions::URI, "https://dlfnext.cn-hangzhou.aliyuncs.com"}, + {CatalogOptions::TOKEN_PROVIDER, "dlf"}, + {CatalogOptions::DLF_ACCESS_KEY_ID, "ak"}, + {CatalogOptions::DLF_ACCESS_KEY_SECRET, "sk"}, + {CatalogOptions::DLF_SECURITY_TOKEN, "sts"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr provider, AuthProvider::Create(options)); + RestAuthParameter parameter = RestAuthParameter::Create("GET", "/v1/config", {}, ""); + ASSERT_OK_AND_ASSIGN(StringMap headers, provider->MergeAuthHeader({{"Authorization", "old"}, + {"x-acs-version", "old"}, + {"custom-header", "kept"}}, + parameter)); + ASSERT_NE(std::string::npos, headers.at("Authorization").find("acs ak:")); + ASSERT_EQ("2026-01-18", headers.at("x-acs-version")); + ASSERT_EQ("sts", headers.at("x-acs-security-token")); + ASSERT_EQ("kept", headers.at("custom-header")); + ASSERT_FALSE(provider->AllowsRedirects()); + + BearTokenAuthProvider bear_provider("token"); + ASSERT_TRUE(bear_provider.AllowsRedirects()); + + ASSERT_EQ("openapi", DlfAuthProvider::ParseSigningAlgorithmFromUri(options.at("uri"))); + ASSERT_EQ("default", DlfAuthProvider::ParseSigningAlgorithmFromUri( + "https://cn-hangzhou-vpc.dlf.aliyuncs.com")); + ASSERT_OK_AND_ASSIGN(std::string region, + DlfAuthProvider::ParseRegionFromUri(options.at("uri"))); + ASSERT_EQ("cn-hangzhou", region); + ASSERT_OK_AND_ASSIGN(std::string host, + DlfAuthProvider::ExtractHost("https://example.com:8443/prefix")); + ASSERT_EQ("example.com:8443", host); +} + +TEST(DlfAuthProviderTest, NormalizesSigningHostLikeTransport) { + std::map options = { + {CatalogOptions::URI, " https://dlfnext.cn-hangzhou.aliyuncs.com "}, + {CatalogOptions::TOKEN_PROVIDER, "dlf"}, + {CatalogOptions::DLF_ACCESS_KEY_ID, "ak"}, + {CatalogOptions::DLF_ACCESS_KEY_SECRET, "sk"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr provider, AuthProvider::Create(options)); + RestAuthParameter parameter = RestAuthParameter::Create("GET", "/v1/config", {}, ""); + ASSERT_OK_AND_ASSIGN(StringMap headers, provider->MergeAuthHeader({}, parameter)); + ASSERT_EQ("dlfnext.cn-hangzhou.aliyuncs.com", headers.at("Host")); +} + +TEST(DlfAuthProviderTest, RejectsIncompleteOrUnknownConfiguration) { + const std::map base = { + {CatalogOptions::URI, "https://dlfnext.cn-hangzhou.aliyuncs.com"}, + {CatalogOptions::TOKEN_PROVIDER, "dlf"}}; + ASSERT_NOK_WITH_MSG(AuthProvider::Create(base).status(), "token path or access key"); + + std::map options = base; + options[CatalogOptions::DLF_TOKEN_LOADER] = "ecs"; + options[CatalogOptions::DLF_TOKEN_ECS_METADATA_URL] = "http://metadata/"; + options[CatalogOptions::DLF_TOKEN_ECS_ROLE_NAME] = "role"; + ASSERT_OK(AuthProvider::Create(options)); + + options = base; + options[CatalogOptions::DLF_TOKEN_LOADER] = "unknown"; + ASSERT_NOK_WITH_MSG(AuthProvider::Create(options).status(), "unsupported DLF token loader"); + + options = base; + options[CatalogOptions::DLF_ACCESS_KEY_ID] = "ak"; + options[CatalogOptions::DLF_ACCESS_KEY_SECRET] = "sk"; + options[CatalogOptions::DLF_SIGNING_ALGORITHM] = "unknown"; + ASSERT_NOK_WITH_MSG(AuthProvider::Create(options).status(), + "unsupported DLF signing algorithm"); + + options[CatalogOptions::DLF_SIGNING_ALGORITHM] = "default"; + options[CatalogOptions::URI] = "http://127.0.0.1:8080"; + ASSERT_NOK_WITH_MSG(AuthProvider::Create(options).status(), "DLF region"); +} + +} // namespace paimon::test diff --git a/src/paimon/rest/mock_rest_server.cpp b/src/paimon/rest/mock_rest_server.cpp index 0bf4ec012..e9f0b36b7 100644 --- a/src/paimon/rest/mock_rest_server.cpp +++ b/src/paimon/rest/mock_rest_server.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/mock_rest_server.h b/src/paimon/rest/mock_rest_server.h index 240f661e6..c688ee398 100644 --- a/src/paimon/rest/mock_rest_server.h +++ b/src/paimon/rest/mock_rest_server.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/resource_paths.cpp b/src/paimon/rest/resource_paths.cpp index 484482378..d4b6adba4 100644 --- a/src/paimon/rest/resource_paths.cpp +++ b/src/paimon/rest/resource_paths.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/resource_paths.h b/src/paimon/rest/resource_paths.h index ea36bf480..009e6e073 100644 --- a/src/paimon/rest/resource_paths.h +++ b/src/paimon/rest/resource_paths.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/resource_paths_test.cpp b/src/paimon/rest/resource_paths_test.cpp index 782f3f690..a0ea565a4 100644 --- a/src/paimon/rest/resource_paths_test.cpp +++ b/src/paimon/rest/resource_paths_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_api.cpp b/src/paimon/rest/rest_api.cpp index 393a0938a..a28e19916 100644 --- a/src/paimon/rest/rest_api.cpp +++ b/src/paimon/rest/rest_api.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -22,6 +24,7 @@ #include "fmt/format.h" #include "paimon/catalog_options.h" +#include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/rapidjson_util.h" #include "paimon/common/utils/sensitive_config_utils.h" #include "paimon/logging.h" @@ -59,13 +62,13 @@ RestApi::RestApi(std::unique_ptr client, Result> RestApi::Create(const std::map& options, const std::string& warehouse, bool config_required, const RestHttpClient::Config& http_config) { - auto uri_iter = options.find(CatalogOptions::URI); - if (uri_iter == options.end() || uri_iter->second.empty()) { + Result uri = OptionsUtils::GetNonEmptyValueFromMap(options, CatalogOptions::URI); + if (!uri.ok()) { return Status::Invalid(fmt::format("option '{}' must be configured for the rest catalog", CatalogOptions::URI)); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr client, - RestHttpClient::Create(uri_iter->second, http_config)); + RestHttpClient::Create(uri.value(), http_config)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr auth_provider, AuthProvider::Create(options)); @@ -81,11 +84,12 @@ Result> RestApi::Create(const std::mapMergeAuthHeader(base_headers, auth_parameter)); - PAIMON_ASSIGN_OR_RAISE( - RestHttpClient::Response response, - client->Execute("GET", ResourcePaths::Config(), query_params, headers, "")); + bool follow_redirects = auth_provider->AllowsRedirects(); + PAIMON_ASSIGN_OR_RAISE(RestHttpClient::Response response, + client->Execute("GET", ResourcePaths::Config(), query_params, + headers, "", follow_redirects)); if (!response.IsSuccessful()) { - return ErrorToStatus(response); + return ErrorToStatus(response, follow_redirects); } ConfigResponse config; PAIMON_RETURN_NOT_OK(ParseResponseBody(response.body, ResourcePaths::Config(), &config)); @@ -109,7 +113,19 @@ Result> RestApi::Create(const std::map= 300 && response.code < HttpStatus::kBadRequest) { + std::string message = fmt::format( + "rest endpoint returned redirect status {}, which is not followed for " + "signed requests", + response.code); + std::string request_id = RestUtil::ExtractRequestId(response.headers); + if (request_id != RestUtil::kUnknownRequestId) { + message += fmt::format(" requestId:{}", request_id); + } + return Status::IOError(message).WithDetail( + std::make_shared(response.code)); + } // The code of the parsed error body takes precedence over the http status, which // a gateway may have rewritten. int64_t code = response.code; @@ -188,10 +204,12 @@ Result RestApi::Execute( } PAIMON_ASSIGN_OR_RAISE(StringMap headers, auth_provider_->MergeAuthHeader(request_headers, auth_parameter)); - PAIMON_ASSIGN_OR_RAISE(RestHttpClient::Response response, - client_->Execute(method, path, query_params, headers, body)); + bool follow_redirects = auth_provider_->AllowsRedirects(); + PAIMON_ASSIGN_OR_RAISE( + RestHttpClient::Response response, + client_->Execute(method, path, query_params, headers, body, follow_redirects)); if (!response.IsSuccessful()) { - return ErrorToStatus(response); + return ErrorToStatus(response, follow_redirects); } return response; } diff --git a/src/paimon/rest/rest_api.h b/src/paimon/rest/rest_api.h index c2394cedc..5d642a194 100644 --- a/src/paimon/rest/rest_api.h +++ b/src/paimon/rest/rest_api.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -105,9 +107,11 @@ class RestApi { /// Maps a non-successful http response to a status: 404 becomes `NotExist`, 409 /// becomes `Exist`, 400 becomes `Invalid`, 501 becomes `NotImplemented` and the - /// other codes become `IOError`. The status carries a `RestErrorDetail` with the - /// mapped code. - static Status ErrorToStatus(const RestHttpClient::Response& response); + /// other codes become `IOError`. A redirect returned while `follow_redirects` is + /// false is reported as deliberately rejected for a signed request. The status + /// carries a `RestErrorDetail` with the mapped code. + static Status ErrorToStatus(const RestHttpClient::Response& response, + bool follow_redirects = true); private: RestApi(std::unique_ptr client, std::unique_ptr auth_provider, diff --git a/src/paimon/rest/rest_auth.cpp b/src/paimon/rest/rest_auth.cpp index 65d6d7551..43513b701 100644 --- a/src/paimon/rest/rest_auth.cpp +++ b/src/paimon/rest/rest_auth.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -18,8 +20,10 @@ #include "fmt/format.h" #include "paimon/catalog_options.h" +#include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/common/utils/url_utils.h" +#include "paimon/rest/dlf_auth.h" namespace paimon { @@ -47,25 +51,30 @@ Result> BearTokenAuthProvider::MergeAuthHeade Result> AuthProvider::Create( const std::map& options) { - auto provider_iter = options.find(CatalogOptions::TOKEN_PROVIDER); - if (provider_iter == options.end() || provider_iter->second.empty()) { + Result provider_value = + OptionsUtils::GetNonEmptyValueFromMap(options, CatalogOptions::TOKEN_PROVIDER); + if (!provider_value.ok()) { return Status::Invalid(fmt::format("option '{}' must be configured for the rest catalog", CatalogOptions::TOKEN_PROVIDER)); } - // Matched leniently in lower case; other clients may match the provider name - // case-sensitively, so only the exact "bear" spelling is portable. - std::string provider = StringUtils::ToLowerCase(provider_iter->second); + // Matched leniently in lower case; other clients may match provider names + // case-sensitively, so the exact "bear" and "dlf" spellings are portable. + std::string provider = StringUtils::ToLowerCase(provider_value.value()); if (provider == "bear") { - auto token_iter = options.find(CatalogOptions::TOKEN); - if (token_iter == options.end() || token_iter->second.empty()) { + Result token = + OptionsUtils::GetNonEmptyValueFromMap(options, CatalogOptions::TOKEN); + if (!token.ok()) { return Status::Invalid( fmt::format("option '{}' must be configured for the bear token provider", CatalogOptions::TOKEN)); } - return std::make_unique(token_iter->second); + return std::make_unique(token.value()); + } + if (provider == "dlf") { + return DlfAuthProvider::Create(options); } - return Status::NotImplemented( - fmt::format("unsupported token provider: {}, only 'bear' is supported for now", provider)); + return Status::NotImplemented(fmt::format( + "unsupported token provider: {}, supported providers are 'bear' and 'dlf'", provider)); } } // namespace paimon diff --git a/src/paimon/rest/rest_auth.h b/src/paimon/rest/rest_auth.h index 22da08d99..639273f74 100644 --- a/src/paimon/rest/rest_auth.h +++ b/src/paimon/rest/rest_auth.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -52,6 +54,11 @@ class AuthProvider { const std::map& base_header, const RestAuthParameter& parameter) const = 0; + /// Whether the transport may follow a redirect without regenerating auth headers. + virtual bool AllowsRedirects() const { + return true; + } + /// Creates the provider configured by `CatalogOptions::TOKEN_PROVIDER`. static Result> Create( const std::map& options); diff --git a/src/paimon/rest/rest_catalog.cpp b/src/paimon/rest/rest_catalog.cpp index 5d2bb8238..eb86035c5 100644 --- a/src/paimon/rest/rest_catalog.cpp +++ b/src/paimon/rest/rest_catalog.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -50,7 +52,7 @@ constexpr const char kPathOption[] = "path"; // `BranchManager::IsMainBranch`, which names the branch directory of a table, stays // case-sensitive: this normalization only decides how a table is addressed on the server. std::optional NormalizeBranch(std::optional branch) { - if (branch && StringUtils::ToLowerCase(branch.value()) == Identifier::kDefaultMainBranch) { + if (branch && StringUtils::EqualsIgnoreCase(branch.value(), Identifier::kDefaultMainBranch)) { return std::nullopt; } return branch; diff --git a/src/paimon/rest/rest_catalog.h b/src/paimon/rest/rest_catalog.h index 984aad708..e952263aa 100644 --- a/src/paimon/rest/rest_catalog.h +++ b/src/paimon/rest/rest_catalog.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_catalog_test.cpp b/src/paimon/rest/rest_catalog_test.cpp index fda03cc4a..140f3f4c1 100644 --- a/src/paimon/rest/rest_catalog_test.cpp +++ b/src/paimon/rest/rest_catalog_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -414,7 +416,7 @@ TEST_F(RestCatalogTest, CreateRejectsInvalidOptions) { ASSERT_NOK_WITH_MSG(CreateRestCatalog().status(), "'token.provider' must be configured"); options_ = valid_options; - options_[CatalogOptions::TOKEN_PROVIDER] = "dlf"; + options_[CatalogOptions::TOKEN_PROVIDER] = "unsupported"; Status unsupported_provider = CreateRestCatalog().status(); ASSERT_TRUE(unsupported_provider.IsNotImplemented()) << unsupported_provider.ToString(); ASSERT_NOK_WITH_MSG(unsupported_provider, "unsupported token provider"); @@ -991,6 +993,12 @@ TEST(RestApiErrorTest, ErrorToStatus) { ASSERT_NOK_WITH_MSG(RestApi::ErrorToStatus(response), "rest request failed with code 429"); response.code = 418; ASSERT_NOK_WITH_MSG(RestApi::ErrorToStatus(response), "rest request failed with code 418"); + + response.code = 302; + Status signed_redirect = RestApi::ErrorToStatus(response, /*follow_redirects=*/false); + ASSERT_NOK_WITH_MSG(signed_redirect, "redirect status 302"); + ASSERT_NOK_WITH_MSG(signed_redirect, "not followed for signed requests"); + ASSERT_EQ(302, checked_pointer_cast(signed_redirect.detail())->GetCode()); } TEST(RestApiErrorTest, MalformedSuccessBodyFails) { diff --git a/src/paimon/rest/rest_http_client.cpp b/src/paimon/rest/rest_http_client.cpp index b17b0a281..f4964bf57 100644 --- a/src/paimon/rest/rest_http_client.cpp +++ b/src/paimon/rest/rest_http_client.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -123,7 +125,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; @@ -224,7 +228,8 @@ std::string RestHttpClient::NormalizeUri(const std::string& uri) { while (!normalized.empty() && normalized.back() == '/') { normalized.pop_back(); } - if (normalized.rfind("http://", 0) != 0 && normalized.rfind("https://", 0) != 0) { + if (!StringUtils::StartsWith(normalized, "http://") && + !StringUtils::StartsWith(normalized, "https://")) { normalized = "http://" + normalized; } return normalized; @@ -247,7 +252,7 @@ std::string RestHttpClient::BuildQueryString( Result RestHttpClient::ExecuteOnce( const std::string& method, const std::string& url, const std::map& headers, const std::string& body, - bool* transport_retriable) const { + bool follow_redirects, bool* transport_retriable) const { CURL* curl = handle_pool_->Acquire(); if (curl == nullptr) { return Status::IOError("failed to create curl handle"); @@ -269,10 +274,11 @@ Result RestHttpClient::ExecuteOnce( curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response.body); curl_easy_setopt(curl, CURLOPT_HEADERFUNCTION, WriteHeaderCallback); curl_easy_setopt(curl, CURLOPT_HEADERDATA, &response.headers); - // Follow redirects transparently, restricted to http(s) targets. Without + // Redirects can be disabled for auth headers whose signatures are bound to the + // original request. Otherwise they are restricted to http(s) targets. Without // CURLOPT_POSTREDIR a 301/302 would replay a body-carrying request as a bodyless // GET. A 303 is left to become a GET, which is what it is defined to mean. - curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L); + curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, follow_redirects ? 1L : 0L); curl_easy_setopt(curl, CURLOPT_MAXREDIRS, 50L); curl_easy_setopt(curl, CURLOPT_POSTREDIR, static_cast(CURL_REDIR_POST_301 | // NOLINT(runtime/int) @@ -397,7 +403,8 @@ std::optional RestHttpClient::GetRetryDelayMs(int32_t execution_count, Result RestHttpClient::Execute( const std::string& method, const std::string& path, const std::map& query_params, - const std::map& headers, const std::string& body) const { + const std::map& headers, const std::string& body, + bool follow_redirects) const { if (method != "GET" && method != "POST" && method != "DELETE") { return Status::Invalid(fmt::format("unsupported http method: {}", method)); } @@ -421,7 +428,8 @@ Result RestHttpClient::Execute( while (true) { execution_count++; bool transport_retriable = false; - Result result = ExecuteOnce(method, url, headers, body, &transport_retriable); + Result result = + ExecuteOnce(method, url, headers, body, follow_redirects, &transport_retriable); bool retriable; if (result.ok()) { retriable = IsRetriableCode(result.value().code); diff --git a/src/paimon/rest/rest_http_client.h b/src/paimon/rest/rest_http_client.h index 077127957..3ebd89472 100644 --- a/src/paimon/rest/rest_http_client.h +++ b/src/paimon/rest/rest_http_client.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -50,8 +52,8 @@ struct HttpStatus { /// response header (delta-seconds or HTTP-date form). A backoff sleep is bounded by /// `retry_max_delay_ms` and the whole request by `retry_timeout_ms`; a `Retry-After` /// beyond the remaining budget stops retrying rather than shortening the sleep. -/// Redirects to http(s) targets are followed transparently, keeping the method and -/// body of POST/DELETE requests. +/// Redirects to http(s) targets are followed by default, keeping the method and body +/// of POST/DELETE requests; callers can disable them for request-bound signatures. class RestHttpClient { public: struct Config { @@ -96,11 +98,12 @@ class RestHttpClient { /// final response, which may carry a non-2xx code, or an error status when the /// request could not be transported at all. Only transient transport errors (an /// established connection breaking mid-request or a truncated response body) are - /// retried; every other transport failure fails immediately. + /// retried; every other transport failure fails immediately. `follow_redirects` + /// must be false when authentication headers are bound to the original request. Result Execute(const std::string& method, const std::string& path, const std::map& query_params, const std::map& headers, - const std::string& body) const; + const std::string& body, bool follow_redirects = true) const; const std::string& GetBaseUri() const { return base_uri_; @@ -135,7 +138,8 @@ class RestHttpClient { /// may be retried; see `Execute` for which kinds are not. Result ExecuteOnce(const std::string& method, const std::string& url, const std::map& headers, - const std::string& body, bool* transport_retriable) const; + const std::string& body, bool follow_redirects, + bool* transport_retriable) const; std::optional GetRetryDelayMs(int32_t execution_count, const Response* response, int64_t remaining_budget_ms) const; diff --git a/src/paimon/rest/rest_http_client_test.cpp b/src/paimon/rest/rest_http_client_test.cpp index 34d07fe66..9b93705a2 100644 --- a/src/paimon/rest/rest_http_client_test.cpp +++ b/src/paimon/rest/rest_http_client_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -489,6 +491,25 @@ TEST(RestHttpClientTest, RedirectIsFollowed) { ASSERT_EQ(0, response.headers.count("location")); } +TEST(RestHttpClientTest, RedirectCanBeDisabledForSignedRequests) { + std::atomic request_count{0}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr server, + MockRestServer::Start([&](const MockRestServer::Request& request) { + request_count++; + MockRestServer::Response response; + response.code = 302; + response.headers["Location"] = "/v1/config"; + return response; + })); + ASSERT_OK_AND_ASSIGN(std::unique_ptr client, + RestHttpClient::Create(server->GetBaseUri())); + ASSERT_OK_AND_ASSIGN(RestHttpClient::Response response, + client->Execute("GET", "/old", {}, {{"x-acs-security-token", "secret"}}, + "", /*follow_redirects=*/false)); + ASSERT_EQ(302, response.code); + ASSERT_EQ(1, request_count.load()); +} + TEST(RestHttpClientTest, PostRedirectKeepsMethodAndBody) { std::mutex mutex; MockRestServer::Request last_request; diff --git a/src/paimon/rest/rest_messages.cpp b/src/paimon/rest/rest_messages.cpp index 324a8408b..df6911429 100644 --- a/src/paimon/rest/rest_messages.cpp +++ b/src/paimon/rest/rest_messages.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_messages.h b/src/paimon/rest/rest_messages.h index b943e6e11..0244fa362 100644 --- a/src/paimon/rest/rest_messages.h +++ b/src/paimon/rest/rest_messages.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_messages_test.cpp b/src/paimon/rest/rest_messages_test.cpp index 7a29b77ca..86fbfbd3d 100644 --- a/src/paimon/rest/rest_messages_test.cpp +++ b/src/paimon/rest/rest_messages_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_util.cpp b/src/paimon/rest/rest_util.cpp index b887d2567..1a746987d 100644 --- a/src/paimon/rest/rest_util.cpp +++ b/src/paimon/rest/rest_util.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, @@ -19,6 +21,7 @@ #include #include "fmt/format.h" +#include "paimon/common/utils/options_utils.h" #include "rapidjson/error/en.h" #include "rapidjson/stringbuffer.h" #include "rapidjson/writer.h" @@ -27,13 +30,7 @@ namespace paimon { std::map RestUtil::ExtractPrefixMap( const std::map& options, const std::string& prefix) { - std::map result; - for (const auto& [key, value] : options) { - if (key.size() > prefix.size() && key.compare(0, prefix.size(), prefix) == 0) { - result[key.substr(prefix.size())] = value; - } - } - return result; + return OptionsUtils::FetchOptionsWithPrefix(prefix, options); } std::string RestUtil::ExtractRequestId(const std::map& headers) { diff --git a/src/paimon/rest/rest_util.h b/src/paimon/rest/rest_util.h index 407217c67..1f5708c3a 100644 --- a/src/paimon/rest/rest_util.h +++ b/src/paimon/rest/rest_util.h @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/rest/rest_util_test.cpp b/src/paimon/rest/rest_util_test.cpp index e76baf156..c994e4462 100644 --- a/src/paimon/rest/rest_util_test.cpp +++ b/src/paimon/rest/rest_util_test.cpp @@ -1,11 +1,13 @@ /* - * Copyright 2026-present Alibaba Inc. + * 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 * - * Licensed 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 + * 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, diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index f05a2347b..4566289f4 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -29,6 +29,7 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/util/checked_cast.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -78,11 +79,14 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) override { - // Noted that SetReadSchema only change inner read_schema_, but take no effective on - // NextBatch PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, arrow::ImportSchema(read_schema)); read_schema_ = arrow_schema; + // A real FileBatchReader restarts from the first row and drops its assigned read ranges + // when the read schema is (re)set. Readers that switch schemas mid-file, such as the + // late-materialization reader moving from its probe pass to its payload pass, rely on it. + current_pos_ = 0; + previous_batch_first_row_num_ = std::numeric_limits::max(); return Status::OK(); } @@ -119,8 +123,29 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Result NextBatchWithBitmap() override { while (true) { PAIMON_RETURN_NOT_OK(next_batch_status_); - if (current_pos_ >= read_end_pos_) { - previous_batch_first_row_num_ = current_pos_; + int32_t begin_pos = current_pos_; + int32_t range_end_pos = read_end_pos_; + if (!read_ranges_.empty()) { + // Reading is restricted to the assigned ranges (ascending and half-open), like a + // real format reader, so that a prefetch reader may dispatch disjoint ranges to + // parallel readers. An empty range set means the whole file may be read. + const std::pair* selected = nullptr; + for (const auto& range : read_ranges_) { + if (static_cast(range.second) > begin_pos) { + selected = ⦥ + break; + } + } + if (selected == nullptr) { + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); + return BatchReader::MakeEofBatchWithBitmap(); + } + // Skip the gap in front of the first range that has not been read yet. + begin_pos = std::max(begin_pos, static_cast(selected->first)); + range_end_pos = std::min(range_end_pos, static_cast(selected->second)); + } + if (begin_pos >= read_end_pos_) { + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); return BatchReader::MakeEofBatchWithBitmap(); } int32_t actual_batch_size = batch_size_; @@ -128,21 +153,23 @@ class MockFileBatchReader : public PrefetchFileBatchReader { std::uniform_int_distribution distribution(1, batch_size_); actual_batch_size = distribution(random_engine_); } - int32_t batch_end_pos = std::min(read_end_pos_, current_pos_ + actual_batch_size); - auto slice = data_->Slice(current_pos_, batch_end_pos - current_pos_); + int32_t batch_end_pos = + std::min({read_end_pos_, range_end_pos, begin_pos + actual_batch_size}); + auto slice = data_->Slice(begin_pos, batch_end_pos - begin_pos); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr concat_slice, arrow::Concatenate({slice}, arrow::default_memory_pool())); RoaringBitmap32 bitmap; - for (auto iter = bitmap_.EqualOrLarger(current_pos_); + for (auto iter = bitmap_.EqualOrLarger(begin_pos); iter != bitmap_.End() && *iter < batch_end_pos; ++iter) { - bitmap.Add(*iter - current_pos_); + bitmap.Add(*iter - begin_pos); } - previous_batch_first_row_num_ = current_pos_; + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); current_pos_ = batch_end_pos; if (bitmap.IsEmpty()) { continue; } + PAIMON_ASSIGN_OR_RAISE(concat_slice, ProjectBatch(concat_slice)); std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( @@ -168,7 +195,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Result GetNumberOfRows() const override { return ToReaderRowNumber(read_end_pos_); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { return ToReaderRowNumber(current_pos_); } void Close() override {} @@ -181,7 +208,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader { return false; } - private: + protected: static uint64_t ToReaderRowNumber(int32_t row_number) { if (row_number < 0) { return std::numeric_limits::max(); @@ -189,6 +216,44 @@ class MockFileBatchReader : public PrefetchFileBatchReader { return static_cast(row_number); } + /// Pick the columns requested by `read_schema_` out of `batch`, in the requested order. + /// + /// `batch` is returned as is unless the requested schema is a genuine re-selection of the + /// columns this file has. Requesting a field the file does not have means the read schema is a + /// logical view over some other physical layout, as the shredding and the row tracking readers + /// do, and those map the raw batch themselves. + /// `batch` is expected to have a zero offset, so its validity buffer can be reused as is. + Result> ProjectBatch( + const std::shared_ptr& batch) const { + auto struct_batch = std::dynamic_pointer_cast(batch); + if (struct_batch == nullptr) { + return batch; + } + arrow::ArrayVector children; + arrow::FieldVector fields; + for (const auto& field : read_schema_->fields()) { + std::shared_ptr column = struct_batch->GetFieldByName(field->name()); + if (column == nullptr) { + return batch; + } + children.push_back(column); + fields.push_back(field); + } + const arrow::FieldVector& batch_fields = struct_batch->type()->fields(); + bool keeps_every_column = fields.size() == batch_fields.size(); + for (size_t i = 0; keeps_every_column && i < fields.size(); i++) { + keeps_every_column = fields[i]->name() == batch_fields[i]->name(); + } + if (keeps_every_column) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr projected, + arrow::StructArray::Make(children, fields, struct_batch->null_bitmap(), + struct_batch->null_count())); + return projected; + } + std::shared_ptr data_; std::shared_ptr file_schema_; std::shared_ptr read_schema_; diff --git a/src/paimon/testing/utils/read_result_collector.h b/src/paimon/testing/utils/read_result_collector.h index 8321a838a..558528662 100644 --- a/src/paimon/testing/utils/read_result_collector.h +++ b/src/paimon/testing/utils/read_result_collector.h @@ -138,7 +138,7 @@ class ReadResultCollector { return std::shared_ptr(); } auto& [c_array, c_schema] = batch; - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto array, + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(c_array.get(), c_schema.get())); return DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool()); } @@ -150,6 +150,11 @@ class ReadResultCollector { return std::make_pair(std::move(c_array), std::move(c_schema)); } + static Status CheckBatchOffset(const BatchReader::ReadBatch& batch) { + assert(!BatchReader::IsEofBatch(batch)); + return CheckArrayOffset(batch.first.get()); + } + // Noted that, sort chunked array by multiple key for timestamp type may cause // coredump in arrow, refer to https://github.com/apache/arrow/issues/47252 static Result> SortArray( @@ -197,16 +202,17 @@ class ReadResultCollector { } PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch)); } - auto& [c_array, c_schema] = batch; - assert(c_array->length > 0); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto result_array, - arrow::ImportArray(c_array.get(), c_schema.get())); - return result_array; + assert(batch.first->length > 0); + return ImportReadBatch(std::move(batch)); } - static Status CheckBatchOffset(const BatchReader::ReadBatch& batch) { - assert(!BatchReader::IsEofBatch(batch)); - return CheckArrayOffset(batch.first.get()); + static Result> ImportReadBatch(BatchReader::ReadBatch&& batch) { + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr record_batch, + arrow::ImportRecordBatch(c_array.get(), c_schema.get())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr struct_array, + record_batch->ToStructArray()); + return struct_array; } static Status CheckArrayOffset(const ArrowArray* array) { diff --git a/test/inte/append_compaction_inte_test.cpp b/test/inte/append_compaction_inte_test.cpp index b8a2c87a8..ebeb23361 100644 --- a/test/inte/append_compaction_inte_test.cpp +++ b/test/inte/append_compaction_inte_test.cpp @@ -155,6 +155,9 @@ class AppendCompactionInteTest : public testing::Test, std::vector GetTestValuesForAppendCompactionInteTest() { std::vector values; values.emplace_back("parquet"); +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic"); +#endif #ifdef PAIMON_ENABLE_ORC values.emplace_back("orc"); #endif @@ -254,7 +257,7 @@ TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteFullCompaction) { TEST_P(AppendCompactionInteTest, TestAppendTableStreamWriteFullCompactionWithMapSharedShredding) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } diff --git a/test/inte/blob_table_inte_test.cpp b/test/inte/blob_table_inte_test.cpp index ea951a311..f5d5960b3 100644 --- a/test/inte/blob_table_inte_test.cpp +++ b/test/inte/blob_table_inte_test.cpp @@ -63,6 +63,7 @@ #include "paimon/data/blob.h" #include "paimon/defs.h" #include "paimon/file_store_write.h" +#include "paimon/format/mosaic/mosaic_format_defs.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/global_index/bitmap_global_index_result.h" @@ -299,6 +300,7 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter auto splits = plan->Splits(); ReadContextBuilder read_context_builder(table_path); read_context_builder.SetReadFieldNames(read_schema).SetPredicate(predicate); + read_context_builder.EnableLateMaterializing(false); if (!options.empty()) { read_context_builder.SetOptions(options); } @@ -588,6 +590,9 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter std::vector GetTestValuesForBlobTableInteTest() { std::vector values; values.emplace_back("parquet"); +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic"); +#endif #ifdef PAIMON_ENABLE_ORC values.emplace_back("orc"); #endif @@ -2228,7 +2233,16 @@ TEST_P(BlobTableInteTest, TestPredicate) { // Avro does not have stats. return; } - CreateTable(); + if (GetParam() == "mosaic") { + CreateTable(/*partition_keys=*/{}, {{Options::MANIFEST_FORMAT, "orc"}, + {Options::FILE_FORMAT, GetParam()}, + {Options::FILE_SYSTEM, "local"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}, + {mosaic::MOSAIC_STATS_COLUMNS, "f0,f2"}}); + } else { + CreateTable(); + } std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); auto schema = arrow::schema(fields_); @@ -2365,7 +2379,7 @@ TEST_P(BlobTableInteTest, TestIOException) { TEST_P(BlobTableInteTest, TestReadTableWithDenseStats) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } std::string table_path = @@ -2429,7 +2443,7 @@ TEST_P(BlobTableInteTest, TestReadTableWithDenseStats) { TEST_P(BlobTableInteTest, TestDataEvolutionAndAlterTable) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } std::string table_path = paimon::test::GetDataDir() + file_format + @@ -2724,7 +2738,7 @@ TEST_P(BlobTableInteTest, TestAppendWriteWithNullBlob) { TEST_P(BlobTableInteTest, TestReadTableWithMultiBlobFields) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } std::string table_path = paimon::test::GetDataDir() + file_format + @@ -2796,6 +2810,9 @@ TEST_P(BlobTableInteTest, TestReadTableWithMultiBlobFields) { } TEST_P(BlobTableInteTest, TestBlobDescriptorField) { + if (GetParam() == "mosaic") { + return; + } // Two blob fields configured via BLOB_DESCRIPTOR_FIELD and stored inline as descriptors. arrow::FieldVector fields = {arrow::field("f0", arrow::int32()), BlobUtils::ToArrowField("b0", true), @@ -2849,6 +2866,9 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorField) { } TEST_P(BlobTableInteTest, TestBlobDescriptorFieldPartialInline) { + if (GetParam() == "mosaic") { + return; + } // 4 blob fields: b0,b1 are inline descriptors; b2,b3 are regular blob fields written to // .blob files. arrow::FieldVector fields = { @@ -2910,6 +2930,9 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorFieldPartialInline) { } TEST_P(BlobTableInteTest, TestBlobDescriptorMultiCommitAndShuffledReadSchema) { + if (GetParam() == "mosaic") { + return; + } // Multiple write+commit rounds with a shuffled read schema: b3, b2, b1, b0, f0. arrow::FieldVector fields = { arrow::field("f0", arrow::int32()), BlobUtils::ToArrowField("b0", true), @@ -3035,7 +3058,7 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorMultiCommitAndShuffledReadSchema) { // The shared-shredding map is read from one main data file while the blob payload is read from a // separate blob file with the same row-id range. TEST_P(BlobTableInteTest, TestSharedShreddingWithBlobDataEvolution) { - if (GetParam() == "avro") { + if (GetParam() == "avro" || GetParam() == "mosaic") { return; } @@ -3094,7 +3117,7 @@ TEST_P(BlobTableInteTest, TestSharedShreddingWithBlobDataEvolution) { // Two independent shared-shredding map columns are written into different main files. TEST_P(BlobTableInteTest, TestMultipleSharedShreddingMapsWithBlobDataEvolution) { - if (GetParam() == "avro") { + if (GetParam() == "avro" || GetParam() == "mosaic") { return; } @@ -3156,7 +3179,7 @@ TEST_P(BlobTableInteTest, TestMultipleSharedShreddingMapsWithBlobDataEvolution) // A newer partial data file rewrites only the shared-shredding map for the same row-id range. TEST_P(BlobTableInteTest, TestSharedShreddingMapOverrideWithBlobDataEvolution) { - if (GetParam() == "avro") { + if (GetParam() == "avro" || GetParam() == "mosaic") { return; } @@ -3280,6 +3303,9 @@ TEST_P(BlobTableInteTest, TestOrcMapStorageLayoutEvolutionWithBlobDataEvolution) } TEST_P(BlobTableInteTest, TestDataEvolutionWithBlobDescriptorField) { + if (GetParam() == "mosaic") { + return; + } // Test DataEvolution (split-column write) combined with blob descriptor fields. // Schema: f0(int32), b0/b1(blob descriptor inline), b2/b3(blob). // Commit 1: file A writes (f0, b2, b3) @@ -3401,6 +3427,9 @@ TEST_P(BlobTableInteTest, TestDataEvolutionWithBlobDescriptorField) { } TEST_P(BlobTableInteTest, TestBlobDescriptorFieldWriteRawBytesDirectly) { + if (GetParam() == "mosaic") { + return; + } // Similar to TestBlobDescriptorField but writes raw bytes directly without converting to // descriptor first. Descriptor fields reject values without the descriptor magic header. arrow::FieldVector fields = {arrow::field("f0", arrow::int32()), @@ -3431,7 +3460,7 @@ TEST_P(BlobTableInteTest, TestBlobDescriptorFieldWriteRawBytesDirectly) { TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamTable) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3595,7 +3624,7 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamTable) { TEST_P(BlobTableInteTest, TestForwardBlobViewReference) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3755,6 +3784,9 @@ TEST_P(BlobTableInteTest, TestForwardBlobViewReference) { TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamDescriptorBlob) { auto file_format = GetParam(); + if (file_format == "mosaic") { + return; + } // Upstream table has two blob descriptor fields. The downstream view references cells from // both b0 (field_id=1) and b1 (field_id=2). const std::string upstream_db_name = "upstream_two_blob"; @@ -3879,7 +3911,7 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithUpstreamDescriptorBlob) { TEST_P(BlobTableInteTest, TestBlobViewFieldWithMultipleUpstreamTables) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -4040,6 +4072,9 @@ TEST_P(BlobTableInteTest, TestBlobViewFieldWithMultipleUpstreamTables) { } TEST_P(BlobTableInteTest, TestBlobViewFailsWhenBothPathsAbsent) { + if (GetParam() == "mosaic") { + return; + } auto upstream_dir = UniqueTestDirectory::Create("local"); arrow::FieldVector fields = CreateBlobViewTable(upstream_dir->Str(), /*deletion_vectors_enabled=*/false); @@ -4055,6 +4090,9 @@ TEST_P(BlobTableInteTest, TestBlobViewFailsWhenBothPathsAbsent) { } TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceOfDeletedRow) { + if (GetParam() == "mosaic") { + return; + } auto upstream_dir = UniqueTestDirectory::Create("local"); arrow::FieldVector fields = CreateBlobViewTable(upstream_dir->Str(), /*deletion_vectors_enabled=*/true); @@ -4095,6 +4133,9 @@ TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceOfDeletedRow) { } TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceInEveryRowRangeGroup) { + if (GetParam() == "mosaic") { + return; + } auto upstream_dir = UniqueTestDirectory::Create("local"); arrow::FieldVector fields = CreateBlobViewTable(upstream_dir->Str(), /*deletion_vectors_enabled=*/true); @@ -4136,6 +4177,9 @@ TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceInEveryRowRangeGroup // predicate, so filtering a row out that way does not stop its blob view reference from being // resolved. This asserts the current behavior, not a desired one. TEST_P(BlobTableInteTest, TestBlobViewPreReadHonorsRowRangesNotPredicate) { + if (GetParam() == "mosaic") { + return; + } auto upstream_dir = UniqueTestDirectory::Create("local"); arrow::FieldVector fields = CreateBlobViewTable(upstream_dir->Str(), /*deletion_vectors_enabled=*/false); @@ -4169,6 +4213,9 @@ TEST_P(BlobTableInteTest, TestBlobViewPreReadHonorsRowRangesNotPredicate) { TEST_P(BlobTableInteTest, TestBlobViewWithFallbackPath) { auto file_format = GetParam(); + if (file_format == "mosaic") { + return; + } const std::string upstream_db_name = "fallback_db"; const std::string upstream_table_name = "fallback_table"; arrow::FieldVector upstream_fields = {arrow::field("f0", arrow::int32()), @@ -4280,7 +4327,7 @@ TEST_P(BlobTableInteTest, TestBlobViewWithFallbackPath) { TEST_P(BlobTableInteTest, TestReadBlobDescriptorFieldFromJava) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } std::string table_path = diff --git a/test/inte/data_evolution_table_test.cpp b/test/inte/data_evolution_table_test.cpp index 6232c1e7e..303c6d700 100644 --- a/test/inte/data_evolution_table_test.cpp +++ b/test/inte/data_evolution_table_test.cpp @@ -32,6 +32,7 @@ #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/defs.h" +#include "paimon/format/mosaic/mosaic_format_defs.h" #include "paimon/fs/file_system.h" #include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/global_index/indexed_split.h" @@ -90,22 +91,24 @@ class DataEvolutionTableTest : public ::testing::Test, return CreateTable(/*partition_keys=*/{}); } - Result>> WriteArray( + Result>> WriteArrays( const std::string& table_path, const std::map& partition, const std::vector& write_cols, - const std::shared_ptr& write_array) const { + const std::vector>& write_arrays) const { // write WriteContextBuilder write_builder(table_path, "commit_user_1"); write_builder.WithWriteSchema(write_cols); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr write_context, write_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); - ArrowArray c_array; - EXPECT_TRUE(arrow::ExportArray(*write_array, &c_array).ok()); - auto record_batch = std::make_unique( - partition, /*bucket=*/0, - /*row_kinds=*/std::vector(), &c_array); - PAIMON_RETURN_NOT_OK(file_store_write->Write(std::move(record_batch))); + for (const auto& write_array : write_arrays) { + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*write_array, &c_array)); + auto record_batch = std::make_unique( + partition, /*bucket=*/0, + /*row_kinds=*/std::vector(), &c_array); + PAIMON_RETURN_NOT_OK(file_store_write->Write(std::move(record_batch))); + } PAIMON_ASSIGN_OR_RAISE(auto commit_msgs, file_store_write->PrepareCommit( /*wait_compaction=*/false, /*commit_identifier=*/0)); @@ -113,6 +116,13 @@ class DataEvolutionTableTest : public ::testing::Test, return commit_msgs; } + Result>> WriteArray( + const std::string& table_path, const std::map& partition, + const std::vector& write_cols, + const std::shared_ptr& write_array) const { + return WriteArrays(table_path, partition, write_cols, {write_array}); + } + Result>> WriteArray( const std::string& table_path, const std::vector& write_cols, const std::shared_ptr& write_array) const { @@ -400,10 +410,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); @@ -839,6 +850,9 @@ TEST_P(DataEvolutionTableTest, TestOnlySomeColumns) { } TEST_P(DataEvolutionTableTest, TestMultipleSharedShreddingMapsPartialOverwrite) { + if (FileFormat() == "mosaic") { + return; + } if (FileFormat() == "avro") { return; } @@ -1477,6 +1491,9 @@ TEST_P(DataEvolutionTableTest, TestPartitionWithPredicate) { {Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, FileFormat()}, {Options::FILE_SYSTEM, "local"}, {Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"}, {"parquet.write.max-row-group-length", "1"}}; + if (file_format == "mosaic") { + options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0"); + } CreateTable(partition_keys, options); std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); @@ -1642,6 +1659,9 @@ TEST_P(DataEvolutionTableTest, TestPartitionWithPredicate) { TEST_P(DataEvolutionTableTest, TestAlterTable) { auto file_format = FileFormat(); + if (file_format == "mosaic") { + return; + } if (file_format == "avro") { return; } @@ -1738,6 +1758,9 @@ TEST_P(DataEvolutionTableTest, TestAlterTable) { } TEST_P(DataEvolutionTableTest, TestReadCompactFiles) { + if (FileFormat() == "mosaic") { + return; + } auto file_format = FileFormat(); if (file_format == "avro") { return; @@ -1768,6 +1791,9 @@ TEST_P(DataEvolutionTableTest, TestReadCompactFiles) { } TEST_P(DataEvolutionTableTest, TestReadTableWithDenseStats) { + if (FileFormat() == "mosaic") { + return; + } auto file_format = FileFormat(); if (file_format == "avro") { return; @@ -1849,6 +1875,9 @@ TEST_P(DataEvolutionTableTest, TestReadTableWithDenseStats) { } TEST_P(DataEvolutionTableTest, TestScanAndReadWithIndex) { + if (FileFormat() == "mosaic") { + return; + } auto file_format = FileFormat(); if (file_format == "avro") { return; @@ -1880,7 +1909,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 +1921,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 +1976,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,34 +1990,212 @@ 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) { + if (FileFormat() == "mosaic") { + return; + } + auto file_format = FileFormat(); + if (file_format == "avro") { + return; + } + std::string table_path = paimon::test::GetDataDir() + file_format + + "/data_evolution_with_index.db/data_evolution_with_index"; + + { + // A file without f0 must not interpret the predicate as f0 = null. + auto predicate = PredicateBuilder::Equal( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, "Lily", 4)); + auto read_type = + arrow::struct_({arrow::field("f0", arrow::utf8()), arrow::field("f2", arrow::int32())}); + auto expected_array = std::dynamic_pointer_cast( + arrow::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; + } + + std::map options = {{Options::FILE_INDEX_READ_ENABLED, "false"}, + {Options::WRITE_BATCH_SIZE, "1"}, + {"parquet.page.size", "1"}, + {"parquet.enable-dictionary", "false"}, + {"parquet.write.enable-page-index", "true"}, + {"parquet.read.enable-page-index-filter", "true"}, + {"orc.stripe.size", "1"}, + {"orc.row.index.stride", "1"}}; + if (FileFormat() == "mosaic") { + options.emplace(Options::FILE_BLOCK_SIZE, "1 B"); + options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0"); + } + CreateDataEvolutionTable(/*deletion_vectors_enabled=*/false, options); + std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + + auto input = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([ + [1, "a", "x"], + [2, "b", "y"], + [3, "c", "z"], + [4, "d", "w"] + ])") + .ValueOrDie()); + std::vector> write_arrays; + for (int64_t i = 0; i < input->length(); i++) { + arrow::ArrayVector children; + for (const auto& child : input->fields()) { + children.push_back(child->Slice(i, 1)); + } + write_arrays.push_back(arrow::StructArray::Make(children, fields_).ValueOrDie()); + } + ASSERT_OK_AND_ASSIGN(auto commit_messages, WriteArrays(table_path, /*partition=*/{}, + {"f0", "f1", "f2"}, write_arrays)); + ASSERT_OK(Commit(table_path, commit_messages)); + + auto predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, Literal(3)); + auto expected = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), R"([ + [3, "c", "z"] + ])") + .ValueOrDie()); + ASSERT_OK(ScanAndRead(table_path, {"f0", "f1", "f2"}, expected, predicate, + /*row_ranges=*/{}, + /*check_scan_plan_when_empty_result=*/true, + /*apply_predicate_to_scan=*/true)); +} + TEST_P(DataEvolutionTableTest, TestPredicate) { if (FileFormat() == "avro") { // Avro does not have stats. return; } - CreateTable(); + if (FileFormat() == "mosaic") { + CreateTable(/*partition_keys=*/{}, {{Options::MANIFEST_FORMAT, "orc"}, + {Options::FILE_FORMAT, FileFormat()}, + {Options::FILE_SYSTEM, "local"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}, + {mosaic::MOSAIC_STATS_COLUMNS, "f0,f1,f2"}}); + } else { + CreateTable(); + } std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); auto schema = arrow::schema(fields_); @@ -2135,6 +2336,9 @@ TEST_P(DataEvolutionTableTest, TestWithRowIds) { {Options::FILE_SYSTEM, "local"}, {Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}; + if (FileFormat() == "mosaic") { + options.emplace(mosaic::MOSAIC_STATS_COLUMNS, "f0,f1"); + } CreateTable(/*partition_keys=*/{}, options); std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); @@ -2951,6 +3155,9 @@ std::vector GetTestValuesForDataEvolutionTableTest() { std::vector values; for (bool enable_snapshot_live_manifest_cache : {false, true}) { values.emplace_back("parquet", enable_snapshot_live_manifest_cache); +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic", enable_snapshot_live_manifest_cache); +#endif #ifdef PAIMON_ENABLE_ORC values.emplace_back("orc", enable_snapshot_live_manifest_cache); #endif diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index ccb9ce324..84ab22df8 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -209,7 +209,8 @@ class GlobalIndexTest : public ::testing::Test, public ::testing::WithParamInter ReadContextBuilder read_context_builder(table_path); read_context_builder.SetReadFieldNames(read_schema) .SetPredicate(predicate) - .WithFileSystem(fs_); + .WithFileSystem(fs_) + .EnableLateMaterializing(false); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); diff --git a/test/inte/pk_compaction_inte_test.cpp b/test/inte/pk_compaction_inte_test.cpp index 18e77f9de..8e9150aee 100644 --- a/test/inte/pk_compaction_inte_test.cpp +++ b/test/inte/pk_compaction_inte_test.cpp @@ -445,7 +445,7 @@ TEST_F(PkCompactionInteTest, TestMetadataOnlyLevelUpgradeKeepsValueStats) { // Verify shared-shredding MAP can be read correctly after PK full compaction. TEST_P(PkCompactionInteTest, TestKeyValueTableFullCompactionWithMapSharedShredding) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -517,7 +517,7 @@ TEST_P(PkCompactionInteTest, TestKeyValueTableFullCompactionWithMapSharedShreddi TEST_P(PkCompactionInteTest, TestKeyValueTableDvCompactionWithMapSharedShredding) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3235,7 +3235,7 @@ TEST_F(PkCompactionInteTest, RemoteLookupFileWithSchemaEvolution) { // 6. ScanAndVerify after full compact TEST_P(PkCompactionInteTest, TestLookupCompatibility) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } // Step 1: Copy pk_compact_lookup table to temp dir. @@ -3597,6 +3597,9 @@ TEST_F(PkCompactionInteTest, AggHllAndThetaSketches) { std::vector GetTestValuesForCompactionInteTest() { std::vector values; values.emplace_back("parquet"); +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic"); +#endif #ifdef PAIMON_ENABLE_ORC values.emplace_back("orc"); #endif diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 6d7e55623..343d99605 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) @@ -182,7 +183,7 @@ class ReadInteTest : public testing::Test, public ::testing::WithParamInterface< /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt)); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt)); } auto bucket_str = bucket_path.substr(bucket_path.find("bucket-") + 7); int32_t bucket = std::stoi(bucket_str); @@ -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,13 +2202,14 @@ 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") .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) .SetPredicate(predicate) + .EnableLateMaterializing(false) .EnablePrefetch(param.enable_prefetch); ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); @@ -2197,6 +2266,89 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateOnlyPushdown) { ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); } +// Late materialization reads the predicate columns first and only materializes the remaining +// columns for matched rows. Combined with the top-level predicate filter, the read path returns +// the exact user-predicate match set. +TEST_P(ReadInteTest, TestAppendReadWithLateMaterializing) { + std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), + DataField(0, arrow::field("f0", arrow::utf8())), + DataField(1, arrow::field("f1", arrow::int32()))}; + ASSERT_OK_AND_ASSIGN( + auto predicate, + PredicateBuilder::Or( + {PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(static_cast(15.0))), + PredicateBuilder::IsNull(/*field_index=*/0, /*field_name=*/"f3", FieldType::DOUBLE)})); + + auto param = GetParam(); + std::string path = + paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; + + ReadContextBuilder context_builder(path); + 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") + .AddOption("test.enable-adaptive-prefetch-strategy", + param.enable_adaptive_prefetch_strategy) + .SetPredicate(predicate) + .EnableLateMaterializing(true) + .EnablePrefetch(param.enable_prefetch); + 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_0; + std::vector file_list_1; + std::vector file_list_2; + if (param.file_format == "orc") { + file_list_0 = {"data-d41fd7d1-b3e4-4905-aad9-b20a780e90a2-0.orc"}; + file_list_1 = {"data-4e30d6c0-f109-4300-a010-4ba03047dd9d-0.orc", + "data-10b9eea8-241d-4e4b-8ab8-2a82d72d79a2-0.orc", + "data-e2bb59ee-ae25-4e5b-9bcc-257250bc5fdd-0.orc", + "data-2d5ea1ea-77c1-47ff-bb87-19a509962a37-0.orc"}; + file_list_2 = {"data-db2b44c0-0d73-449d-82a0-4075bd2cb6e3-0.orc", + "data-b913a160-a4d1-4084-af2a-18333c35668e-0.orc"}; + } else if (param.file_format == "parquet") { + file_list_0 = {"data-46e27d5b-4850-4d1e-abb6-b3aabbbc08cb-0.parquet"}; + file_list_1 = {"data-864a052b-a938-4e04-b32c-6c72699a0c92-0.parquet", + "data-c0401350-64a3-4a54-a143-dd125ad9a8e5-0.parquet", + "data-7a912f84-04b7-4bbb-8dc6-53f4a292ea25-0.parquet", + "data-bb891df7-ea12-4b7e-9017-41aabe08c8ec-0.parquet"}; + file_list_2 = {"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=10/bucket-0", + BinaryRowGenerator::GenerateRow({10}, pool_.get()), file_list_0}, + {paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=10/bucket-1", + BinaryRowGenerator::GenerateRow({10}, pool_.get()), file_list_1}, + {paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=20/bucket-0", + BinaryRowGenerator::GenerateRow({20}, pool_.get()), file_list_2}}; + + auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(fields_with_row_kind); + + // "Bob" (f3 = 12.1) is the only row that does not match the predicate. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ + [0, 15.1, "Emily", 10], [0, 16.1, "Alex", 10], [0, 17.1, "David", 10], + [0, 17.1, "Lily", 10], [0, null, "Paul", 20] + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteTest, TestAppendReadWithPredicateAllFiltered) { std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), @@ -2210,7 +2362,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 +2449,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 +2494,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 +2540,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 +2606,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 +2691,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 +2769,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 +2865,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 +2944,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 +3017,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 +3102,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,9 +3190,10 @@ 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) + .EnableLateMaterializing(false) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); @@ -3098,6 +3251,91 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush ASSERT_TRUE(result_array->Equals(*expected_array)); } +TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing) { + std::vector read_fields = {DataField(1, arrow::field("key1", arrow::int32())), + DataField(7, arrow::field("k", arrow::utf8())), + DataField(2, arrow::field("key_2", arrow::int32())), + DataField(4, arrow::field("c", arrow::int32())), + DataField(8, arrow::field("d", arrow::int32())), + DataField(6, arrow::field("a", arrow::int32())), + DataField(0, arrow::field("key0", arrow::int32())), + DataField(9, arrow::field("e", arrow::int32()))}; + auto param = GetParam(); + std::string path = paimon::test::GetDataDir() + "/" + param.file_format + + "/pk_table_with_alter_table.db/pk_table_with_alter_table/"; + // equal is a partition filter and is not pushed into the data files; less_than is pushed down + // and only matches the column added by schema evolution, where the older files yield nulls. + auto equal = PredicateBuilder::Equal(/*field_index=*/6, /*field_name=*/"key0", FieldType::INT, + Literal(0)); + auto less_than = PredicateBuilder::LessThan(/*field_index=*/7, /*field_name=*/"e", + FieldType::INT, Literal(510)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({equal, less_than})); + + ReadContextBuilder context_builder(path); + 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.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); + context_builder.SetPredicate(predicate); + context_builder.EnableLateMaterializing(true) + .EnablePrefetch(param.enable_prefetch) + .AddOption("test.enable-adaptive-prefetch-strategy", + param.enable_adaptive_prefetch_strategy); + 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_0; + std::vector file_list_1; + std::string deletion_file; + if (param.file_format == "orc") { + file_list_0 = {"data-3842c1d6-6b34-4b2c-a648-9e95b4fb941b-0.orc", + "data-d6d370f3-242b-45c9-8739-44bf31b2b449-0.orc"}; + file_list_1 = {"data-7b538b91-5dbb-4e16-a639-1b5c0696db8c-0.orc"}; + deletion_file = "index-51804749-ed6c-4e7b-b3e9-337cfe38499c-1"; + } else if (param.file_format == "parquet") { + file_list_0 = {"data-8969384c-d715-4113-b663-2248c9a8c8d9-0.parquet", + "data-f2f38e80-7d28-4d51-90b3-c28951e5cdc0-0.parquet"}; + file_list_1 = {"data-d7a33230-223e-4d65-8e39-bc7ed26bdd32-0.parquet"}; + deletion_file = "index-c93829f3-1a72-4d88-8401-70663ce46426-1"; + } + + DataSplitsSchemaDv input_data_splits = { + {path + "key0=1/key1=1/bucket-0", + BinaryRowGenerator::GenerateRow({1, 1}, pool_.get()), + file_list_0, + /*schema ids*/ {0, 1}, + /*deletion file*/ + {DeletionFile(path + "index/" + deletion_file, + /*offset=*/1, /*length=*/26, /*cardinality=*/std::nullopt), + std::nullopt}}, + {path + "key0=0/key1=1/bucket-0", BinaryRowGenerator::GenerateRow({0, 1}, pool_.get()), + file_list_1, + /*schema ids*/ {1}, + /*deletion file*/ {std::nullopt}}}; + + auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(fields_with_row_kind); + + // "Paul" is the only row in partition key0 = 0 whose e is not null and matches e < 510. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ + [0, 1, "Bob", 22, 24, null, 26, 1, null], + [0, 1, "Emily", 32, 34, null, 36, 1, null], + [0, 1, "David", 62, 64, null, 66, 1, null], + [0, 1, "Whether I shall turn out to be the hero of my own life.", 72, 74, null, 76, 1, null], + [0, 1, "Paul", 502, 504, 508, 506, 0, 509] +])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateFilter) { std::vector read_fields = {DataField(1, arrow::field("key1", arrow::int32())), DataField(7, arrow::field("k", arrow::utf8())), @@ -3117,7 +3355,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 +3447,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 +3510,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 +3565,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 +3648,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 +3726,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 +3783,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 +3825,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 +3870,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 +3927,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 +4058,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/read_inte_with_index_test.cpp b/test/inte/read_inte_with_index_test.cpp index f1316c6ce..1442f6263 100644 --- a/test/inte/read_inte_with_index_test.cpp +++ b/test/inte/read_inte_with_index_test.cpp @@ -84,7 +84,8 @@ class ReadInteWithIndexTest : public testing::Test, ReadContextBuilder context_builder(table_path); context_builder.AddOption("read.batch-size", "2") .AddOption("test.enable-adaptive-prefetch-strategy", "false") - .SetPredicate(predicate); + .SetPredicate(predicate) + .EnableLateMaterializing(false); if (enable_prefetch) { context_builder.EnablePrefetch(true).SetPrefetchBatchCount(3); } @@ -944,7 +945,8 @@ TEST_P(ReadInteWithIndexTest, TestSimple) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1018,7 +1020,8 @@ TEST_P(ReadInteWithIndexTest, TestReadWithLimits) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1122,7 +1125,8 @@ TEST_P(ReadInteWithIndexTest, TestEmbeddingBitmapIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1185,7 +1189,8 @@ TEST_P(ReadInteWithIndexTest, TestBitmapWithV1) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1224,7 +1229,8 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1232,6 +1238,75 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndex) { CheckResultForBitmapWithSingleRowGroup(path, arrow_data_type, split); } +TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithLateMaterializing) { + auto [file_format, enable_prefetch] = GetParam(); + std::string path = GetDataDir() + "/" + file_format + + "/append_with_bitmap_no_embedding.db/append_with_bitmap_no_embedding/"; + std::string file_name; + if (file_format == "orc") { + file_name = "data-414509f5-e40c-4245-b992-bbf486778ac9-0.orc"; + } else if (file_format == "parquet") { + file_name = "data-783929b2-49d4-4006-a898-194a62e3278d-0.parquet"; + } + + std::vector read_fields = {SpecialFields::ValueKind(), + DataField(0, arrow::field("f0", arrow::utf8())), + DataField(1, arrow::field("f1", arrow::int32())), + DataField(2, arrow::field("f2", arrow::int32())), + DataField(3, arrow::field("f3", arrow::float64()))}; + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(read_fields); + + auto data_file_meta = std::make_shared( + file_name, /*file_size=*/689, + /*row_count=*/8, /*min_key=*/BinaryRow::EmptyRow(), + /*max_key=*/BinaryRow::EmptyRow(), /*key_stats=*/SimpleStats::EmptyStats(), + /*value_stats=*/SimpleStats::EmptyStats(), /*min_sequence_number=*/0, + /*max_sequence_number=*/7, /*schema_id=*/0, + /*level=*/0, + /*extra_files=*/ + std::vector>({file_name + ".index"}), + /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, + /*embedded_index=*/nullptr, FileSource::Append(), + /*value_stats_cols=*/std::nullopt, + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, + /*bucket_path=*/path + "bucket-0/", {data_file_meta}); + ASSERT_OK_AND_ASSIGN(auto split, + builder.WithSnapshot(1).IsStreaming(false).RawConvertible(true).Build()); + + std::string literal_str = "Bob"; + auto predicate = PredicateBuilder::Equal( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + + ReadContextBuilder context_builder(path); + context_builder.AddOption("read.batch-size", "2") + .AddOption("test.enable-adaptive-prefetch-strategy", "false") + .SetPredicate(predicate) + .EnablePredicateFilter(true) + .EnableLateMaterializing(true); + if (enable_prefetch) { + context_builder.EnablePrefetch(true).SetPrefetchBatchCount(3); + } + 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(std::vector>{split})); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + // Only the two "Bob" rows match the predicate. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ +[0, "Bob", 10, 1, 12.1], +[0, "Bob", 10, 1, 16.1] + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndexWithExternalPath) { auto [file_format, enable_prefetch] = GetParam(); std::string path = GetDataDir() + "/" + file_format + @@ -1270,7 +1345,7 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndexWithExternalPath) { /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/external_file_path, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1312,7 +1387,8 @@ TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithDv) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DeletionFile deletion_file(deletion_file_path, /*offset=*/1, /*length=*/24, /*cardinality=*/2); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, @@ -1418,7 +1494,7 @@ TEST_P(ReadInteWithIndexTest, TestWithAlterTable) { /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); }; std::vector embedded_bytes1 = { @@ -1801,7 +1877,8 @@ TEST_P(ReadInteWithIndexTest, TestWithBsiIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -1861,7 +1938,8 @@ TEST_P(ReadInteWithIndexTest, TestWithBloomFilterIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -2067,7 +2145,8 @@ TEST_P(ReadInteWithIndexTest, TestBitmapPushDownWithMultiStripes) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -2177,7 +2256,8 @@ TEST_P(ReadInteWithIndexTest, TestWithBitmapAndBsiAndBloomFilterIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -2259,7 +2339,8 @@ TEST_P(ReadInteWithIndexTest, TestWithIndexWithoutRegistered) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, @@ -2396,7 +2477,8 @@ TEST_P(ReadInteWithIndexTest, TestRangeBitmapIndex) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); @@ -2441,7 +2523,8 @@ TEST_P(ReadInteWithIndexTest, TestRangeBitmapIndexMultiChunk) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); @@ -2483,7 +2566,8 @@ TEST_P(ReadInteWithIndexTest, TestWithIOException) { /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/path + "bucket-0/", {data_file_meta}); ASSERT_OK_AND_ASSIGN(auto split, diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 67cc21d01..4e28f286f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -23,10 +23,13 @@ #include #include #include +#include #include +#include #include #include #include +#include #include #include #include @@ -40,21 +43,31 @@ #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/commit_context.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" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/table/source/realtime_split.h" #include "paimon/core/utils/snapshot_manager.h" #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 +75,142 @@ #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 { + +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 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"); + } + return delegate_->CreateQueryReaders(tracking_view->Delegate(), offset_begin, context); + } + + private: + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +} // 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 +329,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 +351,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 +386,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 +399,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 +415,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,6 +481,104 @@ class RealtimeWriteInteTest : public ::testing::Test { /*watermark=*/std::nullopt); } + 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"); + } + return latest_snapshot->Id(); + } + + 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 { @@ -292,6 +594,27 @@ class RealtimeWriteInteTest : public ::testing::Test { 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, @@ -314,6 +637,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 +720,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 +747,80 @@ 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"); + } + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_)); + SchemaManager schema_manager(core_options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional> table_schema, + schema_manager.Latest()); + if (!table_schema) { + return Status::Invalid("expected a table schema"); + } + auto read_schema = std::make_unique(); + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(table_schema.value()->Fields()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema( + *RealtimePrimaryKeyLayout::CreateSchema(value_schema->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 +846,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 +870,75 @@ 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()); + } + std::unique_ptr dir_; std::string table_path_; std::string commit_user_ = "realtime_commit_user"; @@ -436,6 +948,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,71 +996,1200 @@ TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } -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, TestCommitOrdersPreparedOffsetRanges) { - CreateTable(/*partition_keys=*/{}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); +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 = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + 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)); + 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::vector commits, + 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, commits.size()); - ASSERT_EQ(OffsetRange(0, 3), commits[0].offset_range); + ASSERT_EQ(1, progress.size()); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); - std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + 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)); + 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, TestPkKeylessProjection) { + 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 disk_batch, + MakeBatch({Row{1, "disk", "p0"}}, /*partitioned=*/false)); + 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 memory_batch, + MakeBatch({Row{1, "memory", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadPlanWithSchemaAndCheck(plan, realtime_context, + arrow::schema({arrow::field("payload", arrow::utf8())}), R"([ + [0, "memory"] + ])"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestCompositePkKeylessProjection) { + 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, "key", "disk"}}, /*partitioned=*/false)); + 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 memory_batch, + MakeBatch({Row{1, "key", "memory"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadPlanWithSchemaAndCheck(plan, realtime_context, + arrow::schema({arrow::field("pt", arrow::utf8())}), R"([ + [0, "memory"] + ])"); + 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::pair(0, 1)), first_sequences.at(p0b0)); + ASSERT_EQ((std::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::pair(2, 3)), second_sequences.at(p0b0)); + ASSERT_EQ((std::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 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(); + 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 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))); - 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); + + 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, 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()); + 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))); - 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()); +} + +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::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows; + 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))); + first_rows.insert(first_rows.end(), rows.begin(), rows.end()); + } + ASSERT_OK_AND_ASSIGN(std::vector first_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + 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 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, 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()); + + 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)); + + 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(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()); +} + +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) { @@ -671,52 +2350,6 @@ 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"); - - 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, TestCloseWriterKeepsContextReadable) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -816,6 +2449,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, @@ -833,98 +2504,402 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()), arrow::field("id", arrow::int64())}); - std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + 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, scan_predicate)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_result, + ReadPlan(memory_plan, realtime_context, read_fields, read_predicate, + /*enable_predicate_filter=*/true)); + std::shared_ptr expected_memory = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "value-2", 2] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, memory_result.data); + ASSERT_TRUE( + std::make_shared(expected_memory)->Equals(*memory_result.data)); + + 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 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::shared_ptr memory_plan, - CreatePlan(realtime_context, scan_predicate)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_result, - ReadPlan(memory_plan, realtime_context, read_fields, read_predicate, - /*enable_predicate_filter=*/true)); - std::shared_ptr expected_memory = - arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ - [0, "value-2", 2] - ])") - .ValueOrDie(); - ASSERT_NE(nullptr, memory_result.data); - ASSERT_TRUE( - std::make_shared(expected_memory)->Equals(*memory_result.data)); - 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 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 +3084,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 +3253,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 +3328,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 +3470,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 +3512,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 +3580,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=*/{}); diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 96a7c9a19..11080c9d5 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -724,7 +724,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -744,6 +744,45 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLateMaterializing) { + auto file_format = FileFormat(); + std::string table_path = paimon::test::GetDataDir() + file_format + + "/pk_table_scan_and_read_dv.db/pk_table_scan_and_read_dv/"; + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID, "6"); + + std::string literal_str = "Alice"; + auto not_equal = PredicateBuilder::NotEqual( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + auto greater_than = PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(18.0)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({not_equal, greater_than})); + scan_context_builder.SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder)); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + + ReadContextBuilder read_context_builder(table_path); + AddReadOptionsForPrefetch(&read_context_builder); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(true); + 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 result_plan, table_scan->CreatePlan()); + ASSERT_EQ(result_plan->SnapshotId().value(), 6); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + + // check result: "Lucy" (f3 = 14.1) does not match f3 > 18 and is filtered out. + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type_, R"([ +[0, "Paul", 20, 1, 18.1] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected); + ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); +} + TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot4WithPredicate) { auto file_format = FileFormat(); std::string table_path = paimon::test::GetDataDir() + file_format + @@ -1251,7 +1290,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -1279,6 +1318,55 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +// Same coverage as the deletion-vector case above, for the merge-on-read path where only the +// key part of the predicate is pushed down into the data files. +TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithLateMaterializing) { + auto file_format = FileFormat(); + std::string table_path = paimon::test::GetDataDir() + file_format + + "/pk_table_scan_and_read_mor.db/pk_table_scan_and_read_mor/"; + + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID, "5"); + + std::string literal_str = "Alice"; + auto not_equal = PredicateBuilder::NotEqual( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + std::string literal_str2 = "Lucy"; + auto less_than = PredicateBuilder::LessThan( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str2.data(), literal_str2.size())); + auto less_or_equal = PredicateBuilder::LessOrEqual(/*field_index=*/3, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(30.0)); + ASSERT_OK_AND_ASSIGN(auto predicate, + PredicateBuilder::And({not_equal, less_than, less_or_equal})); + scan_context_builder.SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder)); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + + ReadContextBuilder read_context_builder(table_path); + AddReadOptionsForPrefetch(&read_context_builder); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(true); + 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 result_plan, table_scan->CreatePlan()); + ASSERT_EQ(result_plan->SnapshotId().value(), 5); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + + // check result: only the rows before "Lucy" with f3 <= 30.0 remain. + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type_, R"([ +[0, "Bob", 10, 0, 12.1], +[0, "David", 10, 0, 17.1], +[0, "Emily", 10, 0, 13.1] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected); + ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); +} + TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot3WithPredicate) { auto file_format = FileFormat(); std::string table_path = paimon::test::GetDataDir() + file_format + @@ -2664,6 +2752,97 @@ TEST_P(ScanAndReadInteTest, TestCastTimestampType) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +#ifdef PAIMON_ENABLE_MOSAIC +TEST_F(ScanAndReadInteTest, TestMosaicJavaAndPythonCompatibility) { + TimezoneGuard timezone_guard("UTC"); + std::string timezone = DateTimeUtils::GetLocalTimezoneName(); + arrow::FieldVector fields = { + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("id", arrow::int32()), + arrow::field("f_boolean", arrow::boolean()), + arrow::field("f_tinyint", arrow::int8()), + arrow::field("f_smallint", arrow::int16()), + arrow::field("f_bigint", arrow::int64()), + arrow::field("f_float", arrow::float32()), + arrow::field("f_double", arrow::float64()), + arrow::field("f_char", arrow::utf8()), + arrow::field("f_varchar", arrow::utf8()), + arrow::field("f_binary", arrow::binary()), + arrow::field("f_varbinary", arrow::binary()), + arrow::field("f_date", arrow::date32()), + arrow::field("f_ts_3", arrow::timestamp(arrow::TimeUnit::MILLI)), + arrow::field("f_ts_6", arrow::timestamp(arrow::TimeUnit::MICRO)), + arrow::field("f_ts_9", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("f_ltz_3", arrow::timestamp(arrow::TimeUnit::MILLI, timezone)), + arrow::field("f_ltz_6", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)), + arrow::field("f_ltz_9", arrow::timestamp(arrow::TimeUnit::NANO, timezone)), + arrow::field("f_decimal_1_0", arrow::decimal128(1, 0)), + arrow::field("f_decimal_18_2", arrow::decimal128(18, 2)), + arrow::field("f_decimal_19_2", arrow::decimal128(19, 2)), + arrow::field("f_decimal_38_18", arrow::decimal128(38, 18)), + arrow::field("f_array_int", arrow::list(arrow::int32())), + arrow::field("f_map_numeric", arrow::map(arrow::int8(), arrow::int16())), + arrow::field("f_map_string_bigint", arrow::map(arrow::utf8(), arrow::int64())), + arrow::field("f_array_array_int", arrow::list(arrow::list(arrow::int32()))), + arrow::field("f_array_map", arrow::list(arrow::map(arrow::utf8(), arrow::int32()))), + arrow::field("f_map_array", arrow::map(arrow::utf8(), arrow::list(arrow::int32()))), + }; + std::shared_ptr data_type = arrow::struct_(fields); + std::shared_ptr expected_array = + arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ +[0, 1, true, -5, -1000, 10000000001, 1.25, 10.5, "char0001", "value-1", "bin00001", "\u0001\u0002\u0003", 20000, "1970-01-01 00:00:01.123", "1970-01-01 00:00:01.123456", "1970-01-01 00:00:01.123456789", "1970-01-01 00:01:01.321", "1970-01-01 00:01:01.654321", "1970-01-01 00:01:01.321654987", "-3", "1.25", "12345678901234567.89", "12345678901234567890.123456789012345678", [1, 2], [[0, 0], [10, 1]], [["k0", 1000], ["z0", 2000]], [[1, 2], [3]], [[["nested0", 0]], [["nested10", 10]]], [["array0", [0, 1]]]], +[0, 2, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null], +[0, 10, true, -3, -998, 10000000010, 3.25, 12.5, "char0010", "value-10", "bin00010", "\u000a\u000b\u000c", 20002, "1970-01-01 00:00:10.123", "1970-01-01 00:00:10.123456", "1970-01-01 00:00:10.123456789", "1970-01-01 00:01:10.321", "1970-01-01 00:01:10.654321", "1970-01-01 00:01:10.321654987", "-1", "10.25", "12345678901234569.89", "12345678901234567892.123456789012345678", [10, 11], [[2, 20], [12, 21]], [["k2", 1002], ["z2", 2002]], [[10, 11], [12]], [[["nested2", 2]], [["nested12", 12]]], [["array2", [2, 3]]]], +[0, 11, false, -2, -997, 10000000011, 4.25, 13.5, "char0011", "value-11", "bin00011", "\u000b\u000c\u000d", 20003, "1970-01-01 00:00:11.123", "1970-01-01 00:00:11.123456", "1970-01-01 00:00:11.123456789", "1970-01-01 00:01:11.321", "1970-01-01 00:01:11.654321", "1970-01-01 00:01:11.321654987", "0", "11.25", "12345678901234570.89", "12345678901234567893.123456789012345678", [11, 12], [[3, 30], [13, 31]], [["k3", 1003], ["z3", 2003]], [[11, 12], [13]], [[["nested3", 3]], [["nested13", 13]]], [["array3", [3, 4]]]], +[0, 20, true, -1, -996, 10000000020, 5.25, 14.5, "char0020", "value-20", "bin00020", "\u0014\u0015\u0016", 20004, "1970-01-01 00:00:20.123", "1970-01-01 00:00:20.123456", "1970-01-01 00:00:20.123456789", "1970-01-01 00:01:20.321", "1970-01-01 00:01:20.654321", "1970-01-01 00:01:20.321654987", "1", "20.25", "12345678901234571.89", "12345678901234567894.123456789012345678", [20, 21], [[4, 40], [14, 41]], [["k4", 1004], ["z4", 2004]], [[20, 21], [22]], [[["nested4", 4]], [["nested14", 14]]], [["array4", [4, 5]]]], +[0, 21, false, 0, -995, 10000000021, 6.25, 15.5, "char0021", "value-21", "bin00021", "\u0015\u0016\u0017", 20005, "1970-01-01 00:00:21.123", "1970-01-01 00:00:21.123456", "1970-01-01 00:00:21.123456789", "1970-01-01 00:01:21.321", "1970-01-01 00:01:21.654321", "1970-01-01 00:01:21.321654987", "2", "21.25", "12345678901234572.89", "12345678901234567895.123456789012345678", [21, 22], [[5, 50], [15, 51]], [["k5", 1005], ["z5", 2005]], [[21, 22], [23]], [[["nested5", 5]], [["nested15", 15]]], [["array5", [5, 6]]]] +])") + .ValueOrDie(); + auto expected = std::make_shared(expected_array); + + auto check_compatibility = [](const std::string& table_path, + const std::shared_ptr& predicate, + const std::shared_ptr& expected_result) { + ScanContextBuilder scan_context_builder(table_path); + ReadContextBuilder read_context_builder(table_path); + if (predicate != nullptr) { + scan_context_builder.SetPredicate(predicate); + read_context_builder.SetPredicate(predicate); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, + scan_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, table_scan->CreatePlan()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, + read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_TRUE(expected_result->Equals(actual)) + << "actual: " << (actual == nullptr ? "null" : actual->ToString()) + << "\nexpected: " << expected_result->ToString(); + }; + + std::shared_ptr predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(20)); + std::shared_ptr predicate_without_stats = PredicateBuilder::GreaterOrEqual( + /*field_index=*/4, /*field_name=*/"f_bigint", FieldType::BIGINT, + Literal(int64_t{10000000020LL})); + for (const std::string table_name : {"append_java_compat", "append_python_compat"}) { + SCOPED_TRACE(table_name); + std::string table_path = GetDataDir() + "/mosaic/" + table_name + ".db/" + table_name; + check_compatibility(table_path, /*predicate=*/nullptr, expected); + check_compatibility(table_path, predicate, expected->Slice(4, 2)); + check_compatibility(table_path, predicate_without_stats, expected); + } +} +#endif + TEST_F(ScanAndReadInteTest, TestAvroWithAppendTable) { auto read_data = [](int64_t snapshot_id, const std::string& result_json) { std::string table_path = GetDataDir() + "/avro/append_multiple.db/append_multiple"; diff --git a/test/inte/scan_inte_test.cpp b/test/inte/scan_inte_test.cpp index fdc38ea9f..f626dba6b 100644 --- a/test/inte/scan_inte_test.cpp +++ b/test/inte/scan_inte_test.cpp @@ -148,7 +148,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot1_partition10_bucket1_ = std::make_shared( @@ -164,7 +164,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot1_partition20_bucket0_ = std::make_shared( @@ -180,7 +180,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot2_partition10_bucket1_ = std::make_shared( @@ -196,7 +196,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot2_partition20_bucket0_ = std::make_shared( @@ -212,7 +212,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot3_partition10_bucket1_ = std::make_shared( @@ -228,7 +228,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot4_partition10_bucket1_ = std::make_shared( @@ -244,7 +244,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); std::shared_ptr meta_snapshot5_partition10_bucket1_ = std::make_shared( @@ -260,7 +260,7 @@ class ScanInteTest : public testing::TestWithParam { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); }; TEST(ScanInteManifestCacheTest, TestRepeatedScanReusesManifestCache) { @@ -1273,9 +1273,9 @@ TEST_P(ScanInteTest, TestScanAppendWithStreamWithAndPredicate) { .value()); std::vector>> expected_data_splits = { - {}, {expected_data_split1_2}, {expected_data_split2_1}, {}, {expected_data_split4_1}}; + {}, {expected_data_split1_2}, {expected_data_split2_1}, {expected_data_split4_1}}; - std::vector> expected_snapshot_ids = {std::nullopt, 1, 2, 3, 4}; + std::vector> expected_snapshot_ids = {std::nullopt, 1, 2, 4}; CheckStreamScanResult(table_scan.get(), expected_snapshot_ids, expected_data_splits); } @@ -1422,7 +1422,7 @@ TEST_P(ScanInteTest, TestScanAppendWithSnapshot1WithMultiPartitionKeys) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1( BinaryRowGenerator::GenerateRow({10, 0}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1483,7 +1483,7 @@ TEST_P(ScanInteTest, TestScanAppendComplexDataWithSnapshot4WithPredicateFilter) /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1( BinaryRowGenerator::GenerateRow({10}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1549,7 +1549,7 @@ TEST_P(ScanInteTest, TestScanAppendComplexDataWithSnapshot4WithPredicateFilter2) /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Compact(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1( BinaryRowGenerator::GenerateRow({10}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1591,7 +1591,8 @@ TEST_P(ScanInteTest, TestScanAppendWithSnapshot1WithEnableStatsDenseStore) { /*creation_time=*/Timestamp(1731412938869ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta2 = std::make_shared( "data-c2613568-0412-4cd9-a0c4-1eae8e4ca89b-0.orc", /*file_size=*/575, /*row_count=*/3, /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), @@ -1603,7 +1604,8 @@ TEST_P(ScanInteTest, TestScanAppendWithSnapshot1WithEnableStatsDenseStore) { /*creation_time=*/Timestamp(1731412938891ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-a6d1261a-f798-4fbd-a251-6d6c7d8060dd-0.orc", /*file_size=*/541, /*row_count=*/1, /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), @@ -1615,7 +1617,8 @@ TEST_P(ScanInteTest, TestScanAppendWithSnapshot1WithEnableStatsDenseStore) { /*creation_time=*/Timestamp(1731412938908ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1( BinaryRowGenerator::GenerateRow({10}, pool_.get()), @@ -1689,7 +1692,8 @@ TEST_P(ScanInteTest, TestScanAppendWithSnapshot1WithEnableStatsDenseStore2) { /*creation_time=*/Timestamp(1731412938891ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( BinaryRowGenerator::GenerateRow({10}, pool_.get()), @@ -1787,7 +1791,8 @@ TEST_P(ScanInteTest, TestScanAppendWithAlterTableWithCast) { /*creation_time=*/Timestamp(1732635461460ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRowGenerator::GenerateRow({1, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1833,7 +1838,8 @@ TEST_P(ScanInteTest, TestScanAppendWithAlterTableWithNoCast) { /*creation_time=*/Timestamp(1730458825047ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder1(BinaryRowGenerator::GenerateRow({1, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1860,7 +1866,8 @@ TEST_P(ScanInteTest, TestScanAppendWithAlterTableWithNoCast) { /*creation_time=*/Timestamp(1730459969493ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder2(BinaryRowGenerator::GenerateRow({0, 1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1910,7 +1917,8 @@ TEST_P(ScanInteTest, TestScanAppendWithAlterTableWithDenseField) { /*creation_time=*/Timestamp(1751647880163ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::optional>({"key0", "f0", "f1", "f2"}), - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder( BinaryRowGenerator::GenerateRow({1}, pool_.get()), /*bucket=*/0, /*bucket_path=*/ @@ -1992,7 +2000,8 @@ TEST_P(ScanInteTest, TestScanAppendWithBitmapEmbeddedIndex) { /*creation_time=*/Timestamp(1745000702835ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/table_path + "bucket-0", {file_meta}); ASSERT_OK_AND_ASSIGN(auto expected_data_split, builder.WithTotalBuckets(-1) @@ -2057,7 +2066,8 @@ TEST_P(ScanInteTest, TestScanAppendWithBitmapNoEmbeddedIndex) { /*creation_time=*/Timestamp(1745235371029ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/table_path + "bucket-0", {file_meta}); ASSERT_OK_AND_ASSIGN(auto expected_data_split, builder.WithTotalBuckets(-1) @@ -2124,7 +2134,8 @@ TEST_P(ScanInteTest, TestScanAppendWithBitmapAndAlterTable) { /*creation_time=*/Timestamp(1745253323731ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/table_path + "bucket-0", {file_meta}); @@ -2198,7 +2209,8 @@ TEST_P(ScanInteTest, TestScanAppendWithBitmapAndAlterTable3) { /*creation_time=*/Timestamp(1745251357742ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/table_path + "bucket-0", {file_meta}); @@ -2274,7 +2286,8 @@ TEST_P(ScanInteTest, TestScanAppendWithBitmapAndAlterTable2) { /*creation_time=*/Timestamp(1745251357742ll, 0), /*delete_row_count=*/0, /*embedded_index=*/embedded_index, FileSource::Append(), /*value_stats_cols=*/std::nullopt, - /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/table_path + "bucket-0", {file_meta}); diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp index d41b1a4a1..43ed0d1fb 100644 --- a/test/inte/write_and_read_inte_test.cpp +++ b/test/inte/write_and_read_inte_test.cpp @@ -90,6 +90,25 @@ class WriteAndReadInteTest return new_options; } + std::string LookupTempDirectory() const { + return PathUtil::JoinPath(test_dir_, "tmp"); + } + + Result> CreateLookupTestHelper( + const std::shared_ptr& schema, + const std::map& options) const { + return TestHelper::Create(test_dir_, schema, /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, + /*is_streaming_mode=*/true, /*ignore_if_exists=*/false, + LookupTempDirectory()); + } + + Result> CreateLookupTestHelper( + const std::string& table_path, const std::map& options) const { + return TestHelper::Create(table_path, options, /*is_streaming_mode=*/true, + LookupTempDirectory()); + } + Status WriteNextSchema(const std::vector& fields, int32_t highest_field_id, const std::map& options) const { return TestHelper::WriteNextSchema(dir_->GetFileSystem(), @@ -129,6 +148,33 @@ class WriteAndReadInteTest return file_system->AtomicStore(schema_path, std::string(buffer.GetString())); } + Status CompactAndCommit(const std::string& table_path, + const std::map& options, + int64_t commit_identifier) const { + WriteContextBuilder write_context_builder(table_path, "commit_user"); + write_context_builder.WithTempDirectory(LookupTempDirectory()); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr write_context, + write_context_builder.SetOptions(options).WithStreamingMode(true).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_store_write, + FileStoreWrite::Create(std::move(write_context))); + PAIMON_RETURN_NOT_OK(file_store_write->Compact(/*partition=*/{}, /*bucket=*/0, + /*full_compaction=*/true)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> compact_messages, + file_store_write->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + PAIMON_RETURN_NOT_OK(file_store_write->Close()); + if (compact_messages.empty()) { + return Status::Invalid("expected compaction commit messages"); + } + CommitContextBuilder commit_context_builder(table_path, "commit_user"); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit_context, + commit_context_builder.SetOptions(options).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_store_commit, + FileStoreCommit::Create(std::move(commit_context))); + return file_store_commit->Commit(compact_messages, commit_identifier); + } + Result ReadAndCheckProjectedResult(const std::map& options, const std::vector& read_fields, const std::shared_ptr& expected_type, @@ -309,6 +355,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()), @@ -370,6 +676,594 @@ TEST_P(WriteAndReadInteTest, TestPKSimple) { ASSERT_TRUE(success); } +TEST_P(WriteAndReadInteTest, TestInputChangelogStreamRead) { + arrow::FieldVector fields = { + arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + }; + auto [file_format, file_system] = GetParam(); + 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}, {Options::CHANGELOG_PRODUCER, "input"}, + }; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN( + auto helper, + TestHelper::Create(test_dir_, arrow::schema(fields), /*partition_keys=*/{}, + /*primary_keys=*/{"pk"}, options, /*is_streaming_mode=*/true)); + + ASSERT_OK_AND_ASSIGN(std::vector> initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), R"([["Alice", 10], ["Bob", 20], ["Alice", 11], ["Bob", 21]])", + /*partition_map=*/{}, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_BEFORE, + RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector> changelog_splits, helper->Scan()); + ASSERT_TRUE(changelog_splits.empty()); + ASSERT_OK_AND_ASSIGN(changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + fields[0], + fields[1], + }); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[0, "Alice", 10], [2, "Alice", 11], + [1, "Bob", 20], [3, "Bob", 21]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogStreamRead) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = { + arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + }; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + // Move the initial value to a high level so the next compaction must look it up. + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + + ASSERT_OK_AND_ASSIGN(std::vector> initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(std::vector> changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + fields[0], + fields[1], + }); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 20]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogInitialFullScanExcludesLevelZero) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = { + arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + }; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector> initial_full_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_FALSE(initial_full_splits.empty()); + for (const auto& split : initial_full_splits) { + auto data_split = std::dynamic_pointer_cast(split); + ASSERT_TRUE(data_split); + for (const auto& file : data_split->DataFiles()) { + ASSERT_GT(file->level, 0); + } + } + + auto expected_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + fields[0], + fields[1], + }); + ASSERT_OK_AND_ASSIGN( + bool success, + helper->ReadAndCheckResult(expected_type, initial_full_splits, R"([[0, "Alice", 10]])")); + ASSERT_TRUE(success); + + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(std::vector> changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + ASSERT_OK_AND_ASSIGN(bool changelog_success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 20]])")); + ASSERT_TRUE(changelog_success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogInsertUpdateDelete) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN( + auto initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10], ["Bob", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN( + auto change_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([["Alice", 11], ["Bob", 0], ["Carol", 30], ["Dave", 0]])", + /*partition_map=*/{}, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE})); + ASSERT_OK(helper->WriteAndCommit(std::move(change_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 11], + [3, "Bob", 20], [0, "Carol", 30]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogWithFirstRow) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::MERGE_ENGINE, "first-row"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN( + auto change_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20], ["Bob", 30]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(change_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[0, "Bob", 30]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogWithDeletionVector) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::DELETION_VECTORS_ENABLED, "true"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN(auto update_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + ASSERT_OK_AND_ASSIGN(bool changelog_success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 20]])")); + ASSERT_TRUE(changelog_success); + + ASSERT_OK_AND_ASSIGN(auto batch_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt, + /*is_streaming=*/false)); + ASSERT_OK_AND_ASSIGN(bool batch_success, helper->ReadAndCheckResult(expected_type, batch_splits, + R"([[0, "Alice", 20]])")); + ASSERT_TRUE(batch_success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogRowDeduplicate) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, + {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, "true"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN(auto write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto unchanged_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(unchanged_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + ASSERT_OK_AND_ASSIGN(auto empty_splits, helper->Scan()); + ASSERT_TRUE(empty_splits.empty()); + + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto changed_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(changed_batch), /*commit_identifier=*/4, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/5)); + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 20]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogRowDeduplicateIgnoreFields) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + arrow::field("ignored", arrow::int32())}; + 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}, + {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE, "true"}, + {Options::CHANGELOG_PRODUCER_ROW_DEDUPLICATE_IGNORE_FIELDS, "ignored"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, TestHelper::MakeRecordBatch( + arrow::struct_(fields), R"([["Alice", 10, 100]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN(auto write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN( + auto ignored_change_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10, 200]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(ignored_change_batch), + /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + ASSERT_OK_AND_ASSIGN(auto empty_splits, helper->Scan()); + ASSERT_TRUE(empty_splits.empty()); + + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN( + auto real_change_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20, 300]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(real_change_batch), /*commit_identifier=*/4, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/5)); + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + auto expected_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1], fields[2]}); + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult( + expected_type, changelog_splits, + R"([[1, "Alice", 10, 200], [2, "Alice", 20, 300]])")); + ASSERT_TRUE(success); +} + +TEST_P(WriteAndReadInteTest, TestChangelogWithSchemaEvolution) { + auto [file_format, file_system] = GetParam(); + arrow::FieldVector fields_v0 = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, {Options::CHANGELOG_PRODUCER, "lookup"}}; + if (file_system == "jindo") { + options = AddOptionsForJindo(options); + } + ASSERT_OK_AND_ASSIGN(auto write_helper, + CreateLookupTestHelper(arrow::schema(fields_v0), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, TestHelper::MakeRecordBatch( + arrow::struct_(fields_v0), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(auto scan_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + scan_helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN( + auto batch_v0, TestHelper::MakeRecordBatch(arrow::struct_(fields_v0), R"([["Alice", 11]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(batch_v0), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + arrow::FieldVector fields_v1 = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + arrow::field("extra", arrow::utf8())}; + ASSERT_OK(WriteNextSchema( + {DataField(0, fields_v1[0]), DataField(1, fields_v1[1]), DataField(2, fields_v1[2])}, + /*highest_field_id=*/2, options)); + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto batch_v1, TestHelper::MakeRecordBatch( + arrow::struct_(fields_v1), R"([["Alice", 12, "v1"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(batch_v1), /*commit_identifier=*/4, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/5)); + + arrow::FieldVector fields_v2 = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + arrow::field("renamed_extra", arrow::utf8())}; + ASSERT_OK(WriteNextSchema( + {DataField(0, fields_v2[0]), DataField(1, fields_v2[1]), DataField(2, fields_v2[2])}, + /*highest_field_id=*/2, options)); + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto batch_v2, TestHelper::MakeRecordBatch( + arrow::struct_(fields_v2), R"([["Alice", 13, "v2"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(batch_v2), /*commit_identifier=*/6, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/7)); + + arrow::FieldVector fields_v3 = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int64()), + arrow::field("renamed_extra", arrow::utf8())}; + ASSERT_OK(WriteNextSchema( + {DataField(0, fields_v3[0]), DataField(1, fields_v3[1]), DataField(2, fields_v3[2])}, + /*highest_field_id=*/2, options)); + write_helper.reset(); + ASSERT_OK_AND_ASSIGN(write_helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto batch_v3, TestHelper::MakeRecordBatch( + arrow::struct_(fields_v3), R"([["Alice", 14, "v3"]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(write_helper->WriteAndCommit(std::move(batch_v3), /*commit_identifier=*/8, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/9)); + + auto expected_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), fields_v3[0], fields_v3[1], fields_v3[2]}); + std::vector expected_data = { + R"([[1, "Alice", 10, null], [2, "Alice", 11, null]])", + R"([[1, "Alice", 11, null], [2, "Alice", 12, "v1"]])", + R"([[1, "Alice", 12, "v1"], [2, "Alice", 13, "v2"]])", + R"([[1, "Alice", 13, "v2"], [2, "Alice", 14, "v3"]])"}; + for (int64_t schema_id = 0; schema_id < static_cast(expected_data.size()); + schema_id++) { + ASSERT_OK_AND_ASSIGN(auto changelog_splits, scan_helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + for (const auto& split : changelog_splits) { + auto data_split = std::dynamic_pointer_cast(split); + ASSERT_TRUE(data_split); + for (const auto& file : data_split->DataFiles()) { + ASSERT_EQ(file->schema_id, schema_id); + } + } + ASSERT_OK_AND_ASSIGN(bool success, scan_helper->ReadAndCheckResult( + expected_type, changelog_splits, + expected_data[static_cast(schema_id)])); + ASSERT_TRUE(success); + } +} + +TEST_P(WriteAndReadInteTest, TestLookupChangelogWithExternalPath) { + auto [file_format, file_system] = GetParam(); + if (file_system == "jindo") { + return; + } + std::unique_ptr external_dir = UniqueTestDirectory::Create(file_system); + ASSERT_TRUE(external_dir); + arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32())}; + 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}, + {Options::CHANGELOG_PRODUCER, "lookup"}, + {Options::DATA_FILE_EXTERNAL_PATHS, "FILE://" + external_dir->Str()}, + {Options::DATA_FILE_EXTERNAL_PATHS_STRATEGY, "round-robin"}}; + ASSERT_OK_AND_ASSIGN(auto helper, CreateLookupTestHelper(arrow::schema(fields), options)); + ASSERT_OK_AND_ASSIGN(auto initial_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 10]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(initial_batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + std::string table_path = PathUtil::JoinPath(test_dir_, "foo.db/bar"); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/1)); + + helper.reset(); + ASSERT_OK_AND_ASSIGN(helper, CreateLookupTestHelper(table_path, options)); + ASSERT_OK_AND_ASSIGN(auto initial_splits, + helper->NewScan(StartupMode::Latest(), /*snapshot_id=*/std::nullopt)); + ASSERT_TRUE(initial_splits.empty()); + ASSERT_OK_AND_ASSIGN(auto update_batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([["Alice", 20]])", + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(update_batch), /*commit_identifier=*/2, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK(CompactAndCommit(table_path, options, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(auto changelog_splits, helper->Scan()); + ASSERT_FALSE(changelog_splits.empty()); + bool found_external_changelog = false; + for (const auto& split : changelog_splits) { + auto data_split = std::dynamic_pointer_cast(split); + ASSERT_TRUE(data_split); + for (const auto& file : data_split->DataFiles()) { + ASSERT_TRUE(file->external_path.has_value()); + ASSERT_TRUE(StringUtils::StartsWith(PathUtil::GetName(file->file_name), "changelog-")); + ASSERT_OK_AND_ASSIGN( + bool exists, external_dir->GetFileSystem()->Exists(file->external_path.value())); + ASSERT_TRUE(exists); + found_external_changelog = true; + } + } + ASSERT_TRUE(found_external_changelog); + + auto expected_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(expected_type, changelog_splits, + R"([[1, "Alice", 10], [2, "Alice", 20]])")); + ASSERT_TRUE(success); +} + TEST_P(WriteAndReadInteTest, TestNestedType) { arrow::FieldVector fields = { arrow::field("f1", arrow::map(arrow::int8(), arrow::int16())), @@ -381,6 +1275,9 @@ TEST_P(WriteAndReadInteTest, TestNestedType) { arrow::field("f6", arrow::decimal128(2, 2))}; auto schema = arrow::schema(fields); auto [file_format, file_system] = GetParam(); + if (file_format == "mosaic") { + return; + } std::map options = { {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "-1"}, @@ -430,7 +1327,7 @@ TEST_P(WriteAndReadInteTest, TestNestedType) { TEST_P(WriteAndReadInteTest, TestSchemaEvolutionAddFieldInsideListAndMap) { auto [file_format, file_system] = GetParam(); - if (file_format == "lance" || file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } auto list_struct = @@ -716,6 +1613,9 @@ TEST_P(WriteAndReadInteTest, TestAppendTimestampType) { }; auto schema = arrow::schema(fields); auto [file_format, file_system] = GetParam(); + if (file_format == "mosaic") { + return; + } std::map options = { {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "-1"}, @@ -769,6 +1669,9 @@ TEST_P(WriteAndReadInteTest, TestPkTimestampType) { }; auto schema = arrow::schema(fields); auto [file_format, file_system] = GetParam(); + if (file_format == "mosaic") { + return; + } std::map options = { {Options::MANIFEST_FORMAT, "avro"}, {Options::FILE_FORMAT, file_format}, {Options::TARGET_FILE_SIZE, "1024"}, {Options::BUCKET, "1"}, @@ -815,6 +1718,9 @@ TEST_P(WriteAndReadInteTest, TestPkTimestampType) { /// the reader has to convert milli back to second for every nested leaf. TEST_P(WriteAndReadInteTest, TestAppendNestedTimestampSecondPrecision) { auto [file_format, file_system] = GetParam(); + if (file_format == "mosaic") { + return; + } TimezoneGuard timezone_guard("Asia/Shanghai"); auto timezone = DateTimeUtils::GetLocalTimezoneName(); auto event_type = arrow::struct_({ @@ -870,6 +1776,9 @@ TEST_P(WriteAndReadInteTest, TestAppendNestedTimestampSecondPrecision) { /// file schema differ only in the timezone of those leaves; the micro precision stays unchanged. TEST_P(WriteAndReadInteTest, TestAppendNestedTimestampLtzMicroTimezoneOnly) { auto [file_format, file_system] = GetParam(); + if (file_format == "mosaic") { + return; + } // Pin a non-UTC timezone so the read schema really differs from what the file reports. TimezoneGuard timezone_guard("Asia/Shanghai"); auto timezone = DateTimeUtils::GetLocalTimezoneName(); @@ -1298,6 +2207,9 @@ TEST_P(WriteAndReadInteTest, TestCharVarcharBinaryVarbinaryTypes) { std::vector> GetTestValuesForWriteAndReadInteTest() { std::vector> values = {{"parquet", "local"}}; +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic", "local"); +#endif #if defined(PAIMON_ENABLE_NETWORK_TESTS) && defined(PAIMON_ENABLE_JINDO) values.emplace_back("parquet", "jindo"); #endif @@ -1714,7 +2626,7 @@ TEST_P(WriteAndReadInteTest, TestAppendWithParquetMetadataCache) { TEST_P(WriteAndReadInteTest, TestAppendSharedShreddingMap) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -1771,7 +2683,7 @@ TEST_P(WriteAndReadInteTest, TestAppendSharedShreddingMap) { TEST_P(WriteAndReadInteTest, TestMapSharedShreddingColumnPlacementPolicies) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -1855,7 +2767,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingColumnPlacementPolicies) { TEST_P(WriteAndReadInteTest, TestAppendMapSharedShreddingWithPartitionAndBucket) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -1963,7 +2875,7 @@ TEST_P(WriteAndReadInteTest, TestAppendMapSharedShreddingWithPartitionAndBucket) TEST_P(WriteAndReadInteTest, TestAppendMapSharedShreddingWithPredicate) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2057,7 +2969,7 @@ TEST_P(WriteAndReadInteTest, TestAppendMapSharedShreddingWithPredicate) { TEST_P(WriteAndReadInteTest, TestMapSharedShreddingNewWriterStartsWithMaxColumnCount) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2126,7 +3038,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingNewWriterStartsWithMaxColumnC TEST_P(WriteAndReadInteTest, TestMapSharedShreddingAdaptsAcrossRollingFiles) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2199,7 +3111,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingAdaptsAcrossRollingFiles) { TEST_P(WriteAndReadInteTest, TestMapSharedShreddingSwitchMapLayoutAndUseMaxColumnsWithoutMetadata) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2282,7 +3194,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingSwitchMapLayoutAndUseMaxColum TEST_P(WriteAndReadInteTest, TestMapSharedShreddingReadAfterRenameColumn) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2366,7 +3278,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingReadAfterRenameColumn) { TEST_P(WriteAndReadInteTest, TestSharedShreddingWithSchemaEvolution) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2477,7 +3389,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingWithSchemaEvolution) { // Verify storage-layout evolution: default->shared-shredding. TEST_P(WriteAndReadInteTest, TestMapStorageLayoutDefaultToSharedShredding) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2555,7 +3467,7 @@ TEST_P(WriteAndReadInteTest, TestMapStorageLayoutDefaultToSharedShredding) { // Verify storage-layout evolution: shared-shredding->default. TEST_P(WriteAndReadInteTest, TestMapStorageLayoutSharedShreddingToDefault) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2619,7 +3531,7 @@ TEST_P(WriteAndReadInteTest, TestMapStorageLayoutSharedShreddingToDefault) { TEST_P(WriteAndReadInteTest, TestAppendMapStorageLayoutSharedShreddingToDefaultCompaction) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2685,23 +3597,7 @@ TEST_P(WriteAndReadInteTest, TestAppendMapStorageLayoutSharedShreddingToDefaultC ASSERT_OK(helper->WriteAndCommit(std::move(batch_v1_file3), commit_identifier++, /*expected_commit_messages=*/std::nullopt)); - WriteContextBuilder write_context_builder(table_path, "commit_user"); - ASSERT_OK_AND_ASSIGN( - auto write_context, - write_context_builder.SetOptions(options_v1).WithStreamingMode(true).Finish()); - ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); - ASSERT_OK(file_store_write->Compact(/*partition=*/{}, /*bucket=*/0, - /*full_compaction=*/true)); - ASSERT_OK_AND_ASSIGN(auto compact_messages, file_store_write->PrepareCommit( - /*wait_compaction=*/true, commit_identifier)); - ASSERT_FALSE(compact_messages.empty()); - - CommitContextBuilder commit_context_builder(table_path, "commit_user"); - ASSERT_OK_AND_ASSIGN(auto commit_context, - commit_context_builder.SetOptions(options_v1).Finish()); - ASSERT_OK_AND_ASSIGN(auto file_store_commit, - FileStoreCommit::Create(std::move(commit_context))); - ASSERT_OK(file_store_commit->Commit(compact_messages, commit_identifier)); + ASSERT_OK(CompactAndCommit(table_path, options_v1, commit_identifier)); arrow::FieldVector expected_fields = fields; expected_fields.insert(expected_fields.begin(), arrow::field("_VALUE_KIND", arrow::int8())); @@ -2723,7 +3619,7 @@ TEST_P(WriteAndReadInteTest, TestAppendMapStorageLayoutSharedShreddingToDefaultC // Nested map values through both selected physical columns and overflow. TEST_P(WriteAndReadInteTest, TestSharedShreddingWithStructValue) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2794,7 +3690,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingWithStructValue) { TEST_P(WriteAndReadInteTest, TestMapSharedShreddingWithComplexValue) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -2908,7 +3804,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingWithComplexValue) { TEST_P(WriteAndReadInteTest, TestMapSharedShreddingWithAllSupportedComplexValueTypes) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3058,7 +3954,7 @@ TEST_P(WriteAndReadInteTest, TestMapSharedShreddingWithAllSupportedComplexValueT TEST_P(WriteAndReadInteTest, TestMapSharedShreddingStructValueSchemaEvolutionReadFails) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3270,7 +4166,7 @@ TEST_P(WriteAndReadInteTest, TestOrcDictionaryLazyDecodingWithSharedShredding) { // Verify shared-shredding in the PK read path. TEST_P(WriteAndReadInteTest, TestPkSharedShreddingMap) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3340,7 +4236,7 @@ TEST_P(WriteAndReadInteTest, TestPkSharedShreddingMap) { TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallWithOverflow) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3475,7 +4371,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallWithOverflow) { TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallWithNullOrMissingKey) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3586,7 +4482,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallWithNullOrMissin TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallMultipleColumns) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3698,7 +4594,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingPartialKeyRecallMultipleColumns) TEST_P(WriteAndReadInteTest, TestMapStorageLayoutDefaultToSharedShreddingPartialKeyRecall) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3788,7 +4684,7 @@ TEST_P(WriteAndReadInteTest, TestMapStorageLayoutDefaultToSharedShreddingPartial TEST_P(WriteAndReadInteTest, TestMapStorageLayoutSharedShreddingToDefaultPartialKeyRecall) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3883,7 +4779,7 @@ TEST_P(WriteAndReadInteTest, TestMapStorageLayoutSharedShreddingToDefaultPartial TEST_P(WriteAndReadInteTest, TestSharedShreddingDuplicateSelectedKeys) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -3931,7 +4827,7 @@ TEST_P(WriteAndReadInteTest, TestSharedShreddingDuplicateSelectedKeys) { TEST_P(WriteAndReadInteTest, TestSharedShreddingAllNullMapColumn) { auto [file_format, file_system] = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } diff --git a/test/inte/write_inte_test.cpp b/test/inte/write_inte_test.cpp index 451c3ce29..5b0b8a025 100644 --- a/test/inte/write_inte_test.cpp +++ b/test/inte/write_inte_test.cpp @@ -197,10 +197,10 @@ class WriteInteTest : public testing::Test, public ::testing::WithParamInterface std::shared_ptr ReconstructDataFileMeta( const std::shared_ptr& file_meta) const { - if (GetParam() != "avro") { + if (GetParam() != "avro" && GetParam() != "mosaic") { return file_meta; } - // For the avro format, all stats are null. + // Avro and Mosaic without configured statistics have null statistics. auto new_meta = std::make_shared( file_meta->file_name, file_meta->file_size, file_meta->row_count, file_meta->min_key, file_meta->max_key, file_meta->key_stats, file_meta->value_stats, @@ -208,7 +208,7 @@ class WriteInteTest : public testing::Test, public ::testing::WithParamInterface file_meta->level, file_meta->extra_files, file_meta->creation_time, file_meta->delete_row_count, file_meta->embedded_index, file_meta->file_source, file_meta->value_stats_cols, file_meta->external_path, file_meta->first_row_id, - file_meta->write_cols); + file_meta->write_cols, /*column_max_sequence_numbers=*/std::nullopt); auto generate_null_stats = [this](const SimpleStats& stats) -> SimpleStats { if (stats == SimpleStats::EmptyStats()) { return stats; @@ -383,6 +383,9 @@ class WriteInteTest : public testing::Test, public ::testing::WithParamInterface std::vector GetTestValuesForWriteInteTest() { std::vector values; values.emplace_back("parquet"); +#ifdef PAIMON_ENABLE_MOSAIC + values.emplace_back("mosaic"); +#endif #ifdef PAIMON_ENABLE_ORC values.emplace_back("orc"); #endif @@ -445,7 +448,7 @@ TEST_P(WriteInteTest, TestAppendTableBatchWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -549,7 +552,7 @@ TEST_P(WriteInteTest, TestAppendTableStreamWriteWithOneBucket) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -619,7 +622,7 @@ TEST_P(WriteInteTest, TestAppendTableStreamWriteWithOneBucket) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -805,6 +808,9 @@ TEST_P(WriteInteTest, TestAppendTableStreamWriteWithPartitionAndMultiBuckets) { } TEST_P(WriteInteTest, TestAppendTableWriteWithComplexType) { + if (GetParam() == "mosaic") { + return; + } auto dir = UniqueTestDirectory::Create(); arrow::FieldVector fields = { arrow::field("f1", arrow::map(arrow::int8(), arrow::int16())), @@ -862,7 +868,7 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithComplexType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -926,7 +932,7 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithComplexType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -1011,7 +1017,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_1 = ReconstructDataFileMeta(file_meta_1); DataIncrement data_increment_1({file_meta_1}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -1037,7 +1043,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -1063,7 +1069,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_3 = ReconstructDataFileMeta(file_meta_3); DataIncrement data_increment_3({file_meta_3}, {}, {}); std::shared_ptr expected_commit_message_3 = std::make_shared( @@ -1135,7 +1141,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_4 = ReconstructDataFileMeta(file_meta_4); DataIncrement data_increment_4({file_meta_4}, {}, {}); std::shared_ptr expected_commit_message_4 = std::make_shared( @@ -1161,7 +1167,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_5 = ReconstructDataFileMeta(file_meta_5); DataIncrement data_increment_5({file_meta_5}, {}, {}); std::shared_ptr expected_commit_message_5 = std::make_shared( @@ -1187,7 +1193,7 @@ TEST_P(WriteInteTest, TestPkTableStreamWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_6 = ReconstructDataFileMeta(file_meta_6); DataIncrement data_increment_6({file_meta_6}, {}, {}); std::shared_ptr expected_commit_message_6 = std::make_shared( @@ -1285,7 +1291,7 @@ TEST_P(WriteInteTest, TestPkTableBatchWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_1 = ReconstructDataFileMeta(file_meta_1); DataIncrement data_increment_1({file_meta_1}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -1311,7 +1317,7 @@ TEST_P(WriteInteTest, TestPkTableBatchWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -1337,7 +1343,7 @@ TEST_P(WriteInteTest, TestPkTableBatchWrite) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_3 = ReconstructDataFileMeta(file_meta_3); DataIncrement data_increment_3({file_meta_3}, {}, {}); std::shared_ptr expected_commit_message_3 = std::make_shared( @@ -1452,7 +1458,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithNoPartitionKey) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_1 = ReconstructDataFileMeta(file_meta_1); DataIncrement data_increment_1({file_meta_1}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -1482,7 +1488,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithNoPartitionKey) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -1558,7 +1564,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithNoPartitionKey) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_3 = ReconstructDataFileMeta(file_meta_3); DataIncrement data_increment_3({file_meta_3}, {}, {}); std::shared_ptr expected_commit_message_3 = std::make_shared( @@ -1586,7 +1592,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithNoPartitionKey) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_4 = ReconstructDataFileMeta(file_meta_4); DataIncrement data_increment_4({file_meta_4}, {}, {}); std::shared_ptr expected_commit_message_4 = std::make_shared( @@ -1626,6 +1632,9 @@ TEST_P(WriteInteTest, TestPkTableWriteWithNoPartitionKey) { } TEST_P(WriteInteTest, TestPkTableWriteWithComplexType) { + if (GetParam() == "mosaic") { + return; + } auto dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); arrow::FieldVector fields = { @@ -1696,7 +1705,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithComplexType) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -1769,7 +1778,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithComplexType) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -2036,7 +2045,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_1 = ReconstructDataFileMeta(file_meta_1); DataIncrement data_increment_1({file_meta_1}, {}, {}); std::shared_ptr expected_commit_message_1 = @@ -2065,7 +2074,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = @@ -2094,7 +2103,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/2, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_3 = ReconstructDataFileMeta(file_meta_3); DataIncrement data_increment_3({file_meta_3}, {}, {}); std::shared_ptr expected_commit_message_3 = @@ -2126,7 +2135,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_4 = ReconstructDataFileMeta(file_meta_4); DataIncrement data_increment_4({file_meta_4}, {}, {}); std::shared_ptr expected_commit_message_4 = @@ -2154,7 +2163,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_5 = ReconstructDataFileMeta(file_meta_5); DataIncrement data_increment_5({file_meta_5}, {}, {}); std::shared_ptr expected_commit_message_5 = @@ -2183,7 +2192,7 @@ TEST_P(WriteInteTest, TestPkTableWriteWithIOException) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta_6 = ReconstructDataFileMeta(file_meta_6); DataIncrement data_increment_6({file_meta_6}, {}, {}); std::shared_ptr expected_commit_message_6 = @@ -2257,7 +2266,7 @@ TEST_F(WriteInteTest, TestAppendTableWriteWithAlterTable) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message = std::make_shared( BinaryRowGenerator::GenerateRow({1, 1}, pool_.get()), /*bucket=*/0, @@ -2337,7 +2346,7 @@ TEST_F(WriteInteTest, TestPKTableWriteWithAlterTable) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message = std::make_shared( BinaryRowGenerator::GenerateRow({0, 0}, pool_.get()), /*bucket=*/0, @@ -2473,7 +2482,7 @@ TEST_P(WriteInteTest, TestWriteAndCommitIOException) { TEST_P(WriteInteTest, TestWriteWithFieldId) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } // prepare write schema and write data @@ -2775,7 +2784,8 @@ TEST_P(WriteInteTest, TestAppendTableStreamWriteWithExternalPath) { /*level=*/0, /*extra_files=*/std::vector>(), /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), - /*value_stats_cols=*/std::nullopt, "FILE:/tmp/xxx", std::nullopt, std::nullopt); + /*value_stats_cols=*/std::nullopt, "FILE:/tmp/xxx", std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -2844,7 +2854,8 @@ TEST_P(WriteInteTest, TestAppendTableStreamWriteWithExternalPath) { /*level=*/0, /*extra_files=*/std::vector>(), /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), - /*value_stats_cols=*/std::nullopt, "FILE:/tmp/xxx", std::nullopt, std::nullopt); + /*value_stats_cols=*/std::nullopt, "FILE:/tmp/xxx", std::nullopt, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); file_meta_2 = ReconstructDataFileMeta(file_meta_2); DataIncrement data_increment_2({file_meta_2}, {}, {}); std::shared_ptr expected_commit_message_2 = std::make_shared( @@ -3074,6 +3085,9 @@ TEST_P(WriteInteTest, TestWriteAndReadWithSpecialPartitionValue) { } TEST_P(WriteInteTest, TestWriteWithNestedSchema) { + if (GetParam() == "mosaic") { + return; + } arrow::FieldVector fields = { arrow::field("f0", arrow::struct_({arrow::field("v0", arrow::boolean()), arrow::field("v1", arrow::int64())}))}; @@ -3335,6 +3349,9 @@ TEST_P(WriteInteTest, TestWriteMemoryUse) { } TEST_P(WriteInteTest, TestAppendTableWithAllNull) { + if (GetParam() == "mosaic") { + return; + } auto dir = UniqueTestDirectory::Create(); arrow::FieldVector fields = { arrow::field("f0", arrow::boolean()), @@ -3443,7 +3460,7 @@ TEST_P(WriteInteTest, TestPkTablePostponeBucket) { /*delete_row_count=*/1, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message = std::make_shared( /*partition_map=*/BinaryRow::EmptyRow(), /*bucket=*/-2, @@ -3532,7 +3549,7 @@ TEST_F(WriteInteTest, TestBranchWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment_0({file_meta_0}, {}, {}); auto expected_commit_message_0 = std::make_shared( BinaryRowGenerator::GenerateRow({std::string("20240726")}, pool_.get()), /*bucket=*/0, @@ -3574,7 +3591,7 @@ TEST_F(WriteInteTest, TestBranchWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); DataIncrement data_increment_1({file_meta_1}, {}, {}); auto expected_commit_message_1 = std::make_shared( BinaryRowGenerator::GenerateRow({std::string("20240725")}, pool_.get()), /*bucket=*/0, @@ -3704,7 +3721,7 @@ TEST_P(WriteInteTest, TestDataEvolutionWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); check_meta(commit_msgs1[0], {file_meta1}); commit(commit_msgs1, /*latest_snapshot_id=*/1, /*next_row_id=*/2); check_committed_meta({{0, std::nullopt, 1, 1}}); @@ -3736,7 +3753,7 @@ TEST_P(WriteInteTest, TestDataEvolutionWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/write_cols2); + /*write_cols=*/write_cols2, /*column_max_sequence_numbers=*/std::nullopt); check_meta(commit_msgs2[0], {file_meta2}); commit(commit_msgs2, /*latest_snapshot_id=*/2, /*next_row_id=*/7); check_committed_meta({{0, std::nullopt, 1, 1}, {2, write_cols2, 2, 2}}); @@ -3762,7 +3779,7 @@ TEST_P(WriteInteTest, TestDataEvolutionWrite) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/write_cols3); + /*write_cols=*/write_cols3, /*column_max_sequence_numbers=*/std::nullopt); check_meta(commit_msgs3[0], {file_meta3}); auto commit_msg_impl = std::dynamic_pointer_cast(commit_msgs3[0]); ASSERT_TRUE(commit_msg_impl); @@ -3834,7 +3851,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithBlobType) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"f0", "f1"})); + /*write_cols=*/std::vector({"f0", "f1"}), + /*column_max_sequence_numbers=*/std::nullopt); file_meta1 = ReconstructDataFileMeta(file_meta1); auto file_meta2 = std::make_shared( "data-xxx.blob", /*file_size=*/764, /*row_count=*/3, @@ -3845,7 +3863,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithBlobType) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"blob"})); + /*write_cols=*/std::vector({"blob"}), + /*column_max_sequence_numbers=*/std::nullopt); auto file_meta3 = std::make_shared( "data-xxx.blob", /*file_size=*/3023, /*row_count=*/1, /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), @@ -3855,7 +3874,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithBlobType) { /*creation_time=*/Timestamp(1724090888706ll, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/3, - /*write_cols=*/std::vector({"blob"})); + /*write_cols=*/std::vector({"blob"}), + /*column_max_sequence_numbers=*/std::nullopt); std::vector> expected_meta = {file_meta1, file_meta2, file_meta3}; // NOTE: Due to the write logic of C++ Paimon and Java Paimon is different, the first_row_id in @@ -3924,7 +3944,7 @@ TEST_P(WriteInteTest, TestAppendTableWithDateFieldAsPartitionField) { /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, - /*write_cols=*/std::nullopt); + /*write_cols=*/std::nullopt, /*column_max_sequence_numbers=*/std::nullopt); file_meta = ReconstructDataFileMeta(file_meta); DataIncrement data_increment({file_meta}, {}, {}); std::shared_ptr expected_commit_message_1 = std::make_shared( @@ -4006,7 +4026,7 @@ TEST_P(WriteInteTest, TestNullabilityCheck) { TEST_P(WriteInteTest, TestPkSpillableMapSharedShreddingReadWrite) { auto file_format = GetParam(); - if (file_format == "avro") { + if (file_format == "avro" || file_format == "mosaic") { return; } @@ -4728,7 +4748,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithMultipleBlobFields) { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"f0", "f1"})); + /*write_cols=*/std::vector({"f0", "f1"}), + /*column_max_sequence_numbers=*/std::nullopt); expected_main = ReconstructDataFileMeta(expected_main); // blob1 file: 3 rows, write_cols={"blob1"}, first_row_id=0 @@ -4741,7 +4762,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithMultipleBlobFields) { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"blob1"})); + /*write_cols=*/std::vector({"blob1"}), + /*column_max_sequence_numbers=*/std::nullopt); // blob2 file: 3 rows, write_cols={"blob2"}, first_row_id=0 auto expected_blob2 = std::make_shared( @@ -4753,7 +4775,8 @@ TEST_P(WriteInteTest, TestAppendTableWriteWithMultipleBlobFields) { /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/0, /*embedded_index=*/nullptr, FileSource::Append(), /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, /*first_row_id=*/0, - /*write_cols=*/std::vector({"blob2"})); + /*write_cols=*/std::vector({"blob2"}), + /*column_max_sequence_numbers=*/std::nullopt); ASSERT_OK_AND_ASSIGN(auto commit_msgs, helper->WriteAndCommit(std::move(batch), commit_identifier++, diff --git a/test/test_data/compatibility/commit_message-v13 b/test/test_data/compatibility/commit_message-v13 new file mode 100644 index 000000000..568705a59 Binary files /dev/null and b/test/test_data/compatibility/commit_message-v13 differ diff --git a/test/test_data/compatibility/data_split-v9 b/test/test_data/compatibility/data_split-v9 new file mode 100644 index 000000000..2bfd4ac19 Binary files /dev/null and b/test/test_data/compatibility/data_split-v9 differ diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/README.md b/test/test_data/mosaic/append_java_compat.db/append_java_compat/README.md new file mode 100644 index 000000000..92c7d2089 --- /dev/null +++ b/test/test_data/mosaic/append_java_compat.db/append_java_compat/README.md @@ -0,0 +1,60 @@ +Table: append_java_compat +Writer: Paimon Java at commit 0043a70fd88ac75dcb83a8f2da5e72ce91e22b1f +Mosaic version: 0.2.0 + +Schema: +id INT NOT NULL +f_boolean BOOLEAN +f_tinyint TINYINT +f_smallint SMALLINT +f_bigint BIGINT +f_float FLOAT +f_double DOUBLE +f_char CHAR(8) +f_varchar VARCHAR(64) +f_binary BINARY(8) +f_varbinary VARBINARY(64) +f_date DATE +f_ts_3 TIMESTAMP(3) +f_ts_6 TIMESTAMP(6) +f_ts_9 TIMESTAMP(9) +f_ltz_3 TIMESTAMP(3) WITH LOCAL TIME ZONE +f_ltz_6 TIMESTAMP(6) WITH LOCAL TIME ZONE +f_ltz_9 TIMESTAMP(9) WITH LOCAL TIME ZONE +f_decimal_1_0 DECIMAL(1, 0) +f_decimal_18_2 DECIMAL(18, 2) +f_decimal_19_2 DECIMAL(19, 2) +f_decimal_38_18 DECIMAL(38, 18) +f_array_int ARRAY +f_map_numeric MAP +f_map_string_bigint MAP +f_array_array_int ARRAY> +f_array_map ARRAY> +f_map_array MAP> + +Options: +bucket = -1 +file.block-size = 1 B +file.format = mosaic +manifest.format = avro +mosaic.num-buckets = 4 +mosaic.stats-columns = id,f_varchar,f_date,f_ts_3,f_ltz_9,f_decimal_18_2 +target-file-size = 64 MB +write.batch-size = 2 + +Data is written as three row groups with two rows in each row group. Timestamp values are shown in +UTC. The second row has NULL in every nullable field. + +Row group 0: +Add: (1, true, -5, -1000, 10000000001, 1.25, 10.5, "char0001", "value-1", "bin00001", 0x010203, 2024-10-04, 1970-01-01 00:00:01.123, 1970-01-01 00:00:01.123456, 1970-01-01 00:00:01.123456789, 1970-01-01 00:01:01.321, 1970-01-01 00:01:01.654321, 1970-01-01 00:01:01.321654987, -3, 1.25, 12345678901234567.89, 12345678901234567890.123456789012345678, [1, 2], {0: 0, 10: 1}, {"k0": 1000, "z0": 2000}, [[1, 2], [3]], [{"nested0": 0}, {"nested10": 10}], {"array0": [0, 1]}) +Add: (2, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +Row group 1: +Add: (10, true, -3, -998, 10000000010, 3.25, 12.5, "char0010", "value-10", "bin00010", 0x0a0b0c, 2024-10-06, 1970-01-01 00:00:10.123, 1970-01-01 00:00:10.123456, 1970-01-01 00:00:10.123456789, 1970-01-01 00:01:10.321, 1970-01-01 00:01:10.654321, 1970-01-01 00:01:10.321654987, -1, 10.25, 12345678901234569.89, 12345678901234567892.123456789012345678, [10, 11], {2: 20, 12: 21}, {"k2": 1002, "z2": 2002}, [[10, 11], [12]], [{"nested2": 2}, {"nested12": 12}], {"array2": [2, 3]}) +Add: (11, false, -2, -997, 10000000011, 4.25, 13.5, "char0011", "value-11", "bin00011", 0x0b0c0d, 2024-10-07, 1970-01-01 00:00:11.123, 1970-01-01 00:00:11.123456, 1970-01-01 00:00:11.123456789, 1970-01-01 00:01:11.321, 1970-01-01 00:01:11.654321, 1970-01-01 00:01:11.321654987, 0, 11.25, 12345678901234570.89, 12345678901234567893.123456789012345678, [11, 12], {3: 30, 13: 31}, {"k3": 1003, "z3": 2003}, [[11, 12], [13]], [{"nested3": 3}, {"nested13": 13}], {"array3": [3, 4]}) + +Row group 2: +Add: (20, true, -1, -996, 10000000020, 5.25, 14.5, "char0020", "value-20", "bin00020", 0x141516, 2024-10-08, 1970-01-01 00:00:20.123, 1970-01-01 00:00:20.123456, 1970-01-01 00:00:20.123456789, 1970-01-01 00:01:20.321, 1970-01-01 00:01:20.654321, 1970-01-01 00:01:20.321654987, 1, 20.25, 12345678901234571.89, 12345678901234567894.123456789012345678, [20, 21], {4: 40, 14: 41}, {"k4": 1004, "z4": 2004}, [[20, 21], [22]], [{"nested4": 4}, {"nested14": 14}], {"array4": [4, 5]}) +Add: (21, false, 0, -995, 10000000021, 6.25, 15.5, "char0021", "value-21", "bin00021", 0x151617, 2024-10-09, 1970-01-01 00:00:21.123, 1970-01-01 00:00:21.123456, 1970-01-01 00:00:21.123456789, 1970-01-01 00:01:21.321, 1970-01-01 00:01:21.654321, 1970-01-01 00:01:21.321654987, 2, 21.25, 12345678901234572.89, 12345678901234567895.123456789012345678, [21, 22], {5: 50, 15: 51}, {"k5": 1005, "z5": 2005}, [[21, 22], [23]], [{"nested5": 5}, {"nested15": 15}], {"array5": [5, 6]}) + +Commit - snapshot 1 diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/bucket-0/data-2c5a05a4-4c30-4777-81fa-43a43f7c260d-0.mosaic b/test/test_data/mosaic/append_java_compat.db/append_java_compat/bucket-0/data-2c5a05a4-4c30-4777-81fa-43a43f7c260d-0.mosaic new file mode 100644 index 000000000..8c0babf51 Binary files /dev/null and b/test/test_data/mosaic/append_java_compat.db/append_java_compat/bucket-0/data-2c5a05a4-4c30-4777-81fa-43a43f7c260d-0.mosaic differ diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-9d1301ba-d0b3-4d85-ae44-e93f6194876f-0 b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-9d1301ba-d0b3-4d85-ae44-e93f6194876f-0 new file mode 100644 index 000000000..548f4da29 Binary files /dev/null and b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-9d1301ba-d0b3-4d85-ae44-e93f6194876f-0 differ diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-0 b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-0 new file mode 100644 index 000000000..7730dd6f6 Binary files /dev/null and b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-0 differ diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-1 b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-1 new file mode 100644 index 000000000..de8c63592 Binary files /dev/null and b/test/test_data/mosaic/append_java_compat.db/append_java_compat/manifest/manifest-list-1572ca97-622c-452d-8a80-a992a0684230-1 differ diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/schema/schema-0 b/test/test_data/mosaic/append_java_compat.db/append_java_compat/schema/schema-0 new file mode 100644 index 000000000..00a4d8cde --- /dev/null +++ b/test/test_data/mosaic/append_java_compat.db/append_java_compat/schema/schema-0 @@ -0,0 +1,162 @@ +{ + "version" : 3, + "id" : 0, + "fields" : [ { + "id" : 0, + "name" : "id", + "type" : "INT NOT NULL" + }, { + "id" : 1, + "name" : "f_boolean", + "type" : "BOOLEAN" + }, { + "id" : 2, + "name" : "f_tinyint", + "type" : "TINYINT" + }, { + "id" : 3, + "name" : "f_smallint", + "type" : "SMALLINT" + }, { + "id" : 4, + "name" : "f_bigint", + "type" : "BIGINT" + }, { + "id" : 5, + "name" : "f_float", + "type" : "FLOAT" + }, { + "id" : 6, + "name" : "f_double", + "type" : "DOUBLE" + }, { + "id" : 7, + "name" : "f_char", + "type" : "CHAR(8)" + }, { + "id" : 8, + "name" : "f_varchar", + "type" : "VARCHAR(64)" + }, { + "id" : 9, + "name" : "f_binary", + "type" : "BINARY(8)" + }, { + "id" : 10, + "name" : "f_varbinary", + "type" : "VARBINARY(64)" + }, { + "id" : 11, + "name" : "f_date", + "type" : "DATE" + }, { + "id" : 12, + "name" : "f_ts_3", + "type" : "TIMESTAMP(3)" + }, { + "id" : 13, + "name" : "f_ts_6", + "type" : "TIMESTAMP(6)" + }, { + "id" : 14, + "name" : "f_ts_9", + "type" : "TIMESTAMP(9)" + }, { + "id" : 15, + "name" : "f_ltz_3", + "type" : "TIMESTAMP(3) WITH LOCAL TIME ZONE" + }, { + "id" : 16, + "name" : "f_ltz_6", + "type" : "TIMESTAMP(6) WITH LOCAL TIME ZONE" + }, { + "id" : 17, + "name" : "f_ltz_9", + "type" : "TIMESTAMP(9) WITH LOCAL TIME ZONE" + }, { + "id" : 18, + "name" : "f_decimal_1_0", + "type" : "DECIMAL(1, 0)" + }, { + "id" : 19, + "name" : "f_decimal_18_2", + "type" : "DECIMAL(18, 2)" + }, { + "id" : 20, + "name" : "f_decimal_19_2", + "type" : "DECIMAL(19, 2)" + }, { + "id" : 21, + "name" : "f_decimal_38_18", + "type" : "DECIMAL(38, 18)" + }, { + "id" : 22, + "name" : "f_array_int", + "type" : { + "type" : "ARRAY", + "element" : "INT" + } + }, { + "id" : 23, + "name" : "f_map_numeric", + "type" : { + "type" : "MAP", + "key" : "TINYINT", + "value" : "SMALLINT" + } + }, { + "id" : 24, + "name" : "f_map_string_bigint", + "type" : { + "type" : "MAP", + "key" : "STRING", + "value" : "BIGINT" + } + }, { + "id" : 25, + "name" : "f_array_array_int", + "type" : { + "type" : "ARRAY", + "element" : { + "type" : "ARRAY", + "element" : "INT" + } + } + }, { + "id" : 26, + "name" : "f_array_map", + "type" : { + "type" : "ARRAY", + "element" : { + "type" : "MAP", + "key" : "STRING", + "value" : "INT" + } + } + }, { + "id" : 27, + "name" : "f_map_array", + "type" : { + "type" : "MAP", + "key" : "STRING", + "value" : { + "type" : "ARRAY", + "element" : "INT" + } + } + } ], + "highestFieldId" : 27, + "partitionKeys" : [ ], + "primaryKeys" : [ ], + "options" : { + "bucket" : "-1", + "file.block-size" : "1 B", + "mosaic.num-buckets" : "4", + "mosaic.stats-columns" : "id,f_varchar,f_date,f_ts_3,f_ltz_9,f_decimal_18_2", + "target-file-size" : "64 MB", + "manifest.format" : "avro", + "file.format" : "mosaic", + "write.batch-size" : "2" + }, + "timeMillis" : 1787565827926 +} \ No newline at end of file diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/EARLIEST b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/EARLIEST new file mode 100644 index 000000000..56a6051ca --- /dev/null +++ b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/EARLIEST @@ -0,0 +1 @@ +1 \ No newline at end of file diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/LATEST b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/LATEST new file mode 100644 index 000000000..56a6051ca --- /dev/null +++ b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/LATEST @@ -0,0 +1 @@ +1 \ No newline at end of file diff --git a/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/snapshot-1 b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/snapshot-1 new file mode 100644 index 000000000..946d52e7c --- /dev/null +++ b/test/test_data/mosaic/append_java_compat.db/append_java_compat/snapshot/snapshot-1 @@ -0,0 +1,18 @@ +{ + "version" : 3, + "uuid" : "81405eea-75ce-455b-b5ad-32e5bb9c010d", + "id" : 1, + "schemaId" : 0, + "baseManifestList" : "manifest-list-1572ca97-622c-452d-8a80-a992a0684230-0", + "baseManifestListSize" : 1006, + "deltaManifestList" : "manifest-list-1572ca97-622c-452d-8a80-a992a0684230-1", + "deltaManifestListSize" : 1113, + "commitUser" : "a1009c67-ff1d-4292-a67d-f69edbb1ab45", + "writerVersion" : "java-2.1-SNAPSHOT-UNKNOWN", + "commitIdentifier" : 9223372036854775807, + "commitKind" : "APPEND", + "timeMillis" : 1787565830612, + "totalRecordCount" : 6, + "deltaRecordCount" : 6, + "nextRowId" : 0 +} \ No newline at end of file diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/README.md b/test/test_data/mosaic/append_python_compat.db/append_python_compat/README.md new file mode 100644 index 000000000..4971ae2a7 --- /dev/null +++ b/test/test_data/mosaic/append_python_compat.db/append_python_compat/README.md @@ -0,0 +1,62 @@ +Table: append_python_compat +Writer: Paimon Python at commit 0043a70fd88ac75dcb83a8f2da5e72ce91e22b1f +Mosaic version: 0.2.0 + +This table was created and committed through Paimon Python's batch table-write pipeline. + +Schema: +id INT NOT NULL +f_boolean BOOLEAN +f_tinyint TINYINT +f_smallint SMALLINT +f_bigint BIGINT +f_float FLOAT +f_double DOUBLE +f_char CHAR(8) +f_varchar VARCHAR(64) +f_binary BINARY(8) +f_varbinary VARBINARY(64) +f_date DATE +f_ts_3 TIMESTAMP(3) +f_ts_6 TIMESTAMP(6) +f_ts_9 TIMESTAMP(9) +f_ltz_3 TIMESTAMP(3) WITH LOCAL TIME ZONE +f_ltz_6 TIMESTAMP(6) WITH LOCAL TIME ZONE +f_ltz_9 TIMESTAMP(9) WITH LOCAL TIME ZONE +f_decimal_1_0 DECIMAL(1, 0) +f_decimal_18_2 DECIMAL(18, 2) +f_decimal_19_2 DECIMAL(19, 2) +f_decimal_38_18 DECIMAL(38, 18) +f_array_int ARRAY +f_map_numeric MAP +f_map_string_bigint MAP +f_array_array_int ARRAY> +f_array_map ARRAY> +f_map_array MAP> + +Options: +bucket = -1 +file.block-size = 1 B +file.format = mosaic +manifest.format = avro +mosaic.num-buckets = 4 +mosaic.stats-columns = id,f_varchar,f_date,f_ts_3,f_ltz_9,f_decimal_18_2 +target-file-size = 64 MB +write.batch-size = 2 + +Data is written as three row groups with two rows in each row group. Timestamp values are shown in +UTC. The second row has NULL in every nullable field. + +Row group 0: +Add: (1, true, -5, -1000, 10000000001, 1.25, 10.5, "char0001", "value-1", "bin00001", 0x010203, 2024-10-04, 1970-01-01 00:00:01.123, 1970-01-01 00:00:01.123456, 1970-01-01 00:00:01.123456789, 1970-01-01 00:01:01.321, 1970-01-01 00:01:01.654321, 1970-01-01 00:01:01.321654987, -3, 1.25, 12345678901234567.89, 12345678901234567890.123456789012345678, [1, 2], {0: 0, 10: 1}, {"k0": 1000, "z0": 2000}, [[1, 2], [3]], [{"nested0": 0}, {"nested10": 10}], {"array0": [0, 1]}) +Add: (2, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +Row group 1: +Add: (10, true, -3, -998, 10000000010, 3.25, 12.5, "char0010", "value-10", "bin00010", 0x0a0b0c, 2024-10-06, 1970-01-01 00:00:10.123, 1970-01-01 00:00:10.123456, 1970-01-01 00:00:10.123456789, 1970-01-01 00:01:10.321, 1970-01-01 00:01:10.654321, 1970-01-01 00:01:10.321654987, -1, 10.25, 12345678901234569.89, 12345678901234567892.123456789012345678, [10, 11], {2: 20, 12: 21}, {"k2": 1002, "z2": 2002}, [[10, 11], [12]], [{"nested2": 2}, {"nested12": 12}], {"array2": [2, 3]}) +Add: (11, false, -2, -997, 10000000011, 4.25, 13.5, "char0011", "value-11", "bin00011", 0x0b0c0d, 2024-10-07, 1970-01-01 00:00:11.123, 1970-01-01 00:00:11.123456, 1970-01-01 00:00:11.123456789, 1970-01-01 00:01:11.321, 1970-01-01 00:01:11.654321, 1970-01-01 00:01:11.321654987, 0, 11.25, 12345678901234570.89, 12345678901234567893.123456789012345678, [11, 12], {3: 30, 13: 31}, {"k3": 1003, "z3": 2003}, [[11, 12], [13]], [{"nested3": 3}, {"nested13": 13}], {"array3": [3, 4]}) + +Row group 2: +Add: (20, true, -1, -996, 10000000020, 5.25, 14.5, "char0020", "value-20", "bin00020", 0x141516, 2024-10-08, 1970-01-01 00:00:20.123, 1970-01-01 00:00:20.123456, 1970-01-01 00:00:20.123456789, 1970-01-01 00:01:20.321, 1970-01-01 00:01:20.654321, 1970-01-01 00:01:20.321654987, 1, 20.25, 12345678901234571.89, 12345678901234567894.123456789012345678, [20, 21], {4: 40, 14: 41}, {"k4": 1004, "z4": 2004}, [[20, 21], [22]], [{"nested4": 4}, {"nested14": 14}], {"array4": [4, 5]}) +Add: (21, false, 0, -995, 10000000021, 6.25, 15.5, "char0021", "value-21", "bin00021", 0x151617, 2024-10-09, 1970-01-01 00:00:21.123, 1970-01-01 00:00:21.123456, 1970-01-01 00:00:21.123456789, 1970-01-01 00:01:21.321, 1970-01-01 00:01:21.654321, 1970-01-01 00:01:21.321654987, 2, 21.25, 12345678901234572.89, 12345678901234567895.123456789012345678, [21, 22], {5: 50, 15: 51}, {"k5": 1005, "z5": 2005}, [[21, 22], [23]], [{"nested5": 5}, {"nested15": 15}], {"array5": [5, 6]}) + +Commit - snapshot 1 diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/bucket-0/data-087c1f82-5909-45cd-a6a8-ff0314c8e365-0.mosaic b/test/test_data/mosaic/append_python_compat.db/append_python_compat/bucket-0/data-087c1f82-5909-45cd-a6a8-ff0314c8e365-0.mosaic new file mode 100644 index 000000000..6740a11a7 Binary files /dev/null and b/test/test_data/mosaic/append_python_compat.db/append_python_compat/bucket-0/data-087c1f82-5909-45cd-a6a8-ff0314c8e365-0.mosaic differ diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-82c6d797-6335-462b-b86e-a076ae20578a-0 b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-82c6d797-6335-462b-b86e-a076ae20578a-0 new file mode 100644 index 000000000..de32901f5 Binary files /dev/null and b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-82c6d797-6335-462b-b86e-a076ae20578a-0 differ diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-0 b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-0 new file mode 100644 index 000000000..da8b9dfe0 Binary files /dev/null and b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-0 differ diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-1 b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-1 new file mode 100644 index 000000000..b89a3a145 Binary files /dev/null and b/test/test_data/mosaic/append_python_compat.db/append_python_compat/manifest/manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-1 differ diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/schema/schema-0 b/test/test_data/mosaic/append_python_compat.db/append_python_compat/schema/schema-0 new file mode 100644 index 000000000..55e18d15a --- /dev/null +++ b/test/test_data/mosaic/append_python_compat.db/append_python_compat/schema/schema-0 @@ -0,0 +1,201 @@ +{ + "version": 3, + "id": 0, + "fields": [ + { + "id": 0, + "name": "id", + "type": "INT NOT NULL" + }, + { + "id": 1, + "name": "f_boolean", + "type": "BOOLEAN" + }, + { + "id": 2, + "name": "f_tinyint", + "type": "TINYINT" + }, + { + "id": 3, + "name": "f_smallint", + "type": "SMALLINT" + }, + { + "id": 4, + "name": "f_bigint", + "type": "BIGINT" + }, + { + "id": 5, + "name": "f_float", + "type": "FLOAT" + }, + { + "id": 6, + "name": "f_double", + "type": "DOUBLE" + }, + { + "id": 7, + "name": "f_char", + "type": "CHAR(8)" + }, + { + "id": 8, + "name": "f_varchar", + "type": "VARCHAR(64)" + }, + { + "id": 9, + "name": "f_binary", + "type": "BINARY(8)" + }, + { + "id": 10, + "name": "f_varbinary", + "type": "VARBINARY(64)" + }, + { + "id": 11, + "name": "f_date", + "type": "DATE" + }, + { + "id": 12, + "name": "f_ts_3", + "type": "TIMESTAMP(3)" + }, + { + "id": 13, + "name": "f_ts_6", + "type": "TIMESTAMP(6)" + }, + { + "id": 14, + "name": "f_ts_9", + "type": "TIMESTAMP(9)" + }, + { + "id": 15, + "name": "f_ltz_3", + "type": "TIMESTAMP(3) WITH LOCAL TIME ZONE" + }, + { + "id": 16, + "name": "f_ltz_6", + "type": "TIMESTAMP(6) WITH LOCAL TIME ZONE" + }, + { + "id": 17, + "name": "f_ltz_9", + "type": "TIMESTAMP(9) WITH LOCAL TIME ZONE" + }, + { + "id": 18, + "name": "f_decimal_1_0", + "type": "DECIMAL(1, 0)" + }, + { + "id": 19, + "name": "f_decimal_18_2", + "type": "DECIMAL(18, 2)" + }, + { + "id": 20, + "name": "f_decimal_19_2", + "type": "DECIMAL(19, 2)" + }, + { + "id": 21, + "name": "f_decimal_38_18", + "type": "DECIMAL(38, 18)" + }, + { + "id": 22, + "name": "f_array_int", + "type": { + "type": "ARRAY", + "element": "INT", + "nullable": true + } + }, + { + "id": 23, + "name": "f_map_numeric", + "type": { + "type": "MAP", + "key": "TINYINT NOT NULL", + "value": "SMALLINT", + "nullable": true + } + }, + { + "id": 24, + "name": "f_map_string_bigint", + "type": { + "type": "MAP", + "key": "STRING NOT NULL", + "value": "BIGINT", + "nullable": true + } + }, + { + "id": 25, + "name": "f_array_array_int", + "type": { + "type": "ARRAY", + "element": { + "type": "ARRAY", + "element": "INT", + "nullable": true + }, + "nullable": true + } + }, + { + "id": 26, + "name": "f_array_map", + "type": { + "type": "ARRAY", + "element": { + "type": "MAP", + "key": "STRING NOT NULL", + "value": "INT", + "nullable": true + }, + "nullable": true + } + }, + { + "id": 27, + "name": "f_map_array", + "type": { + "type": "MAP>", + "key": "STRING NOT NULL", + "value": { + "type": "ARRAY", + "element": "INT", + "nullable": true + }, + "nullable": true + } + } + ], + "highestFieldId": 27, + "partitionKeys": [], + "primaryKeys": [], + "options": { + "bucket": "-1", + "file.block-size": "1 B", + "file.format": "mosaic", + "manifest.format": "avro", + "mosaic.num-buckets": "4", + "mosaic.stats-columns": "id,f_varchar,f_date,f_ts_3,f_ltz_9,f_decimal_18_2", + "target-file-size": "64 MB", + "write.batch-size": "2" + }, + "comment": null, + "timeMillis": 1787569308159 +} \ No newline at end of file diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/LATEST b/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/LATEST new file mode 100644 index 000000000..56a6051ca --- /dev/null +++ b/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/LATEST @@ -0,0 +1 @@ +1 \ No newline at end of file diff --git a/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/snapshot-1 b/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/snapshot-1 new file mode 100644 index 000000000..9a7fb6f06 --- /dev/null +++ b/test/test_data/mosaic/append_python_compat.db/append_python_compat/snapshot/snapshot-1 @@ -0,0 +1,15 @@ +{ + "version": 3, + "id": 1, + "schemaId": 0, + "baseManifestList": "manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-0", + "deltaManifestList": "manifest-list-16d5961d-78cd-498a-9cff-80ef60258224-1", + "totalRecordCount": 6, + "deltaRecordCount": 6, + "commitUser": "c6489507-9e13-4fec-abfa-74b1957f8008", + "commitIdentifier": 9223372036854775807, + "commitKind": "APPEND", + "timeMillis": 1787569308179, + "uuid": "1d6d2d3e-abe1-45ac-88d0-d90429aafbb5", + "writerVersion": "python-2.1.dev-0043a70fd88ac75dcb83a8f2da5e72ce91e22b1f" +} \ No newline at end of file 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 000000000..8a2fe25c9 --- /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 000000000..5184c7a9f 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 000000000..fa900ac0d 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 000000000..761ee7341 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 000000000..8610e46e1 Binary files /dev/null and b/test/test_data/parquet/vector_compatibility/rust_vector_nullable.parquet differ diff --git a/third_party/versions.txt b/third_party/versions.txt index 2e20dc72f..1afdf7a1c 100644 --- a/third_party/versions.txt +++ b/third_party/versions.txt @@ -68,6 +68,10 @@ PAIMON_ARROW_BUILD_VERSION=17.0.0 PAIMON_ARROW_BUILD_SHA256_CHECKSUM=9d280d8042e7cf526f8c28d170d93bfab65e50f94569f6a790982a878d8d898d PAIMON_ARROW_PKG_NAME=apache-arrow-${PAIMON_ARROW_BUILD_VERSION}.tar.gz +PAIMON_MOSAIC_BUILD_VERSION=0.2.0 +PAIMON_MOSAIC_BUILD_SHA256_CHECKSUM=8123eaadd293a7904b692eff900484691711dc0cb24878166a2f90a378d37a22 +PAIMON_MOSAIC_PKG_NAME=apache-paimon-mosaic-${PAIMON_MOSAIC_BUILD_VERSION}-src.tgz + PAIMON_AVRO_BUILD_VERSION=c499eefb48aa2db906c7bca14a047223806f36db PAIMON_AVRO_BUILD_SHA256_CHECKSUM=9771f1dcfe3c01aff7ff670e873e66d3406362f71941821d482de65f3d32d780 PAIMON_AVRO_PKG_NAME=avro-${PAIMON_AVRO_BUILD_VERSION}.tar.gz @@ -97,6 +101,10 @@ PAIMON_AWS_S2N_BUILD_VERSION=v1.7.4 PAIMON_AWS_S2N_BUILD_SHA256_CHECKSUM=af5ce0783fd9e05ed1899fda0c76e02fa5dd92128018e5bfe71634312d2ce8e7 PAIMON_AWS_S2N_PKG_NAME=s2n-${PAIMON_AWS_S2N_BUILD_VERSION}.zip +PAIMON_OSS_SDK_V2_BUILD_VERSION=0.1.2 +PAIMON_OSS_SDK_V2_BUILD_SHA256_CHECKSUM=6a6c00692c8fe8461594a359de8ba97f4468d90f92ffb6bcb92b82a9d1b9311b +PAIMON_OSS_SDK_V2_PKG_NAME=alibabacloud-oss-cpp-sdk-v2-${PAIMON_OSS_SDK_V2_BUILD_VERSION}.tar.gz + PAIMON_FMT_BUILD_VERSION=11.2.0 PAIMON_FMT_BUILD_SHA256_CHECKSUM=bc23066d87ab3168f27cef3e97d545fa63314f5c79df5ea444d41d56f962c6af PAIMON_FMT_PKG_NAME=fmt-${PAIMON_FMT_BUILD_VERSION}.tar.gz @@ -158,10 +166,12 @@ DEPENDENCIES=( "PAIMON_LZ4_URL ${PAIMON_LZ4_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/lz4/lz4/archive/${PAIMON_LZ4_BUILD_VERSION}.tar.gz" "PAIMON_PROTOBUF_URL ${PAIMON_PROTOBUF_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/protocolbuffers/protobuf/releases/download/v${PAIMON_PROTOBUF_BUILD_VERSION}/protobuf-all-${PAIMON_PROTOBUF_BUILD_VERSION}.tar.gz" "PAIMON_TBB_URL ${PAIMON_TBB_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/uxlfoundation/oneTBB/archive/refs/tags/${PAIMON_TBB_BUILD_VERSION}.tar.gz" + "PAIMON_OSS_SDK_V2_URL ${PAIMON_OSS_SDK_V2_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/aliyun/alibabacloud-oss-cpp-sdk-v2/archive/refs/tags/${PAIMON_OSS_SDK_V2_BUILD_VERSION}.tar.gz" "PAIMON_ORC_URL ${PAIMON_ORC_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/apache/orc/archive/refs/tags/${PAIMON_ORC_BUILD_VERSION}.tar.gz" "PAIMON_GTEST_URL ${PAIMON_GTEST_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/googletest/archive/release-${PAIMON_GTEST_BUILD_VERSION}.tar.gz" "PAIMON_BENCHMARK_URL ${PAIMON_BENCHMARK_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/benchmark/archive/refs/tags/v${PAIMON_BENCHMARK_BUILD_VERSION}.tar.gz" "PAIMON_ARROW_URL ${PAIMON_ARROW_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/apache/arrow/releases/download/apache-arrow-${PAIMON_ARROW_BUILD_VERSION}/apache-arrow-${PAIMON_ARROW_BUILD_VERSION}.tar.gz" + "PAIMON_MOSAIC_URL ${PAIMON_MOSAIC_PKG_NAME} https://downloads.apache.org/paimon/paimon-mosaic-${PAIMON_MOSAIC_BUILD_VERSION}/${PAIMON_MOSAIC_PKG_NAME}" "PAIMON_AVRO_URL ${PAIMON_AVRO_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/apache/avro/archive/${PAIMON_AVRO_BUILD_VERSION}.tar.gz" "PAIMON_AWS_C_AUTH_URL ${PAIMON_AWS_C_AUTH_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/awslabs/aws-c-auth/archive/refs/tags/${PAIMON_AWS_C_AUTH_BUILD_VERSION}.tar.gz" "PAIMON_AWS_C_CAL_URL ${PAIMON_AWS_C_CAL_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/awslabs/aws-c-cal/archive/refs/tags/${PAIMON_AWS_C_CAL_BUILD_VERSION}.tar.gz"