From 37dce2aa1c62505c077619a68ffb409d6409414a Mon Sep 17 00:00:00 2001 From: zhouyc-ali <301200591+zhouyc-ali@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:01:56 +0000 Subject: [PATCH 1/4] feat(vortex): add predicate pushdown and optimize selective reads --- benchmark/CMakeLists.txt | 5 + benchmark/benchmark_suite.cpp | 7 + cmake_modules/ApplyVortexReadPatch.cmake | 129 ++ cmake_modules/ThirdpartyToolchain.cmake | 10 +- cmake_modules/vortex-fmm-cache-shards.diff | 114 ++ cmake_modules/vortex-point-read.diff | 357 +++++ cmake_modules/vortex.diff | 71 + crates/vortex_callback_io/callback_io.rs | 622 +++++++- docs/source/user_guide/read.rst | 20 + src/paimon/format/vortex/CMakeLists.txt | 2 + src/paimon/format/vortex/vortex_ffi.h | 29 + src/paimon/format/vortex/vortex_ffi_util.h | 2 + .../vortex/vortex_file_batch_reader.cpp | 298 ++-- .../format/vortex/vortex_file_batch_reader.h | 36 +- .../format/vortex/vortex_file_format.cpp | 30 +- .../format/vortex/vortex_file_format_test.cpp | 798 ++++++++++ .../format/vortex/vortex_io_callbacks.h | 4 + .../vortex/vortex_predicate_converter.cpp | 613 ++++++++ .../vortex/vortex_predicate_converter.h | 60 + .../vortex_predicate_converter_test.cpp | 1341 +++++++++++++++++ .../format/vortex/vortex_reader_builder.h | 16 +- test/inte/scan_and_read_inte_test.cpp | 227 +++ 22 files changed, 4596 insertions(+), 195 deletions(-) create mode 100644 cmake_modules/ApplyVortexReadPatch.cmake create mode 100644 cmake_modules/vortex-fmm-cache-shards.diff create mode 100644 cmake_modules/vortex-point-read.diff create mode 100644 src/paimon/format/vortex/vortex_predicate_converter.cpp create mode 100644 src/paimon/format/vortex/vortex_predicate_converter.h create mode 100644 src/paimon/format/vortex/vortex_predicate_converter_test.cpp diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 60ac3fb34..e66323e10 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -35,6 +35,11 @@ if(PAIMON_ENABLE_AVRO) ${PAIMON_AVRO_FILE_FORMAT_STATIC_LINK_LIBS}) endif() +if(PAIMON_ENABLE_VORTEX) + list(APPEND PAIMON_BENCHMARK_STATIC_LINK_LIBS + ${PAIMON_VORTEX_FILE_FORMAT_STATIC_LINK_LIBS}) +endif() + set(PAIMON_BENCHMARK_PLATFORM_LINK_LIBS) if(UNIX AND NOT APPLE) find_library(PAIMON_BENCHMARK_RT_LIBRARY rt) diff --git a/benchmark/benchmark_suite.cpp b/benchmark/benchmark_suite.cpp index 0a48a9173..e9bb80b78 100644 --- a/benchmark/benchmark_suite.cpp +++ b/benchmark/benchmark_suite.cpp @@ -211,6 +211,13 @@ bool IsFileFormatSupported(const std::string& format) { return true; #else return false; +#endif + } + if (format == "vortex") { +#ifdef PAIMON_ENABLE_VORTEX + return true; +#else + return false; #endif } return false; diff --git a/cmake_modules/ApplyVortexReadPatch.cmake b/cmake_modules/ApplyVortexReadPatch.cmake new file mode 100644 index 000000000..90b3c23d8 --- /dev/null +++ b/cmake_modules/ApplyVortexReadPatch.cmake @@ -0,0 +1,129 @@ +# 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. + +# A reused .patched marker can predate the scan/Arrow ABI additions. Validate the base patch +# before changing any reader source; unknown or older sources need a fresh pinned checkout. +set(_base_files vortex-ffi/src/lib.rs vortex-ffi/src/scan.rs vortex-ffi/src/dtype.rs + vortex-ffi/cinclude/vortex.h) +set(_base_after + 3fdbba15c951286ea0e565e3e21d75de1d7d4065a358f16aba2d47750e7ace76 + 00aba4a0d7339fe8d961fb302572deae58adb9f06eb9b29313aa9503e8352d49 + af3c5d2652a37cd8c7e7d1e2c377a290c9a5e0dac158bd4ed8d0b6922552bec0 + 1ad352c8476bfa41972dd597261fb2753b6903f4e90e9d606b85f9b04ed4c895) +foreach(_index RANGE 0 3) + list(GET _base_files ${_index} _file) + list(GET _base_after ${_index} _expected) + file(SHA256 "${VORTEX_SOURCE_DIR}/${_file}" _actual) + if(NOT _actual STREQUAL _expected) + message(FATAL_ERROR "Vortex base patch is missing, outdated or modified: ${_file}; use a fresh pinned source checkout" + ) + endif() +endforeach() + +# Accept clean reader sources, the reader patch alone, or both reader patches. +set(_files + vortex-layout/src/layouts/dict/reader.rs + vortex-ffi/src/session.rs + vortex-layout/src/layouts/struct_/reader.rs + vortex-file/src/v2/file_stats_reader.rs + vortex-layout/src/segments/shared.rs) +set(_before + 760c3de9f14a93b25322441ad835f9ea42286f2f01b0c72c95c03941684b5581 + 0a5bf14492f8e03f6dfa265b30c355b8f6ae166558a1655038ee51ae6d50a6ca + d77641a60e1d4daf84f58f30d5cc0f5f64bc09f30ff1d711862b66606ec9a772 + e75aefbaa363748ab9e75cc21b1ec6d74701e2c973d1eb204e657c4e8076e7f6 + 2da8289cbd44c64bcbdf102cdc2be291448fc3aa77a0ca34d2f7efb7dfde2419) +set(_reader_patched + 4bfb5f4bf8f707027702e9157d2d40d1b64445692d766a3a2d9f685998af3806 + 093cf553ed52c01a31ff131b96f226157f37ffe67593136d3bd67832f85e362e + d77641a60e1d4daf84f58f30d5cc0f5f64bc09f30ff1d711862b66606ec9a772 + e75aefbaa363748ab9e75cc21b1ec6d74701e2c973d1eb204e657c4e8076e7f6 + 2da8289cbd44c64bcbdf102cdc2be291448fc3aa77a0ca34d2f7efb7dfde2419) +set(_after + 4bfb5f4bf8f707027702e9157d2d40d1b64445692d766a3a2d9f685998af3806 + 093cf553ed52c01a31ff131b96f226157f37ffe67593136d3bd67832f85e362e + 79be1013b9f04e683de40f8a9881e4d15e07945dd8d86b7760d3abf01d14ceb6 + 36d599cbea7c0bd1855d560c2e171e9b61117f4e69cfe3bac4dcd0b126cb549f + 83edf1951eb564ab86e28a6e5f1c9fdb49788fc42ba43d7d26db9fd8f08cf40d) +set(_cache_patch "${CMAKE_CURRENT_LIST_DIR}/vortex-fmm-cache-shards.diff") +set(_all_clean TRUE) +set(_all_reader_patched TRUE) +set(_all_after TRUE) +foreach(_index RANGE 0 4) + list(GET _files ${_index} _file) + list(GET _before ${_index} _old) + list(GET _reader_patched ${_index} _intermediate) + list(GET _after ${_index} _new) + file(SHA256 "${VORTEX_SOURCE_DIR}/${_file}" _actual) + if(NOT _actual STREQUAL _old) + set(_all_clean FALSE) + endif() + if(NOT _actual STREQUAL _intermediate) + set(_all_reader_patched FALSE) + endif() + if(NOT _actual STREQUAL _new) + set(_all_after FALSE) + endif() +endforeach() +if(NOT _all_clean + AND NOT _all_reader_patched + AND NOT _all_after) + message(FATAL_ERROR "Vortex read patch: unknown or partly updated source; preserve and inspect local changes" + ) +endif() + +function(_run_vortex_read_patch _path _direction _dry_run) + set(_arguments --batch ${_direction}) + if(_dry_run) + list(APPEND _arguments --dry-run) + endif() + execute_process(COMMAND patch ${_arguments} -p1 -i "${_path}" + WORKING_DIRECTORY "${VORTEX_SOURCE_DIR}" + RESULT_VARIABLE _result + OUTPUT_VARIABLE _stdout + ERROR_VARIABLE _stderr) + if(NOT _result EQUAL 0) + message(FATAL_ERROR "Vortex read patch failed: ${_stdout}${_stderr}") + endif() +endfunction() + +# The patches affect disjoint files, so both dry-runs complete before any write. +if(_all_clean) + _run_vortex_read_patch("${VORTEX_READ_PATCH}" --forward TRUE) +else() + _run_vortex_read_patch("${VORTEX_READ_PATCH}" --reverse TRUE) +endif() +if(_all_after) + _run_vortex_read_patch("${_cache_patch}" --reverse TRUE) +else() + _run_vortex_read_patch("${_cache_patch}" --forward TRUE) +endif() + +if(_all_clean) + _run_vortex_read_patch("${VORTEX_READ_PATCH}" --forward FALSE) +endif() +if(NOT _all_after) + _run_vortex_read_patch("${_cache_patch}" --forward FALSE) +endif() +foreach(_index RANGE 0 4) + list(GET _files ${_index} _file) + list(GET _after ${_index} _expected) + file(SHA256 "${VORTEX_SOURCE_DIR}/${_file}" _actual) + if(NOT _actual STREQUAL _expected) + message(FATAL_ERROR "Vortex read patch output hash differs: ${_file}") + endif() +endforeach() diff --git a/cmake_modules/ThirdpartyToolchain.cmake b/cmake_modules/ThirdpartyToolchain.cmake index ac94fc2b4..1f35dd431 100644 --- a/cmake_modules/ThirdpartyToolchain.cmake +++ b/cmake_modules/ThirdpartyToolchain.cmake @@ -1529,9 +1529,13 @@ macro(build_vortex) # returns immediately when nothing changed, and copy_if_different keeps # the archive's timestamp stable so dependents do not relink. BUILD_ALWAYS TRUE - BUILD_COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${VORTEX_CALLBACK_IO_SRC}" - "/vortex-ffi/src/callback_io.rs" + BUILD_COMMAND ${CMAKE_COMMAND} "-DVORTEX_SOURCE_DIR=" + "-DVORTEX_READ_PATCH=${CMAKE_CURRENT_LIST_DIR}/vortex-point-read.diff" + -P + "${CMAKE_CURRENT_LIST_DIR}/ApplyVortexReadPatch.cmake" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${VORTEX_CALLBACK_IO_SRC}" + "/vortex-ffi/src/callback_io.rs" COMMAND ${CMAKE_COMMAND} -E env "FLATC=${VORTEX_FLATC}" "RUSTUP_TOOLCHAIN=stable" "CC=${CMAKE_C_COMPILER}" "CXX=${CMAKE_CXX_COMPILER}" diff --git a/cmake_modules/vortex-fmm-cache-shards.diff b/cmake_modules/vortex-fmm-cache-shards.diff new file mode 100644 index 000000000..8ce134137 --- /dev/null +++ b/cmake_modules/vortex-fmm-cache-shards.diff @@ -0,0 +1,114 @@ +diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs +--- a/vortex-layout/src/layouts/struct_/reader.rs ++++ b/vortex-layout/src/layouts/struct_/reader.rs +@@ -115,7 +115,8 @@ + expanded_root_expr, + lazy_children, + field_lookup, +- partitioned_expr_cache: Default::default(), ++ // Each file evaluates few expressions; avoid CPU-count-sized shard allocations. ++ partitioned_expr_cache: DashMap::with_hasher_and_shard_amount(Default::default(), 16), + }) + } + +diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs +--- a/vortex-file/src/v2/file_stats_reader.rs ++++ b/vortex-file/src/v2/file_stats_reader.rs +@@ -74,7 +74,8 @@ + file_stats, + struct_fields, + session, +- prune_cache: Default::default(), ++ // File-level pruning reuses only a small set of expressions per reader. ++ prune_cache: DashMap::with_hasher_and_shard_amount(Default::default(), 16), + } + } + +diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs +--- a/vortex-layout/src/segments/shared.rs ++++ b/vortex-layout/src/segments/shared.rs +@@ -32,7 +32,8 @@ + pub fn new(inner: S) -> Self { + Self { + inner, +- in_flight: DashMap::default(), ++ // A per-file request set does not need a shard for every available CPU. ++ in_flight: DashMap::with_hasher_and_shard_amount(Default::default(), 16), + } + } + } +@@ -65,10 +66,16 @@ + + #[cfg(test)] + mod tests { ++ use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; +- ++ use std::thread; ++ ++ use futures::executor::block_on; ++ use futures::future::join_all; + use vortex_buffer::ByteBuffer; ++ use vortex_error::VortexResult; ++ use vortex_error::vortex_err; + + use super::*; + use crate::segments::SegmentSink; +@@ -119,5 +126,56 @@ + } + + #[tokio::test] ++ async fn test_shared_source_deduplicates_many_segments_across_threads() -> VortexResult<()> { ++ let source = CountingSegmentSource::default(); ++ let mut ids = Vec::new(); ++ let mut expected = Vec::new(); ++ for value in 0..64u8 { ++ let data = ByteBuffer::from(vec![value; 4]); ++ ids.push( ++ source ++ .segments ++ .write(SequenceId::root().downgrade(), vec![data.clone()]) ++ .await?, ++ ); ++ expected.push(data); ++ } ++ ++ let shared_source = SharedSegmentSource::new(source.clone()); ++ let barrier = Barrier::new(8); ++ thread::scope(|scope| -> VortexResult<()> { ++ let mut requests = Vec::new(); ++ for _ in 0..8 { ++ let shared_source = &shared_source; ++ let ids = &ids; ++ let expected = &expected; ++ let barrier = &barrier; ++ requests.push(scope.spawn(move || -> VortexResult<()> { ++ let pending = ids ++ .iter() ++ .map(|id| shared_source.request(*id)) ++ .collect::>(); ++ // Hold every request until all threads share the same in-flight futures. ++ barrier.wait(); ++ for (actual, expected) in block_on(join_all(pending)).into_iter().zip(expected) ++ { ++ assert_eq!(actual?.unwrap_host(), *expected); ++ } ++ Ok(()) ++ })); ++ } ++ for request in requests { ++ request ++ .join() ++ .map_err(|_| vortex_err!("segment requester panicked"))??; ++ } ++ Ok(()) ++ })?; ++ ++ assert_eq!(source.request_count.load(Ordering::SeqCst), ids.len()); ++ Ok(()) ++ } ++ ++ #[tokio::test] + async fn test_shared_source_handles_dropped_futures() { + let source = CountingSegmentSource::default(); diff --git a/cmake_modules/vortex-point-read.diff b/cmake_modules/vortex-point-read.diff new file mode 100644 index 000000000..8b311e1cd --- /dev/null +++ b/cmake_modules/vortex-point-read.diff @@ -0,0 +1,357 @@ +diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs +--- a/vortex-layout/src/layouts/dict/reader.rs ++++ b/vortex-layout/src/layouts/dict/reader.rs +@@ -45,8 +45,8 @@ + values_len: usize, + /// Cached dict values array + values_array: OnceLock, +- /// Cache of expression evaluation results on the values array by expression +- values_evals: DashMap, ++ /// Expression results are only cached for filtering, not payload projection. ++ values_evals: OnceLock>, + + values: LayoutReaderRef, + codes: LayoutReaderRef, +@@ -135,12 +135,13 @@ + // shouldn't. + // TODO(joe): fixme + ++ let values_evals = self.values_evals.get_or_init(DashMap::default); + // Check cache first with read-only lock +- if let Some(fut) = self.values_evals.get(&expr) { ++ if let Some(fut) = values_evals.get(&expr) { + return fut.clone(); + } + +- self.values_evals ++ values_evals + .entry(expr.clone()) + .or_insert_with(|| { + self.values_array_uncanonical() +@@ -290,12 +291,15 @@ + use vortex_array::session::ArraySession; + use vortex_array::validity::Validity; + use vortex_error::VortexExpect; ++ use vortex_error::VortexResult; ++ use vortex_error::vortex_err; + use vortex_io::runtime::Handle; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + use vortex_session::VortexSession; + ++ use super::DictReader; + use crate::LayoutId; + use crate::LayoutRef; + use crate::LayoutStrategy; +@@ -373,9 +377,12 @@ + Nullability::NonNullable, + ); + assert!(layout.encoding_id() == LayoutId::new("vortex.dict")); +- let actual = layout ++ let reader = layout + .new_reader("".into(), segments, &session, &Default::default()) +- .unwrap() ++ .unwrap(); ++ let dict_reader = reader.as_any().downcast_ref::().unwrap(); ++ assert!(dict_reader.values_evals.get().is_none()); ++ let actual = reader + .projection_evaluation( + &(0..layout.row_count()), + &expression, +@@ -384,6 +391,7 @@ + .unwrap() + .await + .unwrap(); ++ assert!(dict_reader.values_evals.get().is_none()); + let expected = StructArray::try_new( + FieldNames::from([FieldName::from("top")]), + vec![ +@@ -402,6 +410,92 @@ + .unwrap() + .into_array(); + assert_arrays_eq!(actual, expected); ++ }) ++ } ++ ++ #[test] ++ fn repeated_nullable_filters_share_lazy_expression_cache() -> VortexResult<()> { ++ block_on(|handle| async move { ++ let session = session_with_handle(handle); ++ let strategy = DictStrategy::new( ++ FlatLayoutStrategy::default(), ++ FlatLayoutStrategy::default(), ++ FlatLayoutStrategy::default(), ++ DictLayoutOptions::default(), ++ ); ++ let array = VarBinArray::from_iter( ++ [ ++ Some("abc"), ++ Some("def"), ++ None, ++ Some("abc"), ++ Some("def"), ++ None, ++ Some("abc"), ++ Some("def"), ++ None, ++ ], ++ DType::Utf8(Nullability::Nullable), ++ ) ++ .into_array(); ++ let segments = Arc::new(TestSegments::default()); ++ let (ptr, eof) = SequenceId::root().split(); ++ let layout = strategy ++ .write_stream( ++ ArrayContext::empty(), ++ Arc::::clone(&segments), ++ SequentialStreamAdapter::new( ++ DType::Utf8(Nullability::Nullable), ++ array.to_array_stream().sequenced(ptr), ++ ) ++ .sendable(), ++ eof, ++ &session, ++ ) ++ .await?; ++ assert!(layout.encoding_id() == LayoutId::new("vortex.dict")); ++ let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; ++ let dict_reader = reader ++ .as_any() ++ .downcast_ref::() ++ .ok_or_else(|| vortex_err!("dict fixture did not create a DictReader"))?; ++ assert!(dict_reader.values_evals.get().is_none()); ++ ++ let filter = eq( ++ root(), ++ lit(vortex_array::scalar::Scalar::utf8( ++ "abc", ++ Nullability::Nullable, ++ )), ++ ); ++ let first_mask = reader ++ .filter_evaluation(&(0..9), &filter, MaskFuture::new_true(9))? ++ .await?; ++ let expected = [true, false, false, true, false, false, true, false, false]; ++ assert_arrays_eq!(first_mask.into_array(), BoolArray::from_iter(expected)); ++ let cache = dict_reader ++ .values_evals ++ .get() ++ .ok_or_else(|| vortex_err!("filter did not initialize the expression cache"))?; ++ assert_eq!(cache.len(), 1); ++ let first_future = cache ++ .get(&filter) ++ .ok_or_else(|| vortex_err!("first expression future is missing"))? ++ .value() ++ .clone(); ++ ++ let second_mask = reader ++ .filter_evaluation(&(0..9), &filter, MaskFuture::new_true(9))? ++ .await?; ++ assert_arrays_eq!(second_mask.into_array(), BoolArray::from_iter(expected)); ++ assert_eq!(cache.len(), 1); ++ let second_future = cache ++ .get(&filter) ++ .ok_or_else(|| vortex_err!("repeated expression future is missing"))? ++ .value() ++ .clone(); ++ assert!(first_future.ptr_eq(&second_future)); ++ Ok(()) + }) + } + +diff --git a/vortex-ffi/src/session.rs b/vortex-ffi/src/session.rs +--- a/vortex-ffi/src/session.rs ++++ b/vortex-ffi/src/session.rs +@@ -1,7 +1,10 @@ + // SPDX-License-Identifier: Apache-2.0 + // SPDX-FileCopyrightText: Copyright the Vortex contributors + ++use std::sync::LazyLock; ++ + use vortex::VortexSessionDefault; ++use vortex::array::arrow::ArrowSessionExt; + use vortex::error::VortexResult; + use vortex::error::vortex_ensure; + use vortex::io::runtime::BlockingRuntime; +@@ -16,6 +19,34 @@ + VortexSession, + vx_session + ); ++ ++// This session is private to the Paimon callback reader. Public/configurable ++// sessions must remain fresh: VortexSession::clone shares mutable SessionVars. ++// Single-file callback opens keep file metadata, caches and execution contexts ++// outside this default session. Initialize the only lazy read-path registry ++// before publishing it, and never configure/register against these handles. ++static PAIMON_READ_SESSION: LazyLock = LazyLock::new(|| { ++ let session = VortexSession::default().with_handle(RUNTIME.handle()); ++ drop(session.arrow()); ++ session ++}); ++ ++/// Create an owned handle to Paimon's shared default callback-reader session. ++/// ++/// This private integration entry point does not change the isolation of ++/// [`vx_session_new`] or [`vx_session_new_with`]. Free the owned handle with ++/// [`vx_session_free`]; other reader handles remain valid. ++/// ++/// # Safety ++/// ++/// The handle is only for Paimon's fixed-runtime, single-file callback read and ++/// Arrow export paths. The caller must not configure it, register plugins, change ++/// session variables/allocators/runtime, or pass it to a writer or multi-file ++/// source. Those operations require a fresh [`vx_session_new`] session. ++#[unsafe(no_mangle)] ++pub unsafe extern "C-unwind" fn vx_paimon_read_session_new() -> *mut vx_session { ++ vx_session::new(PAIMON_READ_SESSION.clone()) ++} + + /// Create an FFI session from a configured default session. + pub fn vx_session_new_with( +@@ -58,9 +89,30 @@ + + #[cfg(test)] + mod tests { ++ use std::collections::BTreeSet; ++ use std::thread; ++ ++ use arrow_schema::DataType; ++ use arrow_schema::Field; ++ use vortex::array::IntoArray; ++ use vortex::array::VortexSessionExecute; ++ use vortex::array::arrays::PrimitiveArray; ++ use vortex::array::arrow::ArrowSessionExt; ++ use vortex::array::session::ArraySessionExt; ++ use vortex::array::validity::Validity; ++ use vortex::buffer::buffer; ++ use vortex::error::VortexResult; ++ use vortex::error::vortex_err; ++ use vortex::io::session::RuntimeSessionExt; ++ use vortex::session::registry::Id; ++ use vortex_array::assert_arrays_eq; ++ ++ use crate::session::vx_paimon_read_session_new; ++ use crate::session::vx_session; + use crate::session::vx_session_clone; + use crate::session::vx_session_free; + use crate::session::vx_session_new; ++ use crate::session::vx_session_new_with; + + #[test] + #[cfg_attr(miri, ignore)] +@@ -81,9 +133,113 @@ + + let copy = vx_session_clone(session); + assert!(!copy.is_null()); ++ drop(vx_session::as_ref(copy).clone().allow_unknown()); ++ assert!(vx_session::as_ref(session).allows_unknown()); ++ let fresh = vx_session_new(); ++ assert!(!vx_session::as_ref(fresh).allows_unknown()); ++ vx_session_free(fresh); + vx_session_free(session); + + vx_session_free(copy); + } + } +-} ++ ++ #[test] ++ #[cfg_attr(miri, ignore)] ++ fn test_configured_sessions_are_isolated_from_paimon_read_defaults() { ++ unsafe { ++ let custom_id = Id::new("paimon.session.isolation-test"); ++ let custom = vx_session_new_with(|session| { ++ let registry = session.arrays().registry().clone(); ++ let plugin = registry.items().next(); ++ if let Some(plugin) = plugin { ++ registry.register(custom_id, plugin); ++ } ++ session.allow_unknown() ++ }); ++ let reader = vx_paimon_read_session_new(); ++ let fresh = vx_session_new(); ++ assert!(vx_session::as_ref(custom).allows_unknown()); ++ assert!( ++ vx_session::as_ref(custom) ++ .arrays() ++ .registry() ++ .find(&custom_id) ++ .is_some() ++ ); ++ for handle in [reader, fresh] { ++ let session = vx_session::as_ref(handle); ++ assert!(!session.allows_unknown()); ++ assert!(session.arrays().registry().find(&custom_id).is_none()); ++ } ++ vx_session_free(custom); ++ vx_session_free(reader); ++ vx_session_free(fresh); ++ } ++ } ++ ++ #[test] ++ #[cfg_attr(miri, ignore)] ++ fn test_paimon_read_defaults_survive_handle_lifetimes() { ++ unsafe { ++ let first = vx_paimon_read_session_new(); ++ let second = vx_paimon_read_session_new(); ++ assert_ne!(first, second); ++ let expected = vx_session::as_ref(first) ++ .arrays() ++ .registry() ++ .ids() ++ .collect::>(); ++ vx_session_free(first); ++ let actual = vx_session::as_ref(second) ++ .arrays() ++ .registry() ++ .ids() ++ .collect::>(); ++ assert_eq!(actual, expected); ++ assert!(!actual.is_empty()); ++ drop(vx_session::as_ref(second).handle()); ++ vx_session_free(second); ++ } ++ } ++ ++ #[test] ++ #[cfg_attr(miri, ignore)] ++ fn test_paimon_read_defaults_support_concurrent_arrow_execution() -> VortexResult<()> { ++ let workers = (0..8) ++ .map(|_| { ++ thread::spawn(|| -> VortexResult<()> { ++ for _ in 0..16 { ++ // Keep a session reference after releasing its FFI handle, ++ // just as a file/data-source keeps its own session clone. ++ let session = unsafe { ++ let handle = vx_paimon_read_session_new(); ++ let session = vx_session::as_ref(handle).clone(); ++ vx_session_free(handle); ++ session ++ }; ++ let input = ++ PrimitiveArray::new(buffer![1_i32, 2, 3], Validity::NonNullable) ++ .into_array(); ++ let mut ctx = session.create_execution_ctx(); ++ let exported = ++ session ++ .arrow() ++ .execute_arrow(input.clone(), None, &mut ctx)?; ++ let imported = session ++ .arrow() ++ .from_arrow_array(exported, &Field::new("", DataType::Int32, false))?; ++ assert_arrays_eq!(input, imported); ++ } ++ Ok(()) ++ }) ++ }) ++ .collect::>(); ++ for worker in workers { ++ worker ++ .join() ++ .map_err(|_| vortex_err!("Paimon session worker panicked"))??; ++ } ++ Ok(()) ++ } ++} diff --git a/cmake_modules/vortex.diff b/cmake_modules/vortex.diff index 26aca1037..d1b088f8e 100644 --- a/cmake_modules/vortex.diff +++ b/cmake_modules/vortex.diff @@ -7,3 +7,74 @@ diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs mod array_iterator; mod binary; mod data_source; +diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs +--- a/vortex-ffi/src/scan.rs ++++ b/vortex-ffi/src/scan.rs +@@ -25,8 +25,11 @@ + use vortex::error::VortexResult; + use vortex::error::vortex_bail; + use vortex::error::vortex_ensure; ++use vortex::expr::merge; ++use vortex::expr::pack; + use vortex::expr::root; + use vortex::io::runtime::BlockingRuntime; ++use vortex::layout::layouts::row_idx::row_idx; + use vortex::layout::scan::arrow::RecordBatchIteratorAdapter; + use vortex::scan::DataSourceScan; + use vortex::scan::Partition; +@@ -112,4 +115,7 @@ + pub limit: u64, + /// If true, return in storage order. + pub ordered: bool, ++ /// If true, append a `_row_idx` UInt64 column carrying each returned row's original physical ++ /// row index, stable across `filter` and `selection`. The host strips it before use. ++ pub include_row_idx: bool, + } +@@ -182,5 +188,12 @@ + Ok(ScanRequest { +- projection, ++ // paimon-cpp: when include_row_idx is set, append a hidden `_row_idx` column so the host can ++ // recover each surviving row's original physical index after `filter`/`selection` drop rows. ++ // row_idx() is substituted with the real index array during the scan; the host strips it. ++ projection: if opts.include_row_idx { ++ merge([projection, pack([("_row_idx", row_idx())], false.into())]) ++ } else { ++ projection ++ }, + filter, + row_range, + selection, +@@ -356,3 +369,5 @@ +- let schema = dtype.to_arrow_schema()?; ++ // paimon-cpp: export standard (non-view) string/binary types so the C++ reader skips its ++ // per-row view->standard rebuild; execute_arrow casts to this target internally. ++ let schema = crate::callback_io::normalize_arrow_schema_views(dtype.to_arrow_schema()?); + let schema = Arc::new(schema); + let target = Field::new_struct("", schema.fields().clone(), false); +diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs +--- a/vortex-ffi/src/dtype.rs ++++ b/vortex-ffi/src/dtype.rs +@@ -344,7 +344,9 @@ + ) -> c_int { + try_or(err, 1, || { + let dtype = vx_dtype::as_ref(dtype); +- let arrow_schema = dtype.to_arrow_schema()?; ++ // paimon-cpp: report standard (non-view) string/binary types so the exported schema matches ++ // the data that vx_partition_scan_arrow now produces (see callback_io normalize). ++ let arrow_schema = crate::callback_io::normalize_arrow_schema_views(dtype.to_arrow_schema()?); + let arrow_schema = FFI_ArrowSchema::try_from(&arrow_schema)?; + unsafe { ptr::write(schema, arrow_schema) }; + Ok(0) +diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h +--- a/vortex-ffi/cinclude/vortex.h ++++ b/vortex-ffi/cinclude/vortex.h +@@ -622,4 +622,9 @@ + * If true, return in storage order. + */ + bool ordered; ++ /** ++ * If true, append a `_row_idx` UInt64 column carrying each returned row's original physical ++ * row index, stable across filter and selection. The host strips it before use. ++ */ ++ bool include_row_idx; + } vx_scan_options; diff --git a/crates/vortex_callback_io/callback_io.rs b/crates/vortex_callback_io/callback_io.rs index f4d5fb857..2c49e19ba 100644 --- a/crates/vortex_callback_io/callback_io.rs +++ b/crates/vortex_callback_io/callback_io.rs @@ -46,6 +46,8 @@ //! `#[repr(C)]` struct, a callback `type`, or a `#[no_mangle]` entry point here MUST be mirrored in //! that header in the same change. +use std::ffi::CStr; +use std::ffi::c_char; use std::ffi::c_void; use std::io; use std::sync::Arc; @@ -62,10 +64,14 @@ use vortex::array::stream::ArrayStreamAdapter; use vortex::buffer::Alignment; use vortex::buffer::ByteBufferMut; use vortex::dtype::DType; +use vortex::dtype::Nullability; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; +use vortex::extension::datetime::Date; +use vortex::extension::datetime::TimeUnit; +use vortex::extension::datetime::Timestamp; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; @@ -78,6 +84,9 @@ use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Handle; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::scalar::PValue; +use vortex::scalar::Scalar; +use vortex::scalar::ScalarValue; use crate::RUNTIME; use crate::array::vx_array; @@ -86,6 +95,7 @@ use crate::dtype::vx_dtype; use crate::error::try_or; use crate::error::try_or_default; use crate::error::vx_error; +use crate::scalar::vx_scalar; use crate::session::vx_session; /// Fill `length` bytes starting at `offset` into `dst`. @@ -175,6 +185,7 @@ struct CallbackReadAt { len: u64, handle: Handle, concurrency: usize, + coalesce_config: CoalesceConfig, } /// Number of concurrent host read callbacks Vortex may have in flight for one file. @@ -185,9 +196,7 @@ const DEFAULT_CONCURRENCY: usize = 16; impl VortexReadAt for CallbackReadAt { fn coalesce_config(&self) -> Option { - // Each read costs an FFI hop plus a blocking-pool hand-off, and the host filesystem may be - // remote, so favor fewer and larger reads. - Some(CoalesceConfig::object_storage()) + Some(self.coalesce_config) } fn concurrency(&self) -> usize { @@ -254,6 +263,7 @@ unsafe fn data_source_new_callback( session: *const vx_session, callbacks: vx_input_callbacks, size: u64, + coalesce_config: CoalesceConfig, ) -> VortexResult<*const vx_data_source> { // Take ownership of the host context before anything else, so an early error below still // releases it exactly once when this `Arc` is dropped. @@ -263,6 +273,14 @@ unsafe fn data_source_new_callback( }); vortex_ensure!(!session.is_null()); + vortex_ensure!( + coalesce_config.max_size > 0, + "coalesced max_size must be positive" + ); + vortex_ensure!( + coalesce_config.distance <= coalesce_config.max_size, + "coalesced max_gap must not exceed max_size" + ); let read_at_fn = callbacks .read_at_fn .ok_or_else(|| vortex_err!("vx_input_callbacks.read_at_fn is required"))?; @@ -274,6 +292,7 @@ unsafe fn data_source_new_callback( len: size, handle: session.handle(), concurrency: DEFAULT_CONCURRENCY, + coalesce_config, }); let file = RUNTIME.block_on(async { session.open_options().open(reader).await })?; @@ -302,7 +321,36 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_callback( err: *mut *mut vx_error, ) -> *const vx_data_source { try_or(err, std::ptr::null(), || unsafe { - data_source_new_callback(session, callbacks, size) + data_source_new_callback(session, callbacks, size, CoalesceConfig::object_storage()) + }) +} + +/// Create a callback data source with an explicit I/O coalescing policy. +/// +/// The original entry point retains its object-storage defaults. Low-latency hosts can choose a +/// smaller gap to avoid reading unrelated bytes between sparse requests. `max_size` must be +/// positive and `max_gap` must not exceed it. The callbacks are consumed even on invalid options. +/// +/// # Safety +/// +/// Same requirements as [`vx_data_source_new_callback`]. `err` must be null or point to a null +/// `*mut vx_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_new_callback_with_coalesce( + session: *const vx_session, + callbacks: vx_input_callbacks, + size: u64, + max_gap: u64, + max_size: u64, + err: *mut *mut vx_error, +) -> *const vx_data_source { + try_or(err, std::ptr::null(), || unsafe { + data_source_new_callback( + session, + callbacks, + size, + CoalesceConfig::new(max_gap, max_size), + ) }) } @@ -523,6 +571,124 @@ pub unsafe extern "C-unwind" fn vx_callback_sink_abort(sink: *mut vx_callback_si drop(unsafe { Box::from_raw(sink) }); } +/// Export standard Arrow string/binary types recursively so the declared schema matches the +/// buffer layout expected by Paimon's C++ reader. +pub(crate) fn normalize_arrow_schema_views(schema: arrow_schema::Schema) -> arrow_schema::Schema { + let fields = schema + .fields() + .iter() + .cloned() + .map(normalize_field_views) + .collect::(); + arrow_schema::Schema::new_with_metadata(fields, schema.metadata.clone()) +} + +fn normalize_field_views(field: arrow_schema::FieldRef) -> arrow_schema::FieldRef { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(normalize_dtype_views(field.data_type())), + ) +} + +fn normalize_dtype_views(dtype: &arrow_schema::DataType) -> arrow_schema::DataType { + use arrow_schema::DataType; + match dtype { + DataType::Utf8View => DataType::Utf8, + DataType::BinaryView => DataType::Binary, + DataType::Struct(fields) => { + DataType::Struct(fields.iter().cloned().map(normalize_field_views).collect()) + } + DataType::List(field) => DataType::List(normalize_field_views(field.clone())), + DataType::LargeList(field) => DataType::LargeList(normalize_field_views(field.clone())), + DataType::FixedSizeList(field, size) => { + DataType::FixedSizeList(normalize_field_views(field.clone()), *size) + } + other => other.clone(), + } +} + +/// Time unit of a timestamp scalar literal. Mirrored in `vortex_ffi.h`; the discriminants are an +/// ABI contract, so do not reorder or renumber them. +#[repr(C)] +#[allow(non_camel_case_types, dead_code)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum vx_time_unit { + /// Seconds since the Unix epoch. + VX_TIME_UNIT_SECOND = 0, + /// Milliseconds since the Unix epoch. + VX_TIME_UNIT_MILLISECOND = 1, + /// Microseconds since the Unix epoch. + VX_TIME_UNIT_MICROSECOND = 2, + /// Nanoseconds since the Unix epoch. + VX_TIME_UNIT_NANOSECOND = 3, +} + +/// Create a timestamp scalar literal: `value` units since the Unix epoch under `unit`. +/// +/// `timezone` is a null-terminated UTF-8 IANA time-zone name, or null for a local (wall-clock) +/// timestamp. Vortex's own C ABI has no timestamp/date scalar constructor, so paimon-cpp adds these +/// two here to push down temporal predicates. +/// +/// # Safety +/// +/// `timezone`, when not null, must point to a valid null-terminated UTF-8 string that outlives this +/// call. `err` must be null or point to a null `*mut vx_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scalar_new_timestamp( + value: i64, + unit: vx_time_unit, + timezone: *const c_char, + is_nullable: bool, + err: *mut *mut vx_error, +) -> *mut vx_scalar { + try_or(err, std::ptr::null_mut(), || { + let time_unit = match unit { + vx_time_unit::VX_TIME_UNIT_SECOND => TimeUnit::Seconds, + vx_time_unit::VX_TIME_UNIT_MILLISECOND => TimeUnit::Milliseconds, + vx_time_unit::VX_TIME_UNIT_MICROSECOND => TimeUnit::Microseconds, + vx_time_unit::VX_TIME_UNIT_NANOSECOND => TimeUnit::Nanoseconds, + }; + let tz = if timezone.is_null() { + None + } else { + let name = unsafe { CStr::from_ptr(timezone) } + .to_str() + .map_err(|e| vortex_err!("timezone is not valid utf8: {e}"))?; + Some(Arc::from(name)) + }; + let dtype = DType::Extension( + Timestamp::new_with_tz(time_unit, tz, Nullability::from(is_nullable)).erased(), + ); + Ok(vx_scalar::new(Scalar::try_new( + dtype, + Some(ScalarValue::Primitive(PValue::I64(value))), + )?)) + }) +} + +/// Create a date scalar literal: `days` days since the Unix epoch (1970-01-01). +/// +/// # Safety +/// +/// `err` must be null or point to a null `*mut vx_error`. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scalar_new_date( + days: i32, + is_nullable: bool, + err: *mut *mut vx_error, +) -> *mut vx_scalar { + try_or(err, std::ptr::null_mut(), || { + let dtype = + DType::Extension(Date::new(TimeUnit::Days, Nullability::from(is_nullable)).erased()); + Ok(vx_scalar::new(Scalar::try_new( + dtype, + Some(ScalarValue::Primitive(PValue::I32(days))), + )?)) + }) +} + #[cfg(test)] #[cfg(unix)] mod tests { @@ -531,6 +697,7 @@ mod tests { use std::ptr; use std::sync::Mutex; use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::NamedTempFile; @@ -557,6 +724,144 @@ mod tests { use crate::sink::vx_array_sink_open_file; use crate::sink::vx_array_sink_push; + #[test] + fn scalar_new_timestamp_and_date_build_extension_scalars() { + use crate::scalar::vx_scalar_free; + + unsafe { + let mut error = ptr::null_mut(); + let ts = vx_scalar_new_timestamp( + 1_700_000_000_123, + vx_time_unit::VX_TIME_UNIT_MILLISECOND, + ptr::null(), + false, + &raw mut error, + ); + assert!(error.is_null()); + assert!(!ts.is_null()); + let DType::Extension(ext) = vx_scalar::as_ref(ts).dtype() else { + panic!("timestamp scalar must carry an extension dtype"); + }; + assert!(ext.is::()); + vx_scalar_free(ts); + + let tz = std::ffi::CString::new("America/New_York").unwrap(); + let ts_tz = vx_scalar_new_timestamp( + 0, + vx_time_unit::VX_TIME_UNIT_SECOND, + tz.as_ptr(), + true, + &raw mut error, + ); + assert!(error.is_null()); + assert!(!ts_tz.is_null()); + vx_scalar_free(ts_tz); + + let date = vx_scalar_new_date(19_000, false, &raw mut error); + assert!(error.is_null()); + assert!(!date.is_null()); + let DType::Extension(ext) = vx_scalar::as_ref(date).dtype() else { + panic!("date scalar must carry an extension dtype"); + }; + assert!(ext.is::()); + vx_scalar_free(date); + } + } + + #[test] + fn normalize_arrow_schema_views_rewrites_view_types_recursively() { + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use std::sync::Arc; + + let schema = Schema::new(vec![ + Field::new("plain_utf8", DataType::Utf8, false), + Field::new("utf8_view", DataType::Utf8View, true), + Field::new("binary_view", DataType::BinaryView, true), + Field::new( + "struct_col", + DataType::Struct( + vec![ + Field::new("inner_view", DataType::Utf8View, true), + Field::new("inner_int", DataType::Int32, false), + ] + .into(), + ), + true, + ), + Field::new( + "list_view", + DataType::List(Arc::new(Field::new("item", DataType::Utf8View, true))), + true, + ), + Field::new( + "fsl_view", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::BinaryView, true)), + 3, + ), + true, + ), + ]); + + let normalized = normalize_arrow_schema_views(schema); + + assert_eq!( + normalized + .field_with_name("plain_utf8") + .unwrap() + .data_type(), + &DataType::Utf8 + ); + assert_eq!( + normalized.field_with_name("utf8_view").unwrap().data_type(), + &DataType::Utf8 + ); + assert_eq!( + normalized + .field_with_name("binary_view") + .unwrap() + .data_type(), + &DataType::Binary + ); + match normalized + .field_with_name("struct_col") + .unwrap() + .data_type() + { + DataType::Struct(fields) => { + assert_eq!(fields[0].data_type(), &DataType::Utf8); + assert_eq!(fields[1].data_type(), &DataType::Int32); + } + other => panic!("expected Struct, got {:?}", other), + } + match normalized.field_with_name("list_view").unwrap().data_type() { + DataType::List(field) => assert_eq!(field.data_type(), &DataType::Utf8), + other => panic!("expected List, got {:?}", other), + } + match normalized.field_with_name("fsl_view").unwrap().data_type() { + DataType::FixedSizeList(field, size) => { + assert_eq!(field.data_type(), &DataType::Binary); + assert_eq!(*size, 3); + } + other => panic!("expected FixedSizeList, got {:?}", other), + } + assert_eq!(normalized.fields().len(), 6); + assert!( + normalized + .field_with_name("utf8_view") + .unwrap() + .is_nullable() + ); + assert!( + !normalized + .field_with_name("plain_utf8") + .unwrap() + .is_nullable() + ); + } + /// Host context standing in for paimon's `InputStream`: positional reads against a file. struct FileContext { file: File, @@ -599,12 +904,8 @@ mod tests { let mut error = ptr::null_mut(); let c_path = std::ffi::CString::new(path.clone()).unwrap(); - let sink = vx_array_sink_open_file( - session, - c_path.as_ptr(), - vx_dtype_ptr, - &raw mut error, - ); + let sink = + vx_array_sink_open_file(session, c_path.as_ptr(), vx_dtype_ptr, &raw mut error); assert!(error.is_null()); let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); let vx_array_ptr = vx_array::new(Arc::new(array.into_array())); @@ -615,25 +916,296 @@ mod tests { vx_array_free(vx_array_ptr); vx_dtype_free(vx_dtype_ptr); - let file = File::open(&path).unwrap(); - let size = file.metadata().unwrap().len(); - let context = Box::into_raw(Box::new(FileContext { file })); - let callbacks = vx_input_callbacks { - ctx: context.cast::(), - read_at_fn: Some(file_read_at), - release_fn: Some(file_release), + for local in [false, true] { + let file = File::open(&path).unwrap(); + let size = file.metadata().unwrap().len(); + let context = Box::into_raw(Box::new(FileContext { file })); + let callbacks = vx_input_callbacks { + ctx: context.cast::(), + read_at_fn: Some(file_read_at), + release_fn: Some(file_release), + }; + + let data_source = if local { + vx_data_source_new_callback_with_coalesce( + session, + callbacks, + size, + 64 << 10, + 4 << 20, + &raw mut error, + ) + } else { + vx_data_source_new_callback(session, callbacks, size, &raw mut error) + }; + assert!(error.is_null()); + assert!(!data_source.is_null()); + + let mut row_count = vx_estimate::default(); + vx_data_source_get_row_count(data_source, &raw mut row_count); + assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); + assert_eq!(row_count.estimate, 3); + + vx_data_source_free(data_source); + } + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn callback_coalesce_preserves_bytes_bounds_and_context_lifetime() { + struct MemoryContext { + bytes: Vec, + reads: AtomicUsize, + releases: AtomicUsize, + } + + unsafe extern "C" fn read_memory( + ctx: *mut c_void, + offset: u64, + dst: *mut u8, + length: usize, + ) -> i32 { + // SAFETY: target owns the boxed Arc, and CallbackReadAt validates the range first. + let context = unsafe { &*ctx.cast::>() }; + context.reads.fetch_add(1, Ordering::SeqCst); + let bytes = &context.bytes[offset as usize..offset as usize + length]; + unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, length) }; + 0 + } + + unsafe extern "C" fn release_memory(ctx: *mut c_void) { + // SAFETY: CallbackTarget releases this boxed Arc exactly once after all readers. + let context = unsafe { Box::from_raw(ctx.cast::>()) }; + context.releases.fetch_add(1, Ordering::SeqCst); + } + + for config in [ + CoalesceConfig::object_storage(), + CoalesceConfig::new(64 << 10, 4 << 20), + ] { + let context = Arc::new(MemoryContext { + bytes: (0..200_000).map(|i| (i % 251) as u8).collect(), + reads: AtomicUsize::new(0), + releases: AtomicUsize::new(0), + }); + let reader = Arc::new(CallbackReadAt { + target: Arc::new(CallbackTarget { + ctx: Box::into_raw(Box::new(Arc::clone(&context))).cast::(), + release_fn: Some(release_memory), + }), + read_at_fn: read_memory, + len: context.bytes.len() as u64, + handle: RUNTIME.handle(), + concurrency: DEFAULT_CONCURRENCY, + coalesce_config: config, + }); + let actual_config = reader.coalesce_config().unwrap(); + assert_eq!(actual_config.distance, config.distance); + assert_eq!(actual_config.max_size, config.max_size); + RUNTIME.block_on(async { + for (offset, length) in [(0, 1), (63_999, 4097), (200_000, 0)] { + let buffer = reader + .read_at(offset, length, Alignment::none()) + .await + .unwrap() + .to_host() + .await; + assert_eq!( + buffer.as_ref(), + &context.bytes[offset as usize..offset as usize + length] + ); + } + assert!(reader.read_at(199_999, 2, Alignment::none()).await.is_err()); + assert!( + reader + .read_at(u64::MAX, 1, Alignment::none()) + .await + .is_err() + ); + }); + assert_eq!(context.reads.load(Ordering::SeqCst), 2); + assert_eq!(context.releases.load(Ordering::SeqCst), 0); + drop(reader); + assert_eq!(context.releases.load(Ordering::SeqCst), 1); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn cancelled_in_flight_read_releases_context_after_callback_returns() { + use std::sync::Condvar; + use std::sync::mpsc; + use std::time::Duration; + + use futures::channel::oneshot; + + struct BlockingContext { + entered: Mutex>>, + released: mpsc::Sender<()>, + proceed: Mutex, + wake: Condvar, + completed: AtomicBool, + released_while_active: AtomicBool, + releases: AtomicUsize, + status: i32, + } + + unsafe extern "C" fn read_blocking( + ctx: *mut c_void, + _offset: u64, + dst: *mut u8, + length: usize, + ) -> i32 { + let context = unsafe { &*ctx.cast::>() }; + if let Some(entered) = context.entered.lock().unwrap().take() { + let _ = entered.send(()); + } + let (_proceed, timeout) = context + .wake + .wait_timeout_while( + context.proceed.lock().unwrap(), + Duration::from_secs(20), + |go| !*go, + ) + .unwrap(); + if !timeout.timed_out() && context.status == 0 { + unsafe { ptr::write_bytes(dst, 0x5a, length) }; + } + context.completed.store(true, Ordering::SeqCst); + if timeout.timed_out() { + -99 + } else { + context.status + } + } + + unsafe extern "C" fn release_blocking(ctx: *mut c_void) { + let context = unsafe { Box::from_raw(ctx.cast::>()) }; + context + .released_while_active + .store(!context.completed.load(Ordering::SeqCst), Ordering::SeqCst); + context.releases.fetch_add(1, Ordering::SeqCst); + let _ = context.released.send(()); + } + + for status in [0, -17] { + let (entered_tx, entered_rx) = oneshot::channel(); + let (released_tx, released_rx) = mpsc::channel(); + let context = Arc::new(BlockingContext { + entered: Mutex::new(Some(entered_tx)), + released: released_tx, + proceed: Mutex::new(false), + wake: Condvar::new(), + completed: AtomicBool::new(false), + released_while_active: AtomicBool::new(false), + releases: AtomicUsize::new(0), + status, + }); + let reader = CallbackReadAt { + target: Arc::new(CallbackTarget { + ctx: Box::into_raw(Box::new(Arc::clone(&context))).cast::(), + release_fn: Some(release_blocking), + }), + read_at_fn: read_blocking, + len: 16, + handle: RUNTIME.handle(), + concurrency: DEFAULT_CONCURRENCY, + coalesce_config: CoalesceConfig::object_storage(), }; + RUNTIME.block_on(async { + let future = reader.read_at(0, 16, Alignment::none()); + let pending = RUNTIME.handle().spawn(future); + entered_rx.await.unwrap(); + drop(pending); + drop(reader); + assert_eq!(context.releases.load(Ordering::SeqCst), 0); + assert!(!context.completed.load(Ordering::SeqCst)); + *context.proceed.lock().unwrap() = true; + context.wake.notify_all(); + RUNTIME + .handle() + .spawn_blocking(move || released_rx.recv_timeout(Duration::from_secs(20))) + .await + .unwrap(); + }); + assert!(context.completed.load(Ordering::SeqCst)); + assert!(!context.released_while_active.load(Ordering::SeqCst)); + assert_eq!(context.releases.load(Ordering::SeqCst), 1); + } + } - let data_source = vx_data_source_new_callback(session, callbacks, size, &raw mut error); - assert!(error.is_null()); - assert!(!data_source.is_null()); + #[test] + #[cfg_attr(miri, ignore)] + fn callback_read_error_is_propagated_and_context_released_once() { + unsafe extern "C" fn fail_read( + _ctx: *mut c_void, + _offset: u64, + _dst: *mut u8, + _length: usize, + ) -> i32 { + -31 + } - let mut row_count = vx_estimate::default(); - vx_data_source_get_row_count(data_source, &raw mut row_count); - assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); - assert_eq!(row_count.estimate, 3); + unsafe extern "C" fn release_counter(ctx: *mut c_void) { + let released = unsafe { Box::from_raw(ctx.cast::>()) }; + released.fetch_add(1, Ordering::SeqCst); + } - vx_data_source_free(data_source); + let released = Arc::new(AtomicUsize::new(0)); + let reader = CallbackReadAt { + target: Arc::new(CallbackTarget { + ctx: Box::into_raw(Box::new(Arc::clone(&released))).cast::(), + release_fn: Some(release_counter), + }), + read_at_fn: fail_read, + len: 16, + handle: RUNTIME.handle(), + concurrency: DEFAULT_CONCURRENCY, + coalesce_config: CoalesceConfig::object_storage(), + }; + let error = RUNTIME + .block_on(reader.read_at(0, 16, Alignment::none())) + .unwrap_err(); + assert!(error.to_string().contains("status -31")); + assert_eq!(released.load(Ordering::SeqCst), 0); + drop(reader); + assert_eq!(released.load(Ordering::SeqCst), 1); + } + + #[test] + #[cfg_attr(miri, ignore)] + fn invalid_coalesce_releases_callback_context() { + static RELEASED: AtomicUsize = AtomicUsize::new(0); + + unsafe extern "C" fn mark_released(_ctx: *mut c_void) { + RELEASED.fetch_add(1, Ordering::SeqCst); + } + + unsafe { + let session = vx_session_new(); + for (max_gap, max_size) in [(0, 0), (1025, 1024)] { + let mut error = ptr::null_mut(); + let previous_releases = RELEASED.load(Ordering::SeqCst); + let callbacks = vx_input_callbacks { + ctx: ptr::null_mut(), + read_at_fn: Some(file_read_at), + release_fn: Some(mark_released), + }; + let data_source = vx_data_source_new_callback_with_coalesce( + session, + callbacks, + 0, + max_gap, + max_size, + &raw mut error, + ); + assert!(data_source.is_null()); + assert!(!error.is_null()); + vx_error_free(error); + assert_eq!(RELEASED.load(Ordering::SeqCst), previous_releases + 1); + } vx_session_free(session); } } diff --git a/docs/source/user_guide/read.rst b/docs/source/user_guide/read.rst index a94776353..ff1a9aa00 100644 --- a/docs/source/user_guide/read.rst +++ b/docs/source/user_guide/read.rst @@ -41,6 +41,26 @@ cross-language task scheduling and interaction (e.g., Java and C++), substantial across the two language ecosystems. +Vortex Read Options +------------------- + +For Vortex files, pass these options through ``ReadContextBuilder::AddOption``: + +- ``vortex.read.membership-pushdown`` accepts ``auto`` (default), ``force`` or ``off``. + ``auto`` limits each IN/NOT IN expression to 64 comparison and boolean nodes after + deduplication and integer-range compression. ``force`` bypasses this budget while + retaining type and value safety checks. ``off`` disables only IN/NOT IN pushdown. +- ``vortex.read.io-coalesce`` accepts ``object-storage`` (default) or ``local``. + The default retains Vortex's object-storage coalescing policy. ``local`` uses a + 64 KiB maximum gap and a 4 MiB coalescing limit for low-latency storage. + Choose the mode explicitly; it is not selected from the file path. + +Native pushdown may return candidate rows requiring further filtering. For exact +results, retain the full residual predicate: enable ``EnablePredicateFilter(true)`` +with ``TableRead``, or apply equivalent filtering in the caller. ``force`` does not +bypass this requirement. + + Schema Evolution ----------------------- Scope and Compatibility diff --git a/src/paimon/format/vortex/CMakeLists.txt b/src/paimon/format/vortex/CMakeLists.txt index a03a5842b..99ed3e62a 100644 --- a/src/paimon/format/vortex/CMakeLists.txt +++ b/src/paimon/format/vortex/CMakeLists.txt @@ -25,6 +25,7 @@ if(PAIMON_ENABLE_VORTEX) vortex_format_writer.cpp vortex_ffi_util.cpp vortex_io_callbacks.cpp + vortex_predicate_converter.cpp vortex_stats_extractor.cpp vortex_writer_builder.cpp) @@ -51,6 +52,7 @@ if(PAIMON_ENABLE_VORTEX) add_paimon_test(vortex_format_test SOURCES vortex_file_format_test.cpp + vortex_predicate_converter_test.cpp STATIC_LINK_LIBS paimon_shared test_utils_static diff --git a/src/paimon/format/vortex/vortex_ffi.h b/src/paimon/format/vortex/vortex_ffi.h index d573d1279..f14e71071 100644 --- a/src/paimon/format/vortex/vortex_ffi.h +++ b/src/paimon/format/vortex/vortex_ffi.h @@ -49,6 +49,11 @@ extern "C" { // crates/vortex_callback_io/callback_io.rs); keep in sync with the Rust definitions. extern "C" { +/// Create an owned handle to the shared, fixed-runtime Paimon callback-reader session. +/// Free it with vx_session_free. Do not configure/register against this session or use it +/// for writers/multi-file sources; those require the fresh vx_session_new API. +vx_session* vx_paimon_read_session_new(); + /// Fill `length` bytes starting at `offset` into `dst`. Returns 0 on success and non-zero on /// failure; a partial read must be reported as a failure. Only the status code crosses the /// boundary, so the callback is expected to keep its error detail on `ctx`. @@ -76,6 +81,13 @@ const vx_data_source* vx_data_source_new_callback(const vx_session* session, vx_input_callbacks callbacks, uint64_t size, vx_error** err); +/// Create a callback data source with explicit I/O coalescing. max_size must be positive and +/// max_gap must not exceed max_size. The callback context is released exactly once, even on error. +const vx_data_source* vx_data_source_new_callback_with_coalesce(const vx_session* session, + vx_input_callbacks callbacks, + uint64_t size, uint64_t max_gap, + uint64_t max_size, vx_error** err); + /// Append `length` bytes from `src` to the host sink. Returns 0 on success and non-zero on failure; /// a partial write must be reported as a failure. Writes for one sink are sequential, never /// concurrent. @@ -112,4 +124,21 @@ void vx_callback_sink_close(vx_callback_sink* sink, vx_error** error_out); /// Discard the sink without writing a footer, leaving whatever the host received invalid. void vx_callback_sink_abort(vx_callback_sink* sink); + +/// Timestamp literal unit. Values must match vx_time_unit in callback_io.rs. +enum vx_time_unit { + VX_TIME_UNIT_SECOND = 0, + VX_TIME_UNIT_MILLISECOND = 1, + VX_TIME_UNIT_MICROSECOND = 2, + VX_TIME_UNIT_NANOSECOND = 3, +}; + +/// Create a timestamp literal in `unit` since the Unix epoch. `timezone` is an IANA zone name, +/// or null for wall-clock time. Returns null and sets `err` on failure; free with vx_scalar_free. +vx_scalar* vx_scalar_new_timestamp(int64_t value, vx_time_unit unit, const char* timezone, + bool is_nullable, vx_error** err); + +/// Create a date scalar literal: `days` days since the Unix epoch (1970-01-01). Returns null and +/// sets `err` on failure; free the result with `vx_scalar_free`. +vx_scalar* vx_scalar_new_date(int32_t days, bool is_nullable, vx_error** err); } diff --git a/src/paimon/format/vortex/vortex_ffi_util.h b/src/paimon/format/vortex/vortex_ffi_util.h index 6172b3175..a9c320060 100644 --- a/src/paimon/format/vortex/vortex_ffi_util.h +++ b/src/paimon/format/vortex/vortex_ffi_util.h @@ -40,5 +40,7 @@ using VxScanPtr = std::unique_ptr; using VxPartitionPtr = std::unique_ptr; using VxDtypePtr = std::unique_ptr; using VxArrayPtr = std::unique_ptr; +using VxExpressionPtr = std::unique_ptr; +using VxScalarPtr = std::unique_ptr; } // namespace paimon::vortex diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.cpp b/src/paimon/format/vortex/vortex_file_batch_reader.cpp index 6fa537b9b..b5791fc49 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.cpp +++ b/src/paimon/format/vortex/vortex_file_batch_reader.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/c/bridge.h" @@ -37,136 +38,16 @@ #include "paimon/common/utils/math.h" #include "paimon/core/utils/nested_projection_utils.h" #include "paimon/format/vortex/vortex_ffi_util.h" +#include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/fs/file_system.h" +#include "paimon/utils/roaring_bitmap32.h" namespace paimon::vortex { namespace { -// Vortex's Arrow export represents strings and binaries with the view types (StringView / -// BinaryView). paimon schemas use the standard string/binary types, so map the view types back, -// recursing through nested types. Types that contain no views are returned unchanged. -std::shared_ptr NormalizeViewType(const std::shared_ptr& type) { - switch (type->id()) { - case arrow::Type::STRING_VIEW: - return arrow::utf8(); - case arrow::Type::BINARY_VIEW: - return arrow::binary(); - case arrow::Type::STRUCT: { - arrow::FieldVector fields; - fields.reserve(type->num_fields()); - for (const std::shared_ptr& field : type->fields()) { - fields.push_back(field->WithType(NormalizeViewType(field->type()))); - } - return arrow::struct_(fields); - } - case arrow::Type::LIST: - return arrow::list(type->field(0)->WithType(NormalizeViewType(type->field(0)->type()))); - case arrow::Type::FIXED_SIZE_LIST: - return arrow::fixed_size_list( - type->field(0)->WithType(NormalizeViewType(type->field(0)->type())), - checked_cast(*type).list_size()); - default: - return type; - } -} - -std::shared_ptr NormalizeViewSchema(const std::shared_ptr& schema) { - arrow::FieldVector fields; - fields.reserve(schema->num_fields()); - for (const std::shared_ptr& field : schema->fields()) { - fields.push_back(field->WithType(NormalizeViewType(field->type()))); - } - return arrow::schema(fields, schema->metadata()); -} - -// Rebuilds `array` with every StringView/BinaryView column (which Vortex's Arrow export produces) -// converted to the standard string/binary type. Arrow 17 has no cast kernel between the view and -// standard types, so the leaves are rebuilt with builders and the parents reconstructed, recursing -// through structs and lists. Arrays with no view-typed data are returned unchanged. `array` must be -// offset-normalized (offset 0) so the rebuilt parents and children stay consistent. -Result> NormalizeViewArray(const std::shared_ptr& array, - arrow::MemoryPool* pool) { - // TODO(zhouyc-ali): avoid the per-row copy by exporting standard utf8/binary - switch (array->type_id()) { - case arrow::Type::STRING_VIEW: { - const auto& view = checked_cast(*array); - arrow::StringBuilder builder(pool); - for (int64_t i = 0; i < view.length(); ++i) { - if (view.IsNull(i)) { - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); - } else { - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(view.GetView(i))); - } - } - std::shared_ptr out; - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); - return out; - } - case arrow::Type::BINARY_VIEW: { - const auto& view = checked_cast(*array); - arrow::BinaryBuilder builder(pool); - for (int64_t i = 0; i < view.length(); ++i) { - if (view.IsNull(i)) { - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); - } else { - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(view.GetView(i))); - } - } - std::shared_ptr out; - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); - return out; - } - case arrow::Type::STRUCT: { - const auto& struct_array = checked_cast(*array); - arrow::ArrayVector children; - arrow::FieldVector fields; - children.reserve(struct_array.num_fields()); - fields.reserve(struct_array.num_fields()); - bool changed = false; - for (int32_t i = 0; i < struct_array.num_fields(); ++i) { - const std::shared_ptr& child = struct_array.field(i); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_child, - NormalizeViewArray(child, pool)); - changed = changed || normalized_child.get() != child.get(); - fields.push_back(array->type()->field(i)->WithType(normalized_child->type())); - children.push_back(std::move(normalized_child)); - } - if (!changed) { - return array; - } - return std::make_shared( - arrow::struct_(fields), struct_array.length(), children, struct_array.null_bitmap(), - struct_array.null_count(), struct_array.offset()); - } - case arrow::Type::LIST: { - const auto& list = checked_cast(*array); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, - NormalizeViewArray(list.values(), pool)); - if (values.get() == list.values().get()) { - return array; - } - return std::make_shared( - arrow::list(array->type()->field(0)->WithType(values->type())), list.length(), - list.value_offsets(), values, list.null_bitmap(), list.null_count(), list.offset()); - } - case arrow::Type::FIXED_SIZE_LIST: { - const auto& list = checked_cast(*array); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, - NormalizeViewArray(list.values(), pool)); - if (values.get() == list.values().get()) { - return array; - } - const auto& fsl_type = checked_cast(*array->type()); - return std::make_shared( - arrow::fixed_size_list(array->type()->field(0)->WithType(values->type()), - fsl_type.list_size()), - list.length(), values, list.null_bitmap(), list.null_count(), list.offset()); - } - default: - return array; - } -} +// Must match the physical row index column appended by vortex-ffi scan.rs. +constexpr char kRowIdxColumn[] = "_row_idx"; // Vortex reports the file's zone and paimon asks for the reader's; the instants are the same. std::shared_ptr AlignTimestampTimezones( @@ -236,23 +117,24 @@ VortexFileBatchReader::VortexFileBatchReader( const std::shared_ptr& input, int32_t batch_size, std::shared_ptr input_context, VxSessionPtr session, VxDataSourcePtr data_source, VxScanPtr scan, const std::shared_ptr& file_schema, - const std::shared_ptr& struct_type, uint64_t total_rows, - const std::shared_ptr& arrow_pool) + uint64_t total_rows, const std::shared_ptr& arrow_pool, + const VortexPredicateOptions& predicate_options) : input_(input), batch_size_(batch_size), input_context_(std::move(input_context)), session_(std::move(session)), data_source_(std::move(data_source)), scan_(std::move(scan)), - file_schema_(NormalizeViewSchema(file_schema)), - struct_type_(struct_type), + file_schema_(file_schema), total_rows_(total_rows), arrow_pool_(arrow_pool), + predicate_options_(predicate_options), metrics_(std::make_shared()) {} Result> VortexFileBatchReader::Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& arrow_pool) { + const std::shared_ptr& arrow_pool, + const VortexPredicateOptions& predicate_options, VortexIoCoalesceMode coalesce_mode) { if (input == nullptr || arrow_pool == nullptr || batch_size <= 0) { return Status::Invalid( "Vortex reader requires non-null input and Arrow memory pool, and positive batch size"); @@ -262,14 +144,22 @@ Result> VortexFileBatchReader::Create( auto length = static_cast(signed_length); auto input_context = std::make_shared(input); - VxSessionPtr session(vx_session_new(), vx_session_free); + VxSessionPtr session(vx_paimon_read_session_new(), vx_session_free); if (session == nullptr) { return Status::IOError("failed to create Vortex session"); } vx_error* error = nullptr; + const bool local_coalesce = coalesce_mode == VortexIoCoalesceMode::LOCAL; + // Smaller gaps limit read amplification for explicitly selected low-latency sources. + constexpr uint64_t kLocalMaxGap = 64ULL << 10; + constexpr uint64_t kLocalMaxSize = 4ULL << 20; VxDataSourcePtr data_source( - vx_data_source_new_callback(session.get(), VortexInputContext::MakeCallbacks(input_context), - length, &error), + local_coalesce + ? vx_data_source_new_callback_with_coalesce( + session.get(), VortexInputContext::MakeCallbacks(input_context), length, + kLocalMaxGap, kLocalMaxSize, &error) + : vx_data_source_new_callback( + session.get(), VortexInputContext::MakeCallbacks(input_context), length, &error), vx_data_source_free); if (data_source == nullptr) { return VortexCallbackError("open Vortex callback data source", error, @@ -291,9 +181,6 @@ Result> VortexFileBatchReader::Create( } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_schema, arrow::ImportSchema(&ffi_schema)); - // Built from the raw view-typed schema on purpose: NextBatch imports each scanned batch as this - // exact type, while the constructor normalizes its own copy kept as file_schema_. - std::shared_ptr struct_type = arrow::struct_(file_schema->fields()); vx_estimate row_count{}; vx_data_source_get_row_count(data_source.get(), &row_count); @@ -306,7 +193,7 @@ Result> VortexFileBatchReader::Create( // scan may be consumed only once. return std::unique_ptr(new VortexFileBatchReader( input, batch_size, std::move(input_context), std::move(session), std::move(data_source), - VxScanPtr(nullptr, vx_scan_free), file_schema, struct_type, total_rows, arrow_pool)); + VxScanPtr(nullptr, vx_scan_free), file_schema, total_rows, arrow_pool, predicate_options)); } VortexFileBatchReader::~VortexFileBatchReader() { @@ -314,12 +201,67 @@ VortexFileBatchReader::~VortexFileBatchReader() { } Result VortexFileBatchReader::OpenNextPartitionStream() { + if (has_selection_ && selection_row_ids_.empty()) { + return false; + } if (scan_ == nullptr) { vx_error* error = nullptr; // Scan in storage order: Vortex defaults to ordered=false and may emit chunks out of order, // which would misalign the physical row positions NextBatch assigns. - vx_scan_options options{}; // zero-init: all columns, no filter/row-range/selection/limit + vx_scan_options options{}; options.ordered = true; + // Missing fields are null-filled by AlignToReadSchema and cannot be selected here. + // The scan clones the projection, so these expressions need only outlive scan creation. + VxExpressionPtr root(nullptr, vx_expression_free); + VxExpressionPtr projection(nullptr, vx_expression_free); + std::vector names; + std::vector name_ptrs; + if (read_schema_ != nullptr) { + names.reserve(read_schema_->num_fields()); + // Native merge requires distinct fields; output alignment restores repeated columns. + std::vector selected_columns(file_schema_->num_fields(), 0); + for (const std::shared_ptr& field : read_schema_->fields()) { + const int32_t field_index = file_schema_->GetFieldIndex(field->name()); + if (field_index >= 0 && selected_columns[field_index] == 0) { + names.push_back(field->name()); + selected_columns[field_index] = 1; + } + } + if (!names.empty()) { + name_ptrs.reserve(names.size()); + for (const std::string& name : names) { + name_ptrs.push_back(name.c_str()); + } + root = VxExpressionPtr(vx_expression_root(), vx_expression_free); + projection = VxExpressionPtr( + vx_expression_select(name_ptrs.data(), name_ptrs.size(), root.get()), + vx_expression_free); + options.projection = projection.get(); + } + } + // VX_SELECTION_INCLUDE_RANGE maps to IncludeByIndex, not a contiguous range. + // The scan copies the individual indices from selection_row_ids_. + if (has_selection_) { + options.selection.idx = selection_row_ids_.data(); + options.selection.idx_len = selection_row_ids_.size(); + options.selection.include = VX_SELECTION_INCLUDE_RANGE; + } + // Filtering drops rows; include_row_idx preserves their original physical positions. + // The scan clones the filter, so it can be freed after scan creation. + VxExpressionPtr filter(nullptr, vx_expression_free); + // The hidden index must not collide with a file field or a requested missing field. + const bool has_row_idx_column = + file_schema_->GetFieldIndex(kRowIdxColumn) >= 0 || + (read_schema_ != nullptr && !read_schema_->GetAllFieldIndices(kRowIdxColumn).empty()); + if (!has_row_idx_column) { + filter = VortexPredicateConverter::Convert(predicate_filter_, file_schema_, + predicate_options_); + } + if (filter != nullptr) { + options.filter = filter.get(); + options.include_row_idx = true; + has_filter_ = true; + } VxScanPtr scan( vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), vx_scan_free); @@ -347,6 +289,18 @@ Result VortexFileBatchReader::OpenNextPartitionStream() { } current_stream_ = stream; stream_active_ = true; + // Import batches with the scan's actual output schema (projected columns, standard types). + ::ArrowSchema ffi_schema = {}; + if (current_stream_.get_schema(¤t_stream_, &ffi_schema) != 0) { + Status callback_status = input_context_->GetCallbackStatus(); + ReleaseStream(); + if (!callback_status.ok()) { + return callback_status.WithMessage("read Vortex stream schema", ": ", + callback_status.message()); + } + return Status::IOError("Vortex Arrow stream get_schema failed"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(current_stream_schema_, arrow::ImportSchema(&ffi_schema)); return true; } @@ -379,8 +333,9 @@ Result> VortexFileBatchReader::ReadNextArray() { ReleaseStream(); // Partition exhausted; try the next one. continue; } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, - arrow::ImportArray(&ffi_array, struct_type_)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(&ffi_array, arrow::struct_(current_stream_schema_->fields()))); if (array->length() == 0) { continue; // Skip empty batches. } @@ -393,7 +348,6 @@ Result> VortexFileBatchReader::AlignToReadSchema( if (read_schema_ == nullptr) { return array; } - // TODO(zhouyc-ali): push top-level column pruning into the Vortex scan std::shared_ptr read_type = arrow::struct_(read_schema_->fields()); return NestedProjectionUtils::AlignArrayToReadType( arrow::MakeArray(AlignTimestampTimezones(array->data(), read_type)), read_type, @@ -419,12 +373,18 @@ Result VortexFileBatchReader::NextBatch() { current_batch_->Slice(current_batch_offset_, row_count); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_array, ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); - // Vortex exports strings/binaries as Arrow view types; convert to the standard ones (no cast - // kernel in Arrow 17) after offset normalization, so the array is offset-0. - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output_array, - NormalizeViewArray(normalized_array, arrow_pool_.get())); + if (has_filter_) { + // Capture physical row indices before AlignToReadSchema removes the hidden column. + const auto& struct_type = checked_cast(*normalized_array->type()); + const int32_t row_idx_pos = struct_type.num_fields() - 1; + if (row_idx_pos < 0 || struct_type.field(row_idx_pos)->name() != kRowIdxColumn) { + return Status::Invalid("Vortex scan did not append the _row_idx column"); + } + previous_batch_row_idx_ = + checked_cast(*normalized_array).field(row_idx_pos); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr projected_array, - AlignToReadSchema(output_array)); + AlignToReadSchema(normalized_array)); previous_first_row_ = rows_emitted_; previous_batch_row_count_ = static_cast(row_count); @@ -448,17 +408,33 @@ Result> VortexFileBatchReader::GetFileSchema() co Status VortexFileBatchReader::SetReadSchema( ::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) { + if (closed_) { + return Status::Invalid("Vortex reader is closed"); + } if (read_schema == nullptr) { return Status::Invalid("Vortex read schema is nullptr"); } - // TODO(zhouyc-ali): push the row selection into the Vortex scan - (void)selection_bitmap; - // No pushdown: the whole file is scanned and reading restarts from row 0, so drop the - // single-use scan; the next read re-creates it. + // Keep selected physical row ids to map output positions after selection pushdown. + std::vector selection_row_ids; + if (selection_bitmap) { + selection_row_ids.reserve(static_cast(selection_bitmap->Cardinality())); + for (auto iter = selection_bitmap->Begin(); iter != selection_bitmap->End(); ++iter) { + const auto row_id = static_cast(*iter); + if (row_id >= total_rows_) { + return Status::Invalid( + fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows_)); + } + selection_row_ids.push_back(row_id); + } + } + // Leave the active scan and its row mapping unchanged if validation or import fails. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_read_schema, arrow::ImportSchema(read_schema)); + // Recreate the single-use scan on the next read with the new projection and predicates. read_schema_ = std::move(imported_read_schema); predicate_filter_ = std::dynamic_pointer_cast(predicate); + has_selection_ = selection_bitmap.has_value(); + selection_row_ids_ = std::move(selection_row_ids); ReleaseStream(); current_batch_.reset(); current_batch_offset_ = 0; @@ -466,6 +442,8 @@ Status VortexFileBatchReader::SetReadSchema( rows_emitted_ = 0; previous_first_row_ = std::numeric_limits::max(); previous_batch_row_count_ = 0; + has_filter_ = false; + previous_batch_row_idx_.reset(); return Status::OK(); } @@ -479,7 +457,26 @@ Result VortexFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch 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; + // The hidden column accounts for both filtering and selection; emitted offsets do not. + if (has_filter_) { + if (previous_batch_row_idx_ == nullptr || + batch_row_id >= static_cast(previous_batch_row_idx_->length())) { + return Status::Invalid( + fmt::format("Vortex filter row index unavailable for row {}", batch_row_id)); + } + return checked_cast(*previous_batch_row_idx_) + .Value(static_cast(batch_row_id)); + } + const uint64_t emitted_offset = previous_first_row_ + batch_row_id; + // An ordered selection emits rows in the same order as selection_row_ids_. + if (has_selection_) { + if (emitted_offset >= selection_row_ids_.size()) { + return Status::Invalid(fmt::format("emitted row offset {} exceeds selection size {}", + emitted_offset, selection_row_ids_.size())); + } + return selection_row_ids_[emitted_offset]; + } + return emitted_offset; } Result VortexFileBatchReader::GetNumberOfRows() const { @@ -498,6 +495,7 @@ void VortexFileBatchReader::ReleaseStream() { if (stream_active_) { ArrowArrayStreamRelease(¤t_stream_); current_stream_ = ::ArrowArrayStream{}; + current_stream_schema_.reset(); stream_active_ = false; } } diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.h b/src/paimon/format/vortex/vortex_file_batch_reader.h index 949a9ae45..b0a1180e9 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.h +++ b/src/paimon/format/vortex/vortex_file_batch_reader.h @@ -22,11 +22,13 @@ #include #include #include +#include #include "arrow/type_fwd.h" #include "paimon/format/vortex/vortex_ffi.h" #include "paimon/format/vortex/vortex_ffi_util.h" #include "paimon/format/vortex/vortex_io_callbacks.h" +#include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/reader/file_batch_reader.h" #include "paimon/result.h" @@ -50,7 +52,9 @@ class VortexFileBatchReader : public FileBatchReader { public: static Result> Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& arrow_pool); + const std::shared_ptr& arrow_pool, + const VortexPredicateOptions& predicate_options = {}, + VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE); ~VortexFileBatchReader() override; @@ -63,16 +67,16 @@ class VortexFileBatchReader : public FileBatchReader { std::shared_ptr GetReaderMetrics() const override; void Close() override; bool SupportPreciseBitmapSelection() const override { - return false; + return true; } private: VortexFileBatchReader(const std::shared_ptr& input, int32_t batch_size, std::shared_ptr input_context, VxSessionPtr session, VxDataSourcePtr data_source, VxScanPtr scan, - const std::shared_ptr& file_schema, - const std::shared_ptr& struct_type, uint64_t total_rows, - const std::shared_ptr& arrow_pool); + const std::shared_ptr& file_schema, uint64_t total_rows, + const std::shared_ptr& arrow_pool, + const VortexPredicateOptions& predicate_options); /// Pull the next Arrow array from the current partition stream, opening the next partition's /// stream when the current one is exhausted. Returns nullptr at end of scan. @@ -96,20 +100,16 @@ class VortexFileBatchReader : public FileBatchReader { VxSessionPtr session_; VxDataSourcePtr data_source_; VxScanPtr scan_; - // The file schema with Vortex's string/binary view types normalized to the standard types; this - // is what GetFileSchema() reports to paimon. + // Full file schema with standard Arrow types, as reported by GetFileSchema(). std::shared_ptr file_schema_; - // The RAW view-typed struct that Vortex's Arrow export produces. NextBatch imports each scanned - // batch as this type (so it must match the C export's buffer layout), then NormalizeViewArray - // rewrites it to file_schema_'s types. Keeping this view-typed is load-bearing: normalizing it - // would make ImportArray misread the buffers. - std::shared_ptr struct_type_; // The schema NextBatch must output, set by SetReadSchema (the projected read schema). Null // until SetReadSchema is called, in which case the full file schema is returned. std::shared_ptr read_schema_; uint64_t total_rows_ = 0; ::ArrowArrayStream current_stream_{}; + // Actual scan output schema, including projection and any hidden physical row index. + std::shared_ptr current_stream_schema_; bool stream_active_ = false; std::shared_ptr current_batch_; int64_t current_batch_offset_ = 0; @@ -118,8 +118,18 @@ class VortexFileBatchReader : public FileBatchReader { uint64_t previous_first_row_ = std::numeric_limits::max(); uint64_t previous_batch_row_count_ = 0; + // Selected physical row ids in output order, used when selection is pushed without a filter. + bool has_selection_ = false; + std::vector selection_row_ids_; + + // A pushed filter returns physical row ids in a hidden _row_idx column. Retain the last + // batch's column because emitted offsets no longer identify physical rows after filtering. + bool has_filter_ = false; + std::shared_ptr previous_batch_row_idx_; + std::shared_ptr arrow_pool_; - std::shared_ptr predicate_filter_; // from SetReadSchema; not pushed down + VortexPredicateOptions predicate_options_; + std::shared_ptr predicate_filter_; std::shared_ptr metrics_; bool closed_ = false; }; diff --git a/src/paimon/format/vortex/vortex_file_format.cpp b/src/paimon/format/vortex/vortex_file_format.cpp index 9d85b2fb7..0cec183b3 100644 --- a/src/paimon/format/vortex/vortex_file_format.cpp +++ b/src/paimon/format/vortex/vortex_file_format.cpp @@ -20,7 +20,9 @@ #include "paimon/format/vortex/vortex_file_format.h" #include "arrow/c/bridge.h" +#include "fmt/format.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/format/vortex/vortex_reader_builder.h" #include "paimon/format/vortex/vortex_stats_extractor.h" #include "paimon/format/vortex/vortex_writer_builder.h" @@ -29,7 +31,33 @@ namespace paimon::vortex { Result> VortexFileFormat::CreateReaderBuilder( int32_t batch_size) const { - return std::make_unique(batch_size); + VortexPredicateOptions predicate_options; + const auto membership_mode = options_.find("vortex.read.membership-pushdown"); + if (membership_mode != options_.end()) { + if (membership_mode->second == "auto") { + predicate_options.membership_pushdown = MembershipPushdownMode::AUTO; + } else if (membership_mode->second == "force") { + predicate_options.membership_pushdown = MembershipPushdownMode::FORCE; + } else if (membership_mode->second == "off") { + predicate_options.membership_pushdown = MembershipPushdownMode::OFF; + } else { + return Status::Invalid( + fmt::format("vortex.read.membership-pushdown must be auto, force or off, got: {}", + membership_mode->second)); + } + } + VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE; + const auto coalesce = options_.find("vortex.read.io-coalesce"); + if (coalesce != options_.end()) { + if (coalesce->second == "local") { + coalesce_mode = VortexIoCoalesceMode::LOCAL; + } else if (coalesce->second != "object-storage") { + return Status::Invalid( + fmt::format("vortex.read.io-coalesce must be object-storage or local, got: {}", + coalesce->second)); + } + } + return std::make_unique(batch_size, predicate_options, coalesce_mode); } Result> VortexFileFormat::CreateWriterBuilder( diff --git a/src/paimon/format/vortex/vortex_file_format_test.cpp b/src/paimon/format/vortex/vortex_file_format_test.cpp index cb383f741..c1ddffccc 100644 --- a/src/paimon/format/vortex/vortex_file_format_test.cpp +++ b/src/paimon/format/vortex/vortex_file_format_test.cpp @@ -24,8 +24,11 @@ #include #include #include +#include +#include #include #include +#include #include #include #include @@ -35,13 +38,18 @@ #include "arrow/array/concatenate.h" #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" +#include "arrow/util/key_value_metadata.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/common/factories/io_hook.h" +#include "paimon/common/predicate/predicate_filter.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/common/utils/scope_guard.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" #include "paimon/format/column_stats.h" #include "paimon/format/file_format.h" #include "paimon/format/file_format_factory.h" @@ -53,11 +61,14 @@ #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/reader/file_batch_reader.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" namespace paimon::vortex::test { @@ -157,6 +168,105 @@ class VortexFileFormatTest : public ::testing::Test { } } + Result> CreatePredicateReader( + const std::string& path, const std::shared_ptr& schema, + const std::shared_ptr& predicate, const std::string& mode, + const std::string& io_coalesce_mode = "") const { + std::map options = {{"file.format", "vortex"}}; + if (!mode.empty()) { + options.emplace("vortex.read.membership-pushdown", mode); + } + if (!io_coalesce_mode.empty()) { + options.emplace("vortex.read.io-coalesce", io_coalesce_mode); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr format, + FileFormatFactory::Get("vortex", options)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr builder, + format->CreateReaderBuilder(/*batch_size=*/2)); + builder->WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, file_system_->Open(path)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr 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, /*selection_bitmap=*/std::nullopt)); + return reader; + } + + // This is the native file reader, before the SDK residual predicate. A declined membership + // predicate intentionally returns a superset, including null rows, for the caller to filter. + Result> CollectNativeRowIds( + std::unique_ptr reader) const { + std::vector row_ids; + 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 array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + for (int64_t i = 0; i < array->length(); ++i) { + PAIMON_ASSIGN_OR_RAISE(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + row_ids.push_back(row_id); + } + } + return row_ids; + } + + static std::shared_ptr ReaderTestSchema() { + return arrow::schema( + {arrow::field("id", arrow::int32(), false), arrow::field("name", arrow::utf8())}); + } + + Result WriteReaderTestFile() const { + const std::shared_ptr schema = ReaderTestSchema(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr data, + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(schema->fields()), + R"([[0,"zero"],[1,null],[2,"two"],[3,"three"],[4,null],[5,"five"]])")); + const std::string path = PathUtil::JoinPath(directory_->Str(), "reader.vortex"); + PAIMON_RETURN_NOT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + return path; + } + + Result> CreateReader(const std::string& path, + int32_t batch_size = 2) const { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr builder, + format_->CreateReaderBuilder(batch_size)); + builder->WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, file_system_->Open(path)); + return builder->Build(input); + } + + Status ConfigureReader(FileBatchReader* reader, const std::shared_ptr& schema, + const std::shared_ptr& predicate = nullptr, + const std::optional& selection = std::nullopt) const { + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + ScopeGuard release_schema([&ffi_schema]() { + if (ffi_schema.release != nullptr) { + ffi_schema.release(&ffi_schema); + } + }); + return reader->SetReadSchema(&ffi_schema, predicate, selection); + } + + Result> ReadNextBatch(FileBatchReader* reader) const { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return std::shared_ptr(); + } + PAIMON_RETURN_NOT_OK(paimon::test::ReadResultCollector::CheckBatchOffset(batch)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + return array; + } + protected: std::shared_ptr format_; std::shared_ptr file_system_; @@ -997,4 +1107,692 @@ TEST_F(VortexFileFormatTest, WriteSurfacesInjectedIOError) { ASSERT_TRUE(run_complete); } +// Selection pushdown must preserve physical row ids across output batches. +TEST_F(VortexFileFormatTest, SelectionPushdownReturnsSelectedRowsAndPhysicalRowIds) { + constexpr int64_t kRows = 10; + arrow::Int32Builder id_builder; + arrow::StringBuilder name_builder; + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(id_builder.Append(static_cast(i * 10)).ok()); + ASSERT_TRUE(name_builder.Append(fmt::format("row_{}", i)).ok()); + } + std::shared_ptr id_array; + std::shared_ptr name_array; + ASSERT_TRUE(id_builder.Finish(&id_array).ok()); + ASSERT_TRUE(name_builder.Finish(&name_array).ok()); + 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::StructArray::Make({id_array, name_array}, fields).ValueOrDie(); + + std::string path = PathUtil::JoinPath(directory_->Str(), "selection.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + // Select physical rows {1, 3, 5, 7}; their id values are {10, 30, 50, 70}. + RoaringBitmap32 selection = RoaringBitmap32::From({1, 3, 5, 7}); + + // batch_size=2 spreads the 4 selected rows over two batches, exercising the cross-batch row id. + 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(*schema, &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, selection)); + + std::vector got_ids; + std::vector got_row_ids; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + std::shared_ptr arr = + arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + std::shared_ptr id_col = checked_pointer_cast( + checked_pointer_cast(arr)->GetFieldByName("id")); + for (int64_t i = 0; i < arr->length(); ++i) { + got_ids.push_back(id_col->Value(i)); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + got_row_ids.push_back(row_id); + } + } + ASSERT_EQ(got_ids, (std::vector{10, 30, 50, 70})); + ASSERT_EQ(got_row_ids, (std::vector{1, 3, 5, 7})); +} + +TEST_F(VortexFileFormatTest, FilterPushdownReturnsMatchingRowsAndPhysicalRowIds) { + constexpr int64_t kRows = 10; + arrow::Int32Builder id_builder; + arrow::StringBuilder name_builder; + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(id_builder.Append(static_cast(i * 10)).ok()); + ASSERT_TRUE(name_builder.Append(fmt::format("row_{}", i)).ok()); + } + std::shared_ptr id_array; + std::shared_ptr name_array; + ASSERT_TRUE(id_builder.Finish(&id_array).ok()); + ASSERT_TRUE(name_builder.Finish(&name_array).ok()); + 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::StructArray::Make({id_array, name_array}, fields).ValueOrDie(); + + std::string path = PathUtil::JoinPath(directory_->Str(), "filter.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + // Push down `id > 30`: rows 4..9 (ids 40..90) survive and rows 0..3 are dropped by the scan. + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, "id", FieldType::INT, Literal(static_cast(30))); + + // batch_size=2 spreads the 6 surviving rows over three batches, exercising cross-batch row ids. + 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(*schema, &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + + std::vector got_ids; + std::vector got_row_ids; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + std::shared_ptr arr = + arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + std::shared_ptr id_col = checked_pointer_cast( + checked_pointer_cast(arr)->GetFieldByName("id")); + for (int64_t i = 0; i < arr->length(); ++i) { + got_ids.push_back(id_col->Value(i)); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + got_row_ids.push_back(row_id); + } + } + // Filtered rows retain their physical positions rather than output offsets. + ASSERT_EQ(got_ids, (std::vector{40, 50, 60, 70, 80, 90})); + ASSERT_EQ(got_row_ids, (std::vector{4, 5, 6, 7, 8, 9})); +} + +TEST_F(VortexFileFormatTest, RejectsInvalidMembershipPushdownMode) { + for (const std::string mode : {"", "AUTO", "enabled"}) { + SCOPED_TRACE(fmt::format("mode={}", mode)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr format, + FileFormatFactory::Get( + "vortex", {{"file.format", "vortex"}, {"vortex.read.membership-pushdown", mode}})); + Result> result = format->CreateReaderBuilder(2); + ASSERT_FALSE(result.ok()); + ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); + ASSERT_NE(result.status().ToString().find("vortex.read.membership-pushdown"), + std::string::npos); + } +} + +TEST_F(VortexFileFormatTest, RejectsInvalidIoCoalesceMode) { + for (const std::string mode : {"", "LOCAL", "auto", "none", "local "}) { + SCOPED_TRACE(fmt::format("mode={}", mode)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", {{"file.format", "vortex"}, + {"vortex.read.io-coalesce", mode}})); + Result> result = format->CreateReaderBuilder(2); + ASSERT_FALSE(result.ok()); + ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); + ASSERT_NE(result.status().ToString().find("vortex.read.io-coalesce"), std::string::npos); + } +} + +TEST_F(VortexFileFormatTest, IoCoalesceModesPreserveDataPredicatesAndRowIds) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "io-coalesce-modes.vortex"); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::utf8())}; + const std::shared_ptr schema = arrow::schema(fields); + const std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields), + R"([[null,"zero"],[10,null],[20,"twenty"],[30,"thirty"],[40,"forty"]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + const std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/0, "id", FieldType::INT, Literal(static_cast(30))); + + for (const std::string mode : {"object-storage", "local"}) { + SCOPED_TRACE(fmt::format("mode={}", mode)); + ASSERT_OK_AND_ASSIGN( + format_, FileFormatFactory::Get( + "vortex", {{"file.format", "vortex"}, {"vortex.read.io-coalesce", mode}})); + AssertReadWithBatchSizes(path, schema, data, {1, 2, 8}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreatePredicateReader(path, schema, predicate, "auto", mode)); + ASSERT_OK_AND_ASSIGN(std::vector row_ids, CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(row_ids, (std::vector{3})); + } +} + +TEST_F(VortexFileFormatTest, MembershipPushdownModesStayLocalToEachReader) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-modes.vortex"); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32())}; + const std::shared_ptr schema = arrow::schema(fields); + const std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[null],[0],[1],[2],[3],[4],[5],[1000]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + std::vector sparse; + for (int32_t i = 0; i < 64; ++i) { + sparse.emplace_back(i * 2); + } + const std::vector> predicates = { + PredicateBuilder::In(0, "id", FieldType::INT, sparse), + PredicateBuilder::NotIn(0, "id", FieldType::INT, sparse), + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(3)))}; + const std::vector> filtered_rows = {{1, 3, 5}, {2, 4, 6, 7}, {4}}; + const std::vector all_rows = {0, 1, 2, 3, 4, 5, 6, 7}; + // Empty mode means the option is omitted, exercising the default auto policy. Configure all + // readers before consuming any, so later force/off readers cannot alter an earlier reader. + const std::vector modes = {"force", "off", "", "auto", "force", "off"}; + std::vector> readers; + std::vector> expected; + for (size_t predicate_idx = 0; predicate_idx < predicates.size(); ++predicate_idx) { + for (const std::string& mode : modes) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreatePredicateReader(path, schema, predicates[predicate_idx], mode)); + readers.push_back(std::move(reader)); + expected.push_back(predicate_idx == 2 || mode == "force" ? filtered_rows[predicate_idx] + : all_rows); + } + } + for (size_t i = 0; i < readers.size(); ++i) { + SCOPED_TRACE( + fmt::format("predicate={}, mode={}", i / modes.size(), modes[i % modes.size()])); + ASSERT_OK_AND_ASSIGN(std::vector rows, + CollectNativeRowIds(std::move(readers[i]))); + ASSERT_EQ(rows, expected[i]); + } +} + +TEST_F(VortexFileFormatTest, AutoPushesCompactLargeMembershipPredicates) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-ranges.vortex"); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32())}; + const std::shared_ptr schema = arrow::schema(fields); + const std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), + R"([[null],[0],[1],[2],[3],[4],[5],[1000]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + std::vector contiguous; + for (int32_t i = 0; i < 64; ++i) { + contiguous.emplace_back(i); + } + const std::vector> predicates = { + PredicateBuilder::In(0, "id", FieldType::INT, contiguous), + PredicateBuilder::NotIn(0, "id", FieldType::INT, contiguous)}; + const std::vector> filtered_rows = {{1, 2, 3, 4, 5, 6}, {7}}; + const std::vector all_rows = {0, 1, 2, 3, 4, 5, 6, 7}; + for (size_t predicate_idx = 0; predicate_idx < predicates.size(); ++predicate_idx) { + for (const std::string mode : {"", "auto", "force", "off"}) { + SCOPED_TRACE(fmt::format("predicate={}, mode={}", predicate_idx, mode)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreatePredicateReader(path, schema, predicates[predicate_idx], mode)); + ASSERT_OK_AND_ASSIGN(std::vector rows, + CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(rows, mode == "off" ? all_rows : filtered_rows[predicate_idx]); + } + } +} + +TEST_F(VortexFileFormatTest, FilterPushdownOnDateColumn) { + constexpr int64_t kRows = 10; + arrow::Date32Builder date_builder; + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(date_builder.Append(static_cast(18000 + i)).ok()); + } + std::shared_ptr date_array; + ASSERT_TRUE(date_builder.Finish(&date_array).ok()); + arrow::FieldVector fields = {arrow::field("d", arrow::date32(), false)}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::StructArray::Make({date_array}, fields).ValueOrDie(); + + std::string path = PathUtil::JoinPath(directory_->Str(), "filter_date.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + // DATE literals use epoch days; d > 18005 keeps physical rows 6..9. + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, "d", FieldType::DATE, Literal(FieldType::DATE, 18005)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(/*batch_size=*/3)); + 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()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + + std::vector got_days; + std::vector got_row_ids; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + std::shared_ptr arr = + arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + std::shared_ptr date_column = checked_pointer_cast( + checked_pointer_cast(arr)->GetFieldByName("d")); + for (int64_t i = 0; i < arr->length(); ++i) { + got_days.push_back(date_column->Value(i)); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + got_row_ids.push_back(row_id); + } + } + ASSERT_EQ(got_days, (std::vector{18006, 18007, 18008, 18009})); + ASSERT_EQ(got_row_ids, (std::vector{6, 7, 8, 9})); +} + +TEST_F(VortexFileFormatTest, FilterPushdownOnTimestampColumn) { + constexpr int64_t kRows = 10; + arrow::TimestampBuilder timestamp_builder(arrow::timestamp(arrow::TimeUnit::MILLI), + arrow::default_memory_pool()); + for (int64_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(timestamp_builder.Append(i * 1000).ok()); + } + std::shared_ptr timestamp_array; + ASSERT_TRUE(timestamp_builder.Finish(×tamp_array).ok()); + arrow::FieldVector fields = {arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MILLI))}; + std::shared_ptr schema = arrow::schema(fields); + std::shared_ptr data = + arrow::StructArray::Make({timestamp_array}, fields).ValueOrDie(); + + std::string path = PathUtil::JoinPath(directory_->Str(), "filter_ts.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + + // Timestamp literals use the stored column's unit (milliseconds here). + std::shared_ptr predicate = PredicateBuilder::GreaterThan( + /*field_index=*/0, "ts", FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(5000))); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader_builder, + format_->CreateReaderBuilder(/*batch_size=*/3)); + 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()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, predicate, /*selection_bitmap=*/std::nullopt)); + + std::vector got_ts; + std::vector got_row_ids; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + std::shared_ptr arr = + arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + std::shared_ptr timestamp_column = + checked_pointer_cast( + checked_pointer_cast(arr)->GetFieldByName("ts")); + for (int64_t i = 0; i < arr->length(); ++i) { + got_ts.push_back(timestamp_column->Value(i)); + ASSERT_OK_AND_ASSIGN(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + got_row_ids.push_back(row_id); + } + } + ASSERT_EQ(got_ts, (std::vector{6000, 7000, 8000, 9000})); + ASSERT_EQ(got_row_ids, (std::vector{6, 7, 8, 9})); +} + +TEST_F(VortexFileFormatTest, ReaderWithoutReadSchemaReturnsFullDataAndChecksRowIdBounds) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK_AND_ASSIGN(uint64_t total_rows, reader->GetNumberOfRows()); + ASSERT_EQ(total_rows, uint64_t{6}); + uint64_t rows = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(ReaderTestSchema()->fields()))); + const auto& data = checked_cast(*array); + const auto& ids = checked_cast(*data.field(0)); + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_EQ(physical_row, rows); + ASSERT_EQ(static_cast(ids.Value(i)), rows++); + } + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(array->length()).ok()); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(std::numeric_limits::max()).ok()); + } + ASSERT_EQ(rows, total_rows); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated_eof, ReadNextBatch(reader.get())); + ASSERT_EQ(repeated_eof, nullptr); +} + +TEST_F(VortexFileFormatTest, CombinedSelectionAndFilterPreservePhysicalRowsAfterProjection) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto projection = arrow::schema({arrow::field("name", arrow::utf8())}); + const auto predicate = + PredicateBuilder::GreaterOrEqual(0, "id", FieldType::INT, Literal(static_cast(2))); + ASSERT_OK( + ConfigureReader(reader.get(), projection, predicate, RoaringBitmap32::From({0, 2, 3, 5}))); + std::vector physical_rows; + std::vector> batches; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + physical_rows.push_back(physical_row); + } + batches.push_back(std::move(array)); + } + ASSERT_EQ(physical_rows, (std::vector{2, 3, 5})); + auto actual = arrow::Concatenate(batches, arrow_pool_.get()).ValueOrDie(); + auto expected = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(projection->fields()), + R"([["two"],["three"],["five"]])") + .ValueOrDie(); + ASSERT_TRUE(actual->Equals(expected)); +} + +TEST_F(VortexFileFormatTest, ResetReadSchemaRestartsAndClearsSelectionAndFilter) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto predicate = + PredicateBuilder::GreaterThan(0, "id", FieldType::INT, Literal(static_cast(2))); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), predicate, + RoaringBitmap32::From({1, 3, 5}))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); + ASSERT_EQ(first->length(), 2); + ASSERT_OK_AND_ASSIGN(uint64_t first_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(first_row, uint64_t{3}); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema())); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr restarted, ReadNextBatch(reader.get())); + ASSERT_EQ(restarted->length(), 2); + ASSERT_OK_AND_ASSIGN(first_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(first_row, uint64_t{0}); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr empty, ReadNextBatch(reader.get())); + ASSERT_EQ(empty, nullptr); + ASSERT_OK( + ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({0, 5}))); + ASSERT_OK_AND_ASSIGN(std::vector rows, CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(rows, (std::vector{0, 5})); +} + +TEST_F(VortexFileFormatTest, EmptyAndMissingProjectionsPreserveRowCountsAndPhysicalIds) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const std::vector> projections = { + arrow::schema({}), arrow::schema({arrow::field("missing", arrow::utf8())}), + arrow::schema({arrow::field("_row_idx", arrow::int64())}), + arrow::schema( + {arrow::field("_row_idx", arrow::int64()), arrow::field("_row_idx", arrow::int64())}), + arrow::schema( + {arrow::field("missing", arrow::struct_({arrow::field("x", arrow::int32())}))})}; + for (const auto& projection : projections) { + SCOPED_TRACE(projection->ToString()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), projection)); + uint64_t rows = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); + const auto& values = checked_cast(*array); + for (const auto& child : values.fields()) { + ASSERT_EQ(child->null_count(), array->length()); + } + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_EQ(physical_row, rows++); + } + } + ASSERT_EQ(rows, uint64_t{6}); + const auto predicate = PredicateBuilder::GreaterThan(0, "id", FieldType::INT, + Literal(static_cast(2))); + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, + RoaringBitmap32::From({0, 1, 3, 5}))); + std::vector selected; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + const auto& values = checked_cast(*array); + for (const auto& child : values.fields()) { + ASSERT_EQ(child->null_count(), array->length()); + } + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + // File readers may return candidates; the caller applies the complete predicate. + if (row > 2) { + selected.push_back(row); + } + } + } + ASSERT_EQ(selected, (std::vector{3, 5})); + } +} + +TEST_F(VortexFileFormatTest, RepeatedProjectionFieldsKeepBothValues) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const auto projection = arrow::schema( + {arrow::field("id", arrow::int32(), false), arrow::field("id", arrow::int32(), false)}); + for (const bool filtered : {false, true}) { + SCOPED_TRACE(filtered); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const std::shared_ptr predicate = + filtered ? PredicateBuilder::GreaterThan(0, "id", FieldType::INT, + Literal(static_cast(2))) + : nullptr; + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate)); + uint64_t rows = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); + const auto& values = checked_cast(*array); + ASSERT_TRUE(values.field(0)->Equals(values.field(1))); + rows += array->length(); + } + ASSERT_EQ(rows, filtered ? uint64_t{3} : uint64_t{6}); + if (filtered) { + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, + RoaringBitmap32::From({1, 3, 5}))); + ASSERT_OK_AND_ASSIGN(std::vector selected, + CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(selected, (std::vector{3, 5})); + } + } +} + +TEST_F(VortexFileFormatTest, FilterDoesNotOverwriteAFileColumnNamedRowIdx) { + const auto schema = arrow::schema({arrow::field("id", arrow::int32(), false), + arrow::field("_row_idx", arrow::int64(), false)}); + auto data = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), + R"([[0,100],[1,101],[2,102],[3,103]])") + .ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "row-idx-column.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto predicate = + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(3))); + ASSERT_OK(ConfigureReader(reader.get(), schema, predicate, RoaringBitmap32::From({1, 3}))); + const auto filter = std::dynamic_pointer_cast(predicate); + ASSERT_NE(filter, nullptr); + std::vector matched_rows; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + const auto& values = checked_cast(*array); + const auto& original_column = checked_cast(*values.field(1)); + ASSERT_OK_AND_ASSIGN(std::vector mask, filter->Test(*array, arrow_pool_.get())); + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_EQ(original_column.Value(i), 100 + static_cast(physical_row)); + if (mask[i] != 0) { + matched_rows.push_back(physical_row); + } + } + } + ASSERT_EQ(matched_rows, (std::vector{3})); +} + +TEST_F(VortexFileFormatTest, InvalidSelectionResetDoesNotChangeTheActiveReader) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From({0, 2, 4}))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); + ASSERT_EQ(first->length(), 2); + ASSERT_NOK( + ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({1, 6}))); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(physical_row, uint64_t{0}); + ASSERT_OK_AND_ASSIGN(physical_row, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(physical_row, uint64_t{2}); + ASSERT_OK_AND_ASSIGN(std::vector remaining, CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(remaining, (std::vector{4})); +} + +TEST_F(VortexFileFormatTest, InvalidArrowSchemaResetDoesNotChangeTheActiveReader) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From({0, 2, 4}))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); + ASSERT_EQ(first->length(), 2); + ::ArrowSchema invalid_schema{}; + ASSERT_NOK(reader->SetReadSchema(&invalid_schema, nullptr, RoaringBitmap32::From({1, 3}))); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(physical_row, uint64_t{0}); + ASSERT_OK_AND_ASSIGN(std::vector remaining, CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(remaining, (std::vector{4})); +} + +TEST_F(VortexFileFormatTest, ReturnedBatchOutlivesClosedReaderAndClosedReaderRejectsReset) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, ReadNextBatch(reader.get())); + reader->Close(); + reader->Close(); + ASSERT_NOK(ConfigureReader(reader.get(), ReaderTestSchema())); + ASSERT_FALSE(reader->NextBatch().ok()); + reader.reset(); + ASSERT_TRUE(batch->ValidateFull().ok()); + auto expected = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(ReaderTestSchema()->fields()), R"([[0,"zero"],[1,null]])") + .ValueOrDie(); + ASSERT_TRUE(batch->Equals(expected)); +} + +TEST_F(VortexFileFormatTest, RawBatchOutlivesScanResetAndReaderDestruction) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema())); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch retained, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(retained)); + const auto projection = arrow::schema({arrow::field("name", arrow::utf8())}); + const auto predicate = + PredicateBuilder::GreaterThan(0, "id", FieldType::INT, Literal(static_cast(2))); + ASSERT_OK( + ConfigureReader(reader.get(), projection, predicate, RoaringBitmap32::From({1, 3, 5}))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr projected, ReadNextBatch(reader.get())); + reader.reset(); + ASSERT_OK(paimon::test::ReadResultCollector::CheckBatchOffset(retained)); + auto original = arrow::ImportArray(retained.first.get(), retained.second.get()).ValueOrDie(); + ASSERT_TRUE(original->ValidateFull().ok()); + ASSERT_TRUE(projected->ValidateFull().ok()); + auto expected_original = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(ReaderTestSchema()->fields()), + R"([[0,"zero"],[1,null]])") + .ValueOrDie(); + auto expected_projected = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(projection->fields()), R"([["three"],["five"]])") + .ValueOrDie(); + ASSERT_TRUE(original->Equals(expected_original)); + ASSERT_TRUE(projected->Equals(expected_projected)); +} + +TEST_F(VortexFileFormatTest, ProjectionPreservesRequestedFieldMetadataAndNullability) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto metadata = arrow::key_value_metadata({"projection-tag"}, {"retained"}); + const auto projection = arrow::schema({arrow::field("name", arrow::utf8(), false, metadata), + arrow::field("id", arrow::int32(), true, metadata)}); + ASSERT_OK(ConfigureReader(reader.get(), projection, nullptr, RoaringBitmap32::From({0, 2}))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> file_schema, reader->GetFileSchema()); + auto imported_file_schema = arrow::ImportSchema(file_schema.get()).ValueOrDie(); + ASSERT_TRUE(imported_file_schema->Equals(*ReaderTestSchema())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE( + array->type()->Equals(arrow::struct_(projection->fields()), /*check_metadata=*/true)); + auto expected = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(projection->fields()), + R"([["zero",0],["two",2]])") + .ValueOrDie(); + ASSERT_TRUE(array->Equals(expected)); +} + +TEST_F(VortexFileFormatTest, EmptyFileAndZeroMatchFilterReturnEof) { + const auto schema = ReaderTestSchema(); + auto empty = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), "[]") + .ValueOrDie(); + const std::string empty_path = PathUtil::JoinPath(directory_->Str(), "empty.vortex"); + ASSERT_OK(WriteFile(empty_path, schema, empty, /*batch_size=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr empty_reader, CreateReader(empty_path)); + ASSERT_OK_AND_ASSIGN(uint64_t count, empty_reader->GetNumberOfRows()); + ASSERT_EQ(count, uint64_t{0}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, ReadNextBatch(empty_reader.get())); + ASSERT_EQ(batch, nullptr); + ASSERT_FALSE(empty_reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto predicate = + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(100))); + ASSERT_OK(ConfigureReader(reader.get(), schema, predicate)); + ASSERT_OK_AND_ASSIGN(batch, ReadNextBatch(reader.get())); + ASSERT_EQ(batch, nullptr); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); +} + } // namespace paimon::vortex::test diff --git a/src/paimon/format/vortex/vortex_io_callbacks.h b/src/paimon/format/vortex/vortex_io_callbacks.h index 3a095d409..afa46fb49 100644 --- a/src/paimon/format/vortex/vortex_io_callbacks.h +++ b/src/paimon/format/vortex/vortex_io_callbacks.h @@ -36,6 +36,10 @@ class OutputStream; namespace paimon::vortex { +/// I/O coalescing policy. Low-latency sources must explicitly opt into LOCAL; +/// the reader does not infer filesystem semantics from InputStream URIs. +enum class VortexIoCoalesceMode { OBJECT_STORAGE, LOCAL }; + /// Bridges Vortex's positional reads to a paimon `InputStream`. class VortexInputContext { public: diff --git a/src/paimon/format/vortex/vortex_predicate_converter.cpp b/src/paimon/format/vortex/vortex_predicate_converter.cpp new file mode 100644 index 000000000..b74bc2f88 --- /dev/null +++ b/src/paimon/format/vortex/vortex_predicate_converter.cpp @@ -0,0 +1,613 @@ +/* + * 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/vortex/vortex_predicate_converter.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/type.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/predicate/compound_predicate.h" +#include "paimon/predicate/function.h" +#include "paimon/predicate/leaf_predicate.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" + +namespace paimon::vortex { +namespace { + +VxExpressionPtr NullExpr() { + return VxExpressionPtr(nullptr, vx_expression_free); +} + +void ClearError(vx_error* error) { + if (error != nullptr) { + vx_error_free(error); + } +} + +// vx_expression_not returns an owned handle typed const; reclaim it as an owned non-const pointer. +VxExpressionPtr MakeNotExpr(const vx_expression* child) { + return VxExpressionPtr(const_cast(vx_expression_not(child)), + vx_expression_free); +} + +// Reject precision loss and overflow before constructing a native timestamp scalar. +bool GetTimestampValue(const Literal& literal, const std::shared_ptr& arrow_type, + int64_t* result, vx_time_unit* result_unit) { + if (arrow_type == nullptr || arrow_type->id() != arrow::Type::TIMESTAMP) { + return false; + } + const auto& timestamp_type = checked_cast(*arrow_type); + if (!timestamp_type.timezone().empty() && timestamp_type.timezone() != "UTC") { + return false; + } + const auto timestamp = literal.GetValue(); + int64_t value = timestamp.GetMillisecond(); + const int32_t nanos = timestamp.GetNanoOfMillisecond(); + vx_time_unit unit = VX_TIME_UNIT_MILLISECOND; + int64_t factor = 1; + int64_t fraction = 0; + switch (timestamp_type.unit()) { + case arrow::TimeUnit::SECOND: + if (value % 1000 != 0 || nanos != 0) { + return false; + } + value /= 1000; + unit = VX_TIME_UNIT_SECOND; + break; + case arrow::TimeUnit::MILLI: + if (nanos != 0) { + return false; + } + break; + case arrow::TimeUnit::MICRO: + if (nanos % 1000 != 0) { + return false; + } + factor = 1000; + fraction = nanos / 1000; + unit = VX_TIME_UNIT_MICROSECOND; + break; + case arrow::TimeUnit::NANO: + factor = 1000000; + fraction = nanos; + unit = VX_TIME_UNIT_NANOSECOND; + break; + } + if (value < std::numeric_limits::min() / factor || + value > std::numeric_limits::max() / factor) { + return false; + } + value *= factor; + if (value > std::numeric_limits::max() - fraction) { + return false; + } + *result = value + fraction; + *result_unit = unit; + return true; +} + +VxScalarPtr MakeLiteralScalar(const Literal& literal, + const std::shared_ptr& arrow_type) { + VxScalarPtr scalar(nullptr, vx_scalar_free); + if (literal.IsNull()) { + return scalar; + } + vx_error* error = nullptr; + switch (literal.GetType()) { + case FieldType::BOOLEAN: + scalar.reset(vx_scalar_new_bool(literal.GetValue(), false)); + break; + case FieldType::TINYINT: + scalar.reset(vx_scalar_new_i8(literal.GetValue(), false)); + break; + case FieldType::SMALLINT: + scalar.reset(vx_scalar_new_i16(literal.GetValue(), false)); + break; + case FieldType::INT: + scalar.reset(vx_scalar_new_i32(literal.GetValue(), false)); + break; + // Vortex stores a date column as a date[days] extension and rejects comparing it against a + // plain int scalar, so build a date-typed scalar. + case FieldType::DATE: + scalar.reset(vx_scalar_new_date(literal.GetValue(), false, &error)); + break; + case FieldType::BIGINT: + scalar.reset(vx_scalar_new_i64(literal.GetValue(), false)); + break; + case FieldType::FLOAT: + case FieldType::DOUBLE: + // Native totalOrder distinguishes NaN signs and payloads; Paimon treats all NaNs alike. + return scalar; + case FieldType::STRING: { + const auto value = literal.GetValue(); + scalar.reset(vx_scalar_new_utf8(value.data(), value.size(), false, &error)); + break; + } + case FieldType::DECIMAL: { + const auto decimal = literal.GetValue(); + const auto precision = static_cast(decimal.Precision()); + const auto scale = static_cast(decimal.Scale()); + if (decimal.IsCompact()) { + scalar.reset(vx_scalar_new_decimal_i64(decimal.ToUnscaledLong(), precision, scale, + false, &error)); + } else { + // Little-endian 16-byte two's complement, low word first. + uint8_t bytes16[16]; + const uint64_t low = decimal.LowBits(); + const uint64_t high = decimal.HighBits(); + std::memcpy(bytes16, &low, sizeof(low)); + std::memcpy(bytes16 + sizeof(low), &high, sizeof(high)); + scalar.reset( + vx_scalar_new_decimal_i128_le(bytes16, precision, scale, false, &error)); + } + break; + } + case FieldType::TIMESTAMP: { + int64_t value = 0; + vx_time_unit unit = VX_TIME_UNIT_MILLISECOND; + if (!GetTimestampValue(literal, arrow_type, &value, &unit)) { + return scalar; + } + const auto& timestamp_type = checked_cast(*arrow_type); + const char* zone = timestamp_type.timezone().empty() ? nullptr : "UTC"; + scalar.reset(vx_scalar_new_timestamp(value, unit, zone, false, &error)); + break; + } + default: + return scalar; + } + ClearError(error); + return scalar; +} + +VxExpressionPtr MakeLiteralExpr(const Literal& literal, + const std::shared_ptr& arrow_type) { + VxScalarPtr scalar = MakeLiteralScalar(literal, arrow_type); + if (scalar == nullptr) { + return NullExpr(); + } + vx_error* literal_error = nullptr; + vx_expression* expr = vx_expression_literal(scalar.get(), &literal_error); + ClearError(literal_error); + return VxExpressionPtr(expr, vx_expression_free); +} + +// Reject implicit coercions that could narrow values before the schema-evolution/residual path. +bool LiteralMatchesColumn(const Literal& literal, FieldType field_type, + const std::shared_ptr& arrow_type) { + if (literal.IsNull() || literal.GetType() != field_type || arrow_type == nullptr) { + return false; + } + switch (field_type) { + case FieldType::BOOLEAN: + return arrow_type->id() == arrow::Type::BOOL; + case FieldType::TINYINT: + return arrow_type->id() == arrow::Type::INT8; + case FieldType::SMALLINT: + return arrow_type->id() == arrow::Type::INT16; + case FieldType::INT: + return arrow_type->id() == arrow::Type::INT32; + case FieldType::BIGINT: + return arrow_type->id() == arrow::Type::INT64; + case FieldType::STRING: + return arrow_type->id() == arrow::Type::STRING; + case FieldType::DATE: + return arrow_type->id() == arrow::Type::DATE32; + case FieldType::DECIMAL: { + if (arrow_type->id() != arrow::Type::DECIMAL128) { + return false; + } + const auto& decimal_type = checked_cast(*arrow_type); + const auto decimal = literal.GetValue(); + return decimal.Precision() == decimal_type.precision() && + decimal.Scale() == decimal_type.scale(); + } + case FieldType::TIMESTAMP: { + if (arrow_type->id() != arrow::Type::TIMESTAMP) { + return false; + } + const auto& timestamp_type = checked_cast(*arrow_type); + return timestamp_type.timezone().empty() || timestamp_type.timezone() == "UTC"; + } + default: + // NaN and signed-zero equivalence needs separate validation; keep the residual path. + return false; + } +} + +struct LiteralPointerHash { + size_t operator()(const Literal* literal) const { + return literal->HashCode(); + } +}; + +struct LiteralPointerEqual { + bool operator()(const Literal* lhs, const Literal* rhs) const { + return *lhs == *rhs; + } +}; + +// All literals must be validated before their values are inspected or used as hash keys. Decimal +// scale/precision is identical to the column, so Literal's same-scale hash equality is safe here. +bool ValidateAndDeduplicate(const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type, + std::vector* literals) { + if (leaf->Literals().empty()) { + return false; + } + for (const Literal& literal : leaf->Literals()) { + if (!LiteralMatchesColumn(literal, leaf->GetFieldType(), arrow_type)) { + return false; + } + if (leaf->GetFieldType() == FieldType::TIMESTAMP) { + int64_t value = 0; + vx_time_unit unit = VX_TIME_UNIT_MILLISECOND; + if (!GetTimestampValue(literal, arrow_type, &value, &unit)) { + return false; + } + } + } + std::unordered_set unique; + unique.reserve(leaf->Literals().size()); + literals->reserve(leaf->Literals().size()); + for (const Literal& literal : leaf->Literals()) { + if (unique.insert(&literal).second) { + literals->push_back(&literal); + } + } + return true; +} + +VxExpressionPtr ConvertInList(const std::vector& literals, + const std::shared_ptr& arrow_type, + const vx_expression* column) { + std::vector scalars; + std::vector raw; + scalars.reserve(literals.size()); + raw.reserve(literals.size()); + for (const Literal* literal : literals) { + VxScalarPtr scalar = MakeLiteralScalar(*literal, arrow_type); + if (scalar == nullptr) { + return NullExpr(); + } + raw.push_back(scalar.get()); + scalars.push_back(std::move(scalar)); + } + vx_error* error = nullptr; + VxScalarPtr list( + vx_scalar_new_list(vx_scalar_dtype(raw[0]), raw.data(), raw.size(), false, &error), + vx_scalar_free); + ClearError(error); + if (list == nullptr) { + return NullExpr(); + } + error = nullptr; + VxExpressionPtr literal_expr(vx_expression_literal(list.get(), &error), vx_expression_free); + ClearError(error); + if (literal_expr == nullptr) { + return NullExpr(); + } + return VxExpressionPtr(vx_expression_list_contains(literal_expr.get(), column), + vx_expression_free); +} + +bool IsSignedInteger(FieldType field_type) { + return field_type == FieldType::TINYINT || field_type == FieldType::SMALLINT || + field_type == FieldType::INT || field_type == FieldType::BIGINT; +} + +// Only called for non-null signed integers whose type has already been validated. +int64_t SignedIntegerValue(const Literal& literal) { + switch (literal.GetType()) { + case FieldType::TINYINT: + return literal.GetValue(); + case FieldType::SMALLINT: + return literal.GetValue(); + case FieldType::INT: + return literal.GetValue(); + case FieldType::BIGINT: + return literal.GetValue(); + default: + assert(false); + return 0; + } +} + +using IntegerRange = std::pair; + +std::vector PlanIntegerRanges(std::vector literals) { + std::sort(literals.begin(), literals.end(), [](const Literal* lhs, const Literal* rhs) { + return SignedIntegerValue(*lhs) < SignedIntegerValue(*rhs); + }); + std::vector ranges; + for (const Literal* literal : literals) { + const int64_t value = SignedIntegerValue(*literal); + const int64_t previous = ranges.empty() ? 0 : SignedIntegerValue(*ranges.back().second); + // Guard before addition, including when the stored type is BIGINT at INT64_MAX. + if (!ranges.empty() && previous != std::numeric_limits::max() && + value == previous + 1) { + ranges.back().second = literal; + } else { + ranges.emplace_back(literal, literal); + } + } + return ranges; +} + +// Exclude input/literal nodes. Each range costs two comparisons and one connective; R ranges +// need R-1 additional connectives. +uint64_t IntegerRangeNodes(const std::vector& ranges) { + uint64_t nodes = ranges.size() - 1; + for (const auto& [lower, upper] : ranges) { + nodes += lower == upper ? 1 : 3; + } + return nodes; +} + +VxExpressionPtr ConvertIntegerRanges(const std::vector& ranges, + const std::shared_ptr& arrow_type, + const vx_expression* column, bool negate) { + std::vector expressions; + expressions.reserve(ranges.size()); + for (const auto& [lower, upper] : ranges) { + VxExpressionPtr lower_literal = MakeLiteralExpr(*lower, arrow_type); + if (lower_literal == nullptr) { + return NullExpr(); + } + if (lower == upper) { + VxExpressionPtr equal(vx_expression_binary(negate ? VX_OPERATOR_NOT_EQ : VX_OPERATOR_EQ, + column, lower_literal.get()), + vx_expression_free); + if (equal == nullptr) { + return NullExpr(); + } + expressions.push_back(std::move(equal)); + continue; + } + VxExpressionPtr upper_literal = MakeLiteralExpr(*upper, arrow_type); + if (upper_literal == nullptr) { + return NullExpr(); + } + VxExpressionPtr lower_bound(vx_expression_binary(negate ? VX_OPERATOR_LT : VX_OPERATOR_GTE, + column, lower_literal.get()), + vx_expression_free); + VxExpressionPtr upper_bound(vx_expression_binary(negate ? VX_OPERATOR_GT : VX_OPERATOR_LTE, + column, upper_literal.get()), + vx_expression_free); + if (lower_bound == nullptr || upper_bound == nullptr) { + return NullExpr(); + } + const vx_expression* bounds[] = {lower_bound.get(), upper_bound.get()}; + VxExpressionPtr expression( + negate ? vx_expression_or(bounds, 2) : vx_expression_and(bounds, 2), + vx_expression_free); + if (expression == nullptr) { + return NullExpr(); + } + expressions.push_back(std::move(expression)); + } + if (expressions.size() == 1) { + return std::move(expressions.front()); + } + std::vector raw; + raw.reserve(expressions.size()); + for (const VxExpressionPtr& expression : expressions) { + raw.push_back(expression.get()); + } + // IN is the union of closed ranges; NOT IN intersects their open complements. + return VxExpressionPtr(negate ? vx_expression_and(raw.data(), raw.size()) + : vx_expression_or(raw.data(), raw.size()), + vx_expression_free); +} + +VxExpressionPtr ConvertMembership(const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type, + const vx_expression* column, bool negate, + const VortexPredicateOptions& options) { + if (options.membership_pushdown == MembershipPushdownMode::OFF) { + return NullExpr(); + } + std::vector literals; + if (!ValidateAndDeduplicate(leaf, arrow_type, &literals)) { + return NullExpr(); + } + std::vector ranges; + if (IsSignedInteger(leaf->GetFieldType())) { + ranges = PlanIntegerRanges(literals); + } + // ListContains lowers N values to N equality nodes and N-1 OR nodes. NOT IN adds NOT plus + // the IS NOT NULL guard (IS NULL, NOT, AND); a complement-of-ranges needs only the guard. + const uint64_t list_nodes = 2 * static_cast(literals.size()) - 1 + (negate ? 4 : 0); + const uint64_t range_nodes = + ranges.empty() ? list_nodes : IntegerRangeNodes(ranges) + (negate ? 3 : 0); + const bool use_ranges = !ranges.empty() && range_nodes < list_nodes; + const uint64_t nodes = use_ranges ? range_nodes : list_nodes; + if (options.membership_pushdown == MembershipPushdownMode::AUTO && + nodes > options.max_membership_nodes) { + return NullExpr(); + } + // Allocate native scalars only after the plan passes the budget; ranges need only endpoints. + VxExpressionPtr expression = use_ranges + ? ConvertIntegerRanges(ranges, arrow_type, column, negate) + : ConvertInList(literals, arrow_type, column); + if (expression == nullptr || !negate) { + return expression; + } + if (!use_ranges) { + expression = MakeNotExpr(expression.get()); + } + VxExpressionPtr is_null(vx_expression_is_null(column), vx_expression_free); + if (expression == nullptr || is_null == nullptr) { + return NullExpr(); + } + VxExpressionPtr is_not_null = MakeNotExpr(is_null.get()); + if (is_not_null == nullptr) { + return NullExpr(); + } + // Neither IN nor NOT IN may retain null rows. Validation also rejects null/empty lists. + const vx_expression* children[] = {is_not_null.get(), expression.get()}; + return VxExpressionPtr(vx_expression_and(children, 2), vx_expression_free); +} + +bool ComparisonOperator(Function::Type type, vx_binary_operator* op) { + switch (type) { + case Function::Type::EQUAL: + *op = VX_OPERATOR_EQ; + return true; + case Function::Type::NOT_EQUAL: + *op = VX_OPERATOR_NOT_EQ; + return true; + case Function::Type::GREATER_THAN: + *op = VX_OPERATOR_GT; + return true; + case Function::Type::GREATER_OR_EQUAL: + *op = VX_OPERATOR_GTE; + return true; + case Function::Type::LESS_THAN: + *op = VX_OPERATOR_LT; + return true; + case Function::Type::LESS_OR_EQUAL: + *op = VX_OPERATOR_LTE; + return true; + default: + return false; + } +} + +VxExpressionPtr ConvertLeaf(const std::shared_ptr& leaf, + const std::shared_ptr& file_schema, + const vx_expression* root, const VortexPredicateOptions& options) { + const int32_t field_index = + file_schema != nullptr ? file_schema->GetFieldIndex(leaf->FieldName()) : -1; + if (field_index < 0) { + // Missing fields are null-filled by the reader, but cannot bind to a native file scan. + return NullExpr(); + } + const std::shared_ptr& arrow_type = file_schema->field(field_index)->type(); + const Function::Type type = leaf->GetFunction().GetType(); + VxExpressionPtr column(vx_expression_get_item(leaf->FieldName().c_str(), root), + vx_expression_free); + if (column == nullptr) { + return NullExpr(); + } + if (type == Function::Type::IS_NULL) { + return VxExpressionPtr(vx_expression_is_null(column.get()), vx_expression_free); + } + if (type == Function::Type::IS_NOT_NULL) { + VxExpressionPtr is_null(vx_expression_is_null(column.get()), vx_expression_free); + if (is_null == nullptr) { + return NullExpr(); + } + return MakeNotExpr(is_null.get()); + } + if (type == Function::Type::IN) { + return ConvertMembership(leaf, arrow_type, column.get(), false, options); + } + if (type == Function::Type::NOT_IN) { + return ConvertMembership(leaf, arrow_type, column.get(), true, options); + } + vx_binary_operator op = VX_OPERATOR_EQ; + if (!ComparisonOperator(type, &op)) { + return NullExpr(); + } + if (leaf->Literals().empty() || + !LiteralMatchesColumn(leaf->Literals()[0], leaf->GetFieldType(), arrow_type)) { + return NullExpr(); + } + VxExpressionPtr literal_expr = MakeLiteralExpr(leaf->Literals()[0], arrow_type); + if (literal_expr == nullptr) { + return NullExpr(); + } + return VxExpressionPtr(vx_expression_binary(op, column.get(), literal_expr.get()), + vx_expression_free); +} + +VxExpressionPtr ConvertImpl(const std::shared_ptr& predicate, + const std::shared_ptr& file_schema, + const vx_expression* root, const VortexPredicateOptions& options) { + if (predicate == nullptr) { + return NullExpr(); + } + const Function::Type type = predicate->GetFunction().GetType(); + if (type == Function::Type::AND || type == Function::Type::OR) { + const auto compound = std::dynamic_pointer_cast(predicate); + if (compound == nullptr) { + return NullExpr(); + } + std::vector children; + children.reserve(compound->Children().size()); + for (const std::shared_ptr& child : compound->Children()) { + VxExpressionPtr converted = ConvertImpl(child, file_schema, root, options); + if (converted != nullptr) { + children.push_back(std::move(converted)); + } else if (type == Function::Type::OR) { + // Dropping an OR branch could exclude matching rows. + return NullExpr(); + } + // AND: an unpushable child is dropped; the residual filter still applies it. + } + if (children.empty()) { + return NullExpr(); + } + std::vector raw; + raw.reserve(children.size()); + for (const VxExpressionPtr& child : children) { + raw.push_back(child.get()); + } + vx_expression* combined = type == Function::Type::AND + ? vx_expression_and(raw.data(), raw.size()) + : vx_expression_or(raw.data(), raw.size()); + return VxExpressionPtr(combined, vx_expression_free); + } + const auto leaf = std::dynamic_pointer_cast(predicate); + if (leaf == nullptr) { + return NullExpr(); + } + return ConvertLeaf(leaf, file_schema, root, options); +} + +} // namespace + +VxExpressionPtr VortexPredicateConverter::Convert(const std::shared_ptr& predicate, + const std::shared_ptr& file_schema, + const VortexPredicateOptions& options) { + if (predicate == nullptr) { + return NullExpr(); + } + VxExpressionPtr root(vx_expression_root(), vx_expression_free); + if (root == nullptr) { + return NullExpr(); + } + return ConvertImpl(predicate, file_schema, root.get(), options); +} + +} // namespace paimon::vortex diff --git a/src/paimon/format/vortex/vortex_predicate_converter.h b/src/paimon/format/vortex/vortex_predicate_converter.h new file mode 100644 index 000000000..39ed7e9c8 --- /dev/null +++ b/src/paimon/format/vortex/vortex_predicate_converter.h @@ -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. + */ + +#pragma once + +#include +#include + +#include "paimon/format/vortex/vortex_ffi_util.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { +class Predicate; +} // namespace paimon + +namespace paimon::vortex { + +enum class MembershipPushdownMode { AUTO, FORCE, OFF }; + +struct VortexPredicateOptions { + MembershipPushdownMode membership_pushdown = MembershipPushdownMode::AUTO; + // Per-leaf comparison and boolean node budget after deduplication/range compression. + uint32_t max_membership_nodes = 64; +}; + +/// Converts Paimon predicates into Vortex scan filters, leaving unsupported conditions to the +/// caller's residual filter. +class VortexPredicateConverter { + public: + VortexPredicateConverter() = delete; + ~VortexPredicateConverter() = delete; + + /// @param predicate The predicate to convert. + /// @param file_schema Stored column types, including timestamp units and zones. + /// @param options Membership pushdown policy and expression budget. + /// @return An owned expression, or nullptr if no condition can be pushed. + static VxExpressionPtr Convert(const std::shared_ptr& predicate, + const std::shared_ptr& file_schema, + const VortexPredicateOptions& options = {}); +}; + +} // namespace paimon::vortex diff --git a/src/paimon/format/vortex/vortex_predicate_converter_test.cpp b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp new file mode 100644 index 000000000..3759218ea --- /dev/null +++ b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp @@ -0,0 +1,1341 @@ +/* + * 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/vortex/vortex_predicate_converter.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/format/vortex/vortex_io_callbacks.h" +#include "paimon/fs/file_system.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::vortex::test { +namespace { + +std::shared_ptr TestSchema() { + return arrow::schema({ + arrow::field("i", arrow::int32()), + arrow::field("s", arrow::utf8()), + arrow::field("ts_ms", arrow::timestamp(arrow::TimeUnit::MILLI)), + arrow::field("ts_ltz", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC")), + arrow::field("l", arrow::int64()), + arrow::field("b", arrow::boolean()), + arrow::field("dt", arrow::date32()), + arrow::field("dec_compact", arrow::decimal128(10, 2)), + arrow::field("dec_wide", arrow::decimal128(19, 2)), + }); +} + +bool Pushed(const std::shared_ptr& predicate, + const std::shared_ptr& schema, + const VortexPredicateOptions& options = {}) { + return VortexPredicateConverter::Convert(predicate, schema, options) != nullptr; +} + +struct NativeInCase { + const char* name; + std::shared_ptr type; + FieldType field_type; + const char* values_json; + std::vector literals; +}; + +std::string NativeErrorMessage(const vx_error* error) { + if (error == nullptr) { + return ""; + } + const vx_string* message = vx_error_get_message(error); + return message == nullptr ? "unknown Vortex error" + : std::string(vx_string_ptr(message), vx_string_len(message)); +} + +// Check native rows without a Paimon residual filter so conversion alone cannot pass the test. +void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptr& predicate, + const std::vector& expected_rows, bool project_rowid_only = false, + const VortexPredicateOptions& options = {}, bool expect_native_filter = true, + const vx_expression* native_filter_override = nullptr) { + SCOPED_TRACE(test_case.name); + auto values_result = + arrow::ipc::internal::json::ArrayFromJSON(test_case.type, test_case.values_json); + ASSERT_TRUE(values_result.ok()) << values_result.status().ToString(); + auto ids_result = arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[0,1,2,3]"); + ASSERT_TRUE(ids_result.ok()) << ids_result.status().ToString(); + const arrow::FieldVector fields = {arrow::field("rowid", arrow::int64(), false), + arrow::field("value", test_case.type)}; + const std::shared_ptr schema = arrow::schema(fields); + auto array_result = + arrow::StructArray::Make({ids_result.ValueOrDie(), values_result.ValueOrDie()}, fields); + ASSERT_TRUE(array_result.ok()) << array_result.status().ToString(); + + std::unique_ptr directory = + paimon::test::UniqueTestDirectory::Create(); + ASSERT_NE(directory, nullptr); + const std::string path = PathUtil::JoinPath(directory->Str(), "values.vortex"); + vx_error* error = nullptr; + ScopeGuard release_error([&error]() { + if (error != nullptr) { + vx_error_free(error); + } + }); + VxSessionPtr session(vx_session_new(), vx_session_free); + ASSERT_NE(session, nullptr); + ::ArrowArray exported_array{}; + ::ArrowSchema exported_schema{}; + ScopeGuard release_export([&exported_array, &exported_schema]() { + if (exported_array.release != nullptr) { + exported_array.release(&exported_array); + } + if (exported_schema.release != nullptr) { + exported_schema.release(&exported_schema); + } + }); + ASSERT_TRUE( + arrow::ExportArray(*array_result.ValueOrDie(), &exported_array, &exported_schema).ok()); + // Import consumes both Arrow handles; the guard also covers an incomplete failed import. + VxArrayPtr array(vx_array_from_arrow(&exported_array, &exported_schema, false, &error), + vx_array_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(array, nullptr); + vx_file_write_array(session.get(), path.c_str(), array.get(), &error); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + LocalFileSystem file_system; + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system.Open(path)); + ASSERT_OK_AND_ASSIGN(int64_t length, input->Length()); + auto input_context = std::make_shared(input); + VxDataSourcePtr source( + vx_data_source_new_callback(session.get(), VortexInputContext::MakeCallbacks(input_context), + static_cast(length), &error), + vx_data_source_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(source, nullptr); + VxExpressionPtr filter = VortexPredicateConverter::Convert(predicate, schema, options); + ASSERT_EQ(filter != nullptr, expect_native_filter); + vx_scan_options scan_options{}; + scan_options.filter = native_filter_override == nullptr ? filter.get() : native_filter_override; + VxExpressionPtr root(nullptr, vx_expression_free); + VxExpressionPtr projection(nullptr, vx_expression_free); + if (project_rowid_only) { + root = VxExpressionPtr(vx_expression_root(), vx_expression_free); + ASSERT_NE(root, nullptr); + const char* names[] = {"rowid"}; + projection = + VxExpressionPtr(vx_expression_select(names, 1, root.get()), vx_expression_free); + ASSERT_NE(projection, nullptr); + scan_options.projection = projection.get(); + } + VxScanPtr scan(vx_data_source_scan(source.get(), &scan_options, nullptr, &error), vx_scan_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(scan, nullptr); + std::vector row_ids; + while (true) { + VxPartitionPtr partition(vx_scan_next_partition(scan.get(), &error), vx_partition_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + if (partition == nullptr) { + break; + } + ::ArrowArrayStream stream{}; + ScopeGuard release_stream([&stream]() { + if (stream.release != nullptr) { + stream.release(&stream); + } + }); + // The scan consumes the partition even when it fails. + const int32_t result = + vx_partition_scan_arrow(session.get(), partition.release(), &stream, &error); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_EQ(result, 0); + auto reader_result = arrow::ImportRecordBatchReader(&stream); + ASSERT_TRUE(reader_result.ok()) << reader_result.status().ToString(); + const std::shared_ptr reader = reader_result.ValueOrDie(); + while (true) { + std::shared_ptr batch; + const arrow::Status status = reader->ReadNext(&batch); + ASSERT_TRUE(status.ok()) << status.ToString(); + if (batch == nullptr) { + break; + } + ASSERT_EQ(batch->num_columns(), project_rowid_only ? 1 : 2); + const std::shared_ptr ids = batch->GetColumnByName("rowid"); + ASSERT_NE(ids, nullptr); + ASSERT_EQ(ids->type_id(), arrow::Type::INT64); + const auto& typed_ids = checked_cast(*ids); + for (int64_t row = 0; row < batch->num_rows(); ++row) { + ASSERT_FALSE(typed_ids.IsNull(row)); + row_ids.push_back(typed_ids.Value(row)); + } + } + } + std::sort(row_ids.begin(), row_ids.end()); + ASSERT_EQ(row_ids, expected_rows); +} + +// Compare an unsafe direct native expression with the converter's conservative fallback. +void AssertNativeComparisonRows(const NativeInCase& test_case, + const std::shared_ptr& predicate, + const vx_scalar* scalar, vx_binary_operator op, + const std::vector& expected_rows) { + VxExpressionPtr root(vx_expression_root(), vx_expression_free); + ASSERT_NE(root, nullptr); + VxExpressionPtr column(vx_expression_get_item("value", root.get()), vx_expression_free); + ASSERT_NE(column, nullptr); + vx_error* error = nullptr; + ScopeGuard release_error([&error]() { + if (error != nullptr) { + vx_error_free(error); + } + }); + VxExpressionPtr literal(vx_expression_literal(scalar, &error), vx_expression_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(literal, nullptr); + VxExpressionPtr filter(vx_expression_binary(op, column.get(), literal.get()), + vx_expression_free); + ASSERT_NE(filter, nullptr); + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, predicate, expected_rows, false, {}, false, filter.get())); +} + +template +NativeInCase IntegerRangeBoundaryCase(const char* name, + const std::shared_ptr& type, + FieldType field_type, const char* values_json) { + std::vector literals; + // Two disjoint four-value ranges at the extrema. Keep them unsorted and repeat endpoints. + for (int32_t offset = 0; offset < 4; ++offset) { + literals.emplace_back(static_cast(std::numeric_limits::max() - offset)); + literals.emplace_back(static_cast(std::numeric_limits::min() + offset)); + } + literals.emplace_back(std::numeric_limits::max()); + literals.emplace_back(std::numeric_limits::min()); + return {name, type, field_type, values_json, std::move(literals)}; +} + +TEST(VortexPredicateConverterTest, ComparisonsOnSupportedTypesArePushed) { + const std::shared_ptr schema = TestSchema(); + ASSERT_TRUE( + Pushed(PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(static_cast(42))), + schema)); + ASSERT_TRUE( + Pushed(PredicateBuilder::NotEqual(0, "i", FieldType::INT, Literal(static_cast(1))), + schema)); + ASSERT_TRUE(Pushed( + PredicateBuilder::GreaterThan(0, "l", FieldType::BIGINT, Literal(static_cast(7))), + schema)); + ASSERT_TRUE(Pushed( + PredicateBuilder::GreaterOrEqual(0, "i", FieldType::INT, Literal(static_cast(1))), + schema)); + ASSERT_TRUE( + Pushed(PredicateBuilder::LessThan(0, "i", FieldType::INT, Literal(static_cast(1))), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::Equal(0, "b", FieldType::BOOLEAN, Literal(true)), schema)); + ASSERT_TRUE(Pushed( + PredicateBuilder::Equal(0, "s", FieldType::STRING, Literal(FieldType::STRING, "abc", 3)), + schema)); + ASSERT_TRUE( + Pushed(PredicateBuilder::Equal(0, "dt", FieldType::DATE, Literal(FieldType::DATE, 19000)), + schema)); + // Compact decimal (precision <= 18) and wide decimal (precision > 18) both push. + ASSERT_TRUE(Pushed(PredicateBuilder::Equal(0, "dec_compact", FieldType::DECIMAL, + Literal(Decimal::FromUnscaledLong(12345, 10, 2))), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::Equal(0, "dec_wide", FieldType::DECIMAL, + Literal(Decimal::FromUnscaledLong(12345, 19, 2))), + schema)); +} + +TEST(VortexPredicateConverterTest, IsNullAndIsNotNullArePushed) { + const std::shared_ptr schema = TestSchema(); + ASSERT_TRUE(Pushed(PredicateBuilder::IsNull(0, "i", FieldType::INT), schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::IsNotNull(0, "i", FieldType::INT), schema)); +} + +TEST(VortexPredicateConverterTest, UnknownOrAmbiguousFileColumnsFallBack) { + const std::shared_ptr equal = + PredicateBuilder::Equal(0, "value", FieldType::INT, Literal(int32_t{1})); + const std::shared_ptr is_null = PredicateBuilder::IsNull(0, "value", FieldType::INT); + const std::shared_ptr is_not_null = + PredicateBuilder::IsNotNull(0, "value", FieldType::INT); + for (const std::shared_ptr& schema : + {std::shared_ptr(), arrow::schema({arrow::field("other", arrow::int32())}), + arrow::schema( + {arrow::field("value", arrow::int32()), arrow::field("value", arrow::int32())})}) { + ASSERT_FALSE(Pushed(equal, schema)); + ASSERT_FALSE(Pushed(is_null, schema)); + ASSERT_FALSE(Pushed(is_not_null, schema)); + } + const std::shared_ptr present = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int32_t{1})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({present, is_null})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({present, is_null})); + ASSERT_TRUE(Pushed(conjunction, TestSchema())); + ASSERT_FALSE(Pushed(disjunction, TestSchema())); +} + +TEST(VortexPredicateConverterTest, ComparisonsRequireMatchingFileAndLiteralTypes) { + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + ASSERT_FALSE(Pushed(build(0, "i", FieldType::INT, Literal(int64_t{1})), TestSchema())); + ASSERT_FALSE(Pushed(build(0, "i", FieldType::BIGINT, Literal(int64_t{1})), TestSchema())); + ASSERT_FALSE(Pushed(build(0, "dt", FieldType::INT, Literal(int32_t{1})), TestSchema())); + ASSERT_FALSE(Pushed(build(0, "dec_compact", FieldType::DECIMAL, + Literal(Decimal::FromUnscaledLong(12345, 19, 2))), + TestSchema())); + ASSERT_FALSE(Pushed(build(0, "dec_compact", FieldType::DECIMAL, + Literal(Decimal::FromUnscaledLong(12345, 10, 3))), + TestSchema())); + } +} + +TEST(VortexPredicateConverterTest, UnboundPredicatesPreserveNativeCandidateRows) { + const NativeInCase test_case = {"unbound", arrow::int32(), FieldType::INT, "[1,null,2,3]", {}}; + const std::vector> predicates = { + PredicateBuilder::Equal(2, "missing", FieldType::INT, Literal(int32_t{1})), + PredicateBuilder::IsNull(2, "missing", FieldType::INT), + PredicateBuilder::IsNotNull(2, "missing", FieldType::INT), + PredicateBuilder::Equal(1, "value", FieldType::BIGINT, Literal(int64_t{1})), + }; + for (const std::shared_ptr& predicate : predicates) { + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, false, {}, false)); + } +} + +TEST(VortexPredicateConverterTest, RandomizedIntegerMembershipMatchesIndependentEvaluation) { + struct IntegerType { + const char* name; + std::shared_ptr arrow_type; + FieldType field_type; + int64_t minimum; + int64_t maximum; + }; + const std::vector types = { + {"i8", arrow::int8(), FieldType::TINYINT, std::numeric_limits::min(), + std::numeric_limits::max()}, + {"i16", arrow::int16(), FieldType::SMALLINT, std::numeric_limits::min(), + std::numeric_limits::max()}, + {"i32", arrow::int32(), FieldType::INT, std::numeric_limits::min(), + std::numeric_limits::max()}, + {"i64", arrow::int64(), FieldType::BIGINT, std::numeric_limits::min(), + std::numeric_limits::max()}, + }; + std::mt19937_64 generator(0x5EED); + for (const IntegerType& type : types) { + const std::string json = fmt::format("[{},null,0,{}]", type.minimum, type.maximum); + const std::shared_ptr schema = + arrow::schema({arrow::field("value", type.arrow_type)}); + for (int32_t iteration = 0; iteration < 32; ++iteration) { + SCOPED_TRACE(fmt::format("{} iteration {}", type.name, iteration)); + std::vector values; + std::vector literals; + for (uint64_t remaining = 1 + generator() % 80; remaining > 0; --remaining) { + const uint64_t choice = generator() % 5; + const int64_t value = choice == 0 ? type.minimum + : choice == 1 ? type.maximum + : static_cast(generator() % 65) - 32; + values.push_back(value); + switch (type.field_type) { + case FieldType::TINYINT: + literals.emplace_back(static_cast(value)); + break; + case FieldType::SMALLINT: + literals.emplace_back(static_cast(value)); + break; + case FieldType::INT: + literals.emplace_back(static_cast(value)); + break; + default: + literals.emplace_back(value); + break; + } + } + const NativeInCase test_case = {type.name, type.arrow_type, type.field_type, + json.c_str(), literals}; + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "value", type.field_type, literals) + : PredicateBuilder::In(1, "value", type.field_type, literals); + std::vector expected; + const std::vector stored = {type.minimum, 0, type.maximum}; + const std::vector row_ids = {0, 2, 3}; + for (size_t row = 0; row < stored.size(); ++row) { + const bool found = + std::find(values.begin(), values.end(), stored[row]) != values.end(); + if (found != negate) { + expected.push_back(row_ids[row]); + } + } + for (MembershipPushdownMode mode : + {MembershipPushdownMode::AUTO, MembershipPushdownMode::FORCE}) { + const VortexPredicateOptions options{mode, 64}; + const bool pushed = Pushed(predicate, schema, options); + const std::vector candidates = + pushed ? expected : std::vector{0, 1, 2, 3}; + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, predicate, candidates, false, options, pushed)); + } + } + } + } +} + +TEST(VortexPredicateConverterTest, TimestampUsesStoredColumnUnit) { + const std::shared_ptr schema = TestSchema(); + ASSERT_TRUE(Pushed(PredicateBuilder::Equal(0, "ts_ms", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(1234))), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::Equal(0, "ts_ltz", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(1234))), + schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::Equal(0, "absent", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(1234))), + schema)); +} + +TEST(VortexPredicateConverterTest, TimestampComparisonsRequireLosslessRepresentableValues) { + const std::shared_ptr schema = arrow::schema({ + arrow::field("seconds", arrow::timestamp(arrow::TimeUnit::SECOND)), + arrow::field("millis", arrow::timestamp(arrow::TimeUnit::MILLI)), + arrow::field("micros", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC")), + arrow::field("nanos", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("local", arrow::timestamp(arrow::TimeUnit::MICRO, "Asia/Shanghai")), + }); + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + for (const char* name : {"seconds", "millis", "micros", "nanos"}) { + ASSERT_TRUE(Pushed( + build(0, name, FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(-1000))), + schema)); + } + for (const char* name : {"seconds", "millis", "micros"}) { + ASSERT_FALSE(Pushed( + build(0, name, FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(0, 1))), + schema)); + } + for (const char* name : {"micros", "nanos"}) { + for (int64_t millis : + {std::numeric_limits::min(), std::numeric_limits::max()}) { + ASSERT_FALSE(Pushed(build(0, name, FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(millis))), + schema)); + } + } + ASSERT_FALSE(Pushed( + build(0, "local", FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(1000))), + schema)); + } +} + +TEST(VortexPredicateConverterTest, LosslessTimestampComparisonsFilterNativeFiles) { + const std::vector cases = { + {"seconds", + arrow::timestamp(arrow::TimeUnit::SECOND), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(2000))}}, + {"millis", + arrow::timestamp(arrow::TimeUnit::MILLI), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(2))}}, + {"micros", + arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(0, 2000))}}, + {"nanos", + arrow::timestamp(arrow::TimeUnit::NANO), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(0, 2))}}, + }; + for (const NativeInCase& test_case : cases) { + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::Equal(1, "value", FieldType::TIMESTAMP, test_case.literals[0]), {3})); + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, + PredicateBuilder::GreaterOrEqual(1, "value", FieldType::TIMESTAMP, + test_case.literals[0]), + {2, 3})); + } +} + +TEST(VortexPredicateConverterTest, TruncatedTimestampComparisonFallsBackWithoutLosingRows) { + const NativeInCase test_case = {"micros_fraction", + arrow::timestamp(arrow::TimeUnit::MICRO), + FieldType::TIMESTAMP, + "[0,null,1,-1]", + {}}; + const Literal one_ns(Timestamp::FromEpochMillis(0, 1)); + const std::shared_ptr less_than = + PredicateBuilder::LessThan(1, "value", FieldType::TIMESTAMP, one_ns); + ASSERT_OK_AND_ASSIGN(int32_t comparison, + Literal(Timestamp::FromEpochMillis(0)).CompareTo(one_ns)); + ASSERT_LT(comparison, 0); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, less_than, {0, 1, 2, 3}, false, {}, false)); + + vx_error* error = nullptr; + ScopeGuard release_error([&error]() { + if (error != nullptr) { + vx_error_free(error); + } + }); + VxScalarPtr truncated( + vx_scalar_new_timestamp(0, VX_TIME_UNIT_MICROSECOND, nullptr, false, &error), + vx_scalar_free); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(truncated, nullptr); + // Truncating 1ns to 0us loses row 0, which the exact comparison above must retain. + ASSERT_NO_FATAL_FAILURE( + AssertNativeComparisonRows(test_case, less_than, truncated.get(), VX_OPERATOR_LT, {3})); +} + +TEST(VortexPredicateConverterTest, FloatingPointComparisonsFallBackForAllLiterals) { + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (bool single_precision : {false, true}) { + const FieldType field_type = single_precision ? FieldType::FLOAT : FieldType::DOUBLE; + const std::shared_ptr schema = arrow::schema( + {arrow::field("value", single_precision ? arrow::float32() : arrow::float64())}); + for (double value : {-0.0, 0.0, 1.5, std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) { + const Literal literal = + single_precision ? Literal(static_cast(value)) : Literal(value); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + ASSERT_FALSE(Pushed(build(0, "value", field_type, literal), schema)); + } + } + ASSERT_TRUE(Pushed(PredicateBuilder::IsNull(0, "value", field_type), schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::IsNotNull(0, "value", field_type), schema)); + } +} + +TEST(VortexPredicateConverterTest, NativeFloatingPointOrderRequiresConservativeFallback) { + const NativeInCase test_case = { + "floating_order", arrow::float64(), FieldType::DOUBLE, "[NaN,null,-0.0,0.0]", {}}; + const double nan = std::numeric_limits::quiet_NaN(); + ASSERT_OK_AND_ASSIGN(int32_t nan_comparison, Literal(nan).CompareTo(Literal(-nan))); + ASSERT_EQ(nan_comparison, 0); + ASSERT_OK_AND_ASSIGN(int32_t zero_comparison, Literal(-0.0).CompareTo(Literal(0.0))); + ASSERT_LT(zero_comparison, 0); + const std::shared_ptr equal_nan = + PredicateBuilder::Equal(1, "value", FieldType::DOUBLE, Literal(-nan)); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, equal_nan, {0, 1, 2, 3}, false, {}, false)); + VxScalarPtr negative_nan(vx_scalar_new_f64(-nan, false), vx_scalar_free); + ASSERT_NE(negative_nan, nullptr); + // Native totalOrder distinguishes NaN signs, while Paimon must retain row 0. + ASSERT_NO_FATAL_FAILURE( + AssertNativeComparisonRows(test_case, equal_nan, negative_nan.get(), VX_OPERATOR_EQ, {})); + + const std::shared_ptr equal_negative_zero = + PredicateBuilder::Equal(1, "value", FieldType::DOUBLE, Literal(-0.0)); + VxScalarPtr negative_zero(vx_scalar_new_f64(-0.0, false), vx_scalar_free); + ASSERT_NE(negative_zero, nullptr); + // Signed-zero ordering agrees in this case; it does not establish compatible NaN semantics. + ASSERT_NO_FATAL_FAILURE(AssertNativeComparisonRows(test_case, equal_negative_zero, + negative_zero.get(), VX_OPERATOR_EQ, {2})); +} + +TEST(VortexPredicateConverterTest, UnsafeComparisonsKeepSafeAndBranchesAndRejectNativeOr) { + const std::vector cases = { + {"timestamp_fraction", + arrow::timestamp(arrow::TimeUnit::MICRO), + FieldType::TIMESTAMP, + "[0,null,1,-1]", + {Literal(Timestamp::FromEpochMillis(0, 1))}}, + {"floating_nan", + arrow::float64(), + FieldType::DOUBLE, + "[NaN,null,-0.0,0.0]", + {Literal(-std::numeric_limits::quiet_NaN())}}, + }; + for (const NativeInCase& test_case : cases) { + const std::shared_ptr unsafe = + PredicateBuilder::LessThan(1, "value", test_case.field_type, test_case.literals[0]); + const std::shared_ptr safe = + PredicateBuilder::LessThan(0, "rowid", FieldType::BIGINT, Literal(int64_t{3})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({unsafe, safe})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({unsafe, safe})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, and_predicate, {0, 1, 2}, true)); + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, or_predicate, {0, 1, 2, 3}, true, {}, false)); + } +} + +TEST(VortexPredicateConverterTest, UnsupportedFunctionsAndLiteralsAreNotPushed) { + const std::shared_ptr schema = TestSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr starts_with, + PredicateBuilder::StartsWith(0, "s", FieldType::STRING, + Literal(FieldType::STRING, "a", 1))); + ASSERT_FALSE(Pushed(starts_with, schema)); + ASSERT_FALSE(Pushed( + PredicateBuilder::Equal(0, "bin", FieldType::BINARY, Literal(FieldType::BINARY, "x", 1)), + schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(FieldType::INT)), schema)); +} + +TEST(VortexPredicateConverterTest, InSmallAndLargeListsArePushed) { + const std::shared_ptr schema = TestSchema(); + for (const int32_t size : {1, 2, 20, 21, 1024}) { + std::vector literals; + for (int32_t value = 0; value < size; ++value) { + literals.emplace_back(value); + } + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), schema)); + } + ASSERT_TRUE(Pushed( + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int32_t{1})}), + schema)); +} + +TEST(VortexPredicateConverterTest, InSupportedTypesArePushed) { + const std::shared_ptr schema = arrow::schema({ + arrow::field("b", arrow::boolean()), + arrow::field("tiny", arrow::int8()), + arrow::field("small", arrow::int16()), + arrow::field("l", arrow::int64()), + arrow::field("s", arrow::utf8()), + arrow::field("dt", arrow::date32()), + arrow::field("dec", arrow::decimal128(10, 2)), + arrow::field("wide_dec", arrow::decimal128(19, 2)), + }); + ASSERT_TRUE(Pushed( + PredicateBuilder::In(0, "b", FieldType::BOOLEAN, {Literal(false), Literal(true)}), schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "tiny", FieldType::TINYINT, + {Literal(int8_t{-1}), Literal(int8_t{2})}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "small", FieldType::SMALLINT, + {Literal(int16_t{-1}), Literal(int16_t{2})}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "l", FieldType::BIGINT, + {Literal(int64_t{-1}), Literal(int64_t{2})}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "s", FieldType::STRING, + {Literal(FieldType::STRING, "", 0), + Literal(FieldType::STRING, "abc", 3)}), + schema)); + ASSERT_TRUE(Pushed( + PredicateBuilder::In(0, "dt", FieldType::DATE, + {Literal(FieldType::DATE, -1), Literal(FieldType::DATE, 19000)}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 10, 2)), + Literal(Decimal::FromUnscaledLong(-456, 10, 2))}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "wide_dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 19, 2))}), + schema)); +} + +TEST(VortexPredicateConverterTest, InAndNotInSupportedTypesFilterNativeFilesAndExcludeNullRows) { + const Decimal::int128_t wide_value = static_cast(999999999999999999LL) * 10; + const std::vector cases = { + {"boolean", + arrow::boolean(), + FieldType::BOOLEAN, + "[true,null,false,true]", + {Literal(true)}}, + {"tinyint", + arrow::int8(), + FieldType::TINYINT, + "[-1,null,7,2]", + {Literal(int8_t{-1}), Literal(int8_t{2})}}, + {"smallint", + arrow::int16(), + FieldType::SMALLINT, + "[-1,null,7,2]", + {Literal(int16_t{-1}), Literal(int16_t{2})}}, + {"int", + arrow::int32(), + FieldType::INT, + "[-1,null,7,2]", + {Literal(int32_t{-1}), Literal(int32_t{2})}}, + {"bigint", + arrow::int64(), + FieldType::BIGINT, + "[-1,null,7,2]", + {Literal(int64_t{-1}), Literal(int64_t{2})}}, + {"string", + arrow::utf8(), + FieldType::STRING, + R"(["first",null,"miss","last"])", + {Literal(FieldType::STRING, "first", 5), Literal(FieldType::STRING, "last", 4)}}, + {"date", + arrow::date32(), + FieldType::DATE, + "[-1,null,7,19000]", + {Literal(FieldType::DATE, -1), Literal(FieldType::DATE, 19000)}}, + {"decimal", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"7.00","4.56"])", + {Literal(Decimal::FromUnscaledLong(-123, 10, 2)), + Literal(Decimal::FromUnscaledLong(456, 10, 2))}}, + {"wide_decimal", + arrow::decimal128(19, 2), + FieldType::DECIMAL, + R"(["-99999999999999999.90",null,"7.00","99999999999999999.90"])", + {Literal(Decimal(19, 2, -wide_value)), Literal(Decimal(19, 2, wide_value))}}, + {"timestamp_seconds", + arrow::timestamp(arrow::TimeUnit::SECOND), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(-1000)), Literal(Timestamp::FromEpochMillis(2000))}}, + {"timestamp_millis", + arrow::timestamp(arrow::TimeUnit::MILLI), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(-1)), Literal(Timestamp::FromEpochMillis(2))}}, + {"timestamp_micros", + arrow::timestamp(arrow::TimeUnit::MICRO), + FieldType::TIMESTAMP, + "[-1,null,7000,2001]", + {Literal(Timestamp::FromEpochMillis(-1, 999000)), + Literal(Timestamp::FromEpochMillis(2, 1000))}}, + {"timestamp_nanos", + arrow::timestamp(arrow::TimeUnit::NANO), + FieldType::TIMESTAMP, + "[-1,null,7000000,2000001]", + {Literal(Timestamp::FromEpochMillis(-1, 999999)), + Literal(Timestamp::FromEpochMillis(2, 1))}}, + {"timestamp_utc", + arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + FieldType::TIMESTAMP, + "[-1,null,7000,2001]", + {Literal(Timestamp::FromEpochMillis(-1, 999000)), + Literal(Timestamp::FromEpochMillis(2, 1000))}}, + }; + for (const NativeInCase& test_case : cases) { + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", test_case.field_type, test_case.literals), + {0, 3})); + // The null row (rowid=1) must not reappear when the membership test is negated. + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::NotIn(1, "value", test_case.field_type, test_case.literals), {2})); + } +} + +TEST(VortexPredicateConverterTest, NotInSmallAndLargeListsFilterNativeFiles) { + const NativeInCase test_case = { + "not_in_list_sizes", arrow::int32(), FieldType::INT, "[0,null,2048,1]", {}}; + for (const int32_t size : {1, 2, 20, 21, 100, 1024}) { + SCOPED_TRACE(size); + std::vector literals; + for (int32_t value = 0; value < size; ++value) { + literals.emplace_back(value); + } + const std::vector expected_rows = + size == 1 ? std::vector{2, 3} : std::vector{2}; + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::INT, literals), + expected_rows)); + } +} + +TEST(VortexPredicateConverterTest, IntegerRangesAtTypeLimitsFilterNativeFiles) { + const std::vector cases = { + IntegerRangeBoundaryCase("tinyint_range_limits", arrow::int8(), FieldType::TINYINT, + "[-128,null,0,127]"), + IntegerRangeBoundaryCase("smallint_range_limits", arrow::int16(), + FieldType::SMALLINT, "[-32768,null,0,32767]"), + IntegerRangeBoundaryCase("int_range_limits", arrow::int32(), FieldType::INT, + "[-2147483648,null,0,2147483647]"), + IntegerRangeBoundaryCase("bigint_range_limits", arrow::int64(), FieldType::BIGINT, + "[-9223372036854775808,null,0,9223372036854775807]"), + }; + for (const NativeInCase& test_case : cases) { + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", test_case.field_type, test_case.literals), + {0, 3})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::NotIn(1, "value", test_case.field_type, test_case.literals), {2})); + } +} + +TEST(VortexPredicateConverterTest, IntegerRangesPreserveGapsAndSingletonsInNativeFiles) { + NativeInCase test_case = { + "integer_ranges_with_gap", arrow::int32(), FieldType::INT, "[-7,null,0,6]", {}}; + for (int32_t value = -8; value <= -1; ++value) { + test_case.literals.emplace_back(value); + } + for (int32_t value = 3; value <= 8; ++value) { + test_case.literals.emplace_back(value); + } + test_case.literals.emplace_back(int32_t{20}); + test_case.literals.emplace_back(int32_t{-8}); + std::reverse(test_case.literals.begin(), test_case.literals.end()); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", FieldType::INT, test_case.literals), {0, 3})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::INT, test_case.literals), {2})); + test_case.name = "integer_ranges_singleton"; + test_case.values_json = "[-7,null,20,0]"; + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", FieldType::INT, test_case.literals), {0, 2})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::INT, test_case.literals), {3})); +} + +TEST(VortexPredicateConverterTest, FullIntegerDomainAndDuplicateSingletonFilterNativeFiles) { + NativeInCase test_case = { + "tinyint_full_domain", arrow::int8(), FieldType::TINYINT, "[-128,null,0,127]", {}}; + // The loop uses int32_t so incrementing the maximum TINYINT cannot wrap. + for (int32_t value = std::numeric_limits::min(); + value <= std::numeric_limits::max(); ++value) { + test_case.literals.emplace_back(static_cast(value)); + } + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", FieldType::TINYINT, test_case.literals), + {0, 2, 3})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::TINYINT, test_case.literals), + {})); + test_case.name = "duplicate_singleton"; + test_case.literals = std::vector(8, Literal(int8_t{-128})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::In(1, "value", FieldType::TINYINT, test_case.literals), {0})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::TINYINT, test_case.literals), + {2, 3})); +} + +TEST(VortexPredicateConverterTest, IntegerRangesComposeWithNativePredicatesAndProjection) { + NativeInCase test_case = { + "integer_range_compound", arrow::int32(), FieldType::INT, "[-4,null,8,4]", {}}; + for (int32_t value = -4; value <= 4; ++value) { + test_case.literals.emplace_back(value); + } + const std::shared_ptr in = + PredicateBuilder::In(1, "value", FieldType::INT, test_case.literals); + const std::shared_ptr not_in = + PredicateBuilder::NotIn(1, "value", FieldType::INT, test_case.literals); + const std::shared_ptr first_row = + PredicateBuilder::Equal(0, "rowid", FieldType::BIGINT, Literal(int64_t{0})); + const std::shared_ptr first_three = + PredicateBuilder::LessThan(0, "rowid", FieldType::BIGINT, Literal(int64_t{3})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr in_and, + PredicateBuilder::And({in, first_three})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, in_and, {0}, true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr not_in_and, + PredicateBuilder::And({not_in, first_three})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in_and, {2}, true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr in_or, PredicateBuilder::Or({in, first_row})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, in_or, {0, 3}, true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr not_in_or, + PredicateBuilder::Or({not_in, first_row})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in_or, {0, 2}, true)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr not_in_or_null, + PredicateBuilder::Or({not_in, PredicateBuilder::IsNull(1, "value", FieldType::INT)})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in_or_null, {1, 2}, true)); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, in, {0, 3}, true)); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in, {2}, true)); +} + +TEST(VortexPredicateConverterTest, CompressibleIntegerListsStillValidateEveryLiteral) { + const std::shared_ptr schema = TestSchema(); + std::vector safe; + for (int32_t value = 0; value < 16; ++value) { + safe.emplace_back(value); + } + for (const Literal& unsafe : {Literal(FieldType::INT), Literal(int64_t{16})}) { + for (size_t position : {size_t{0}, size_t{8}, safe.size()}) { + std::vector literals = safe; + literals.insert(literals.begin() + position, unsafe); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, literals), schema)); + } + } + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, safe), + arrow::schema({arrow::field("i", arrow::int64())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, safe), + arrow::schema({arrow::field("i", arrow::int64())}))); +} + +TEST(VortexPredicateConverterTest, NotInDuplicatesMissesAndAllExcludedFilterNativeFiles) { + const NativeInCase test_case = { + "not_in_edges", arrow::int32(), FieldType::INT, "[-1,null,7,2]", {}}; + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::NotIn(1, "value", FieldType::INT, + {Literal(int32_t{-1}), Literal(int32_t{-1}), Literal(int32_t{2})}), + {2})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, PredicateBuilder::NotIn(1, "value", FieldType::INT, {Literal(int32_t{99})}), + {0, 2, 3})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::NotIn(1, "value", FieldType::INT, + {Literal(int32_t{-1}), Literal(int32_t{7}), Literal(int32_t{2})}), + {})); + const NativeInCase all_null = { + "not_in_all_null", arrow::int32(), FieldType::INT, "[null,null,null,null]", {}}; + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + all_null, PredicateBuilder::NotIn(1, "value", FieldType::INT, {Literal(int32_t{1})}), {})); +} + +TEST(VortexPredicateConverterTest, NotInCompoundPredicatesAndProjectionFilterNativeFiles) { + const NativeInCase test_case = { + "not_in_compound", arrow::int32(), FieldType::INT, "[-1,null,7,2]", {}}; + const std::shared_ptr not_in = PredicateBuilder::NotIn( + 1, "value", FieldType::INT, {Literal(int32_t{-1}), Literal(int32_t{2})}); + const std::shared_ptr first_row = + PredicateBuilder::Equal(0, "rowid", FieldType::BIGINT, Literal(int64_t{0})); + const std::shared_ptr first_three = + PredicateBuilder::LessThan(0, "rowid", FieldType::BIGINT, Literal(int64_t{3})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({not_in, first_three})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, and_predicate, {2})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({not_in, first_row})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, or_predicate, {0, 2})); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr or_null, + PredicateBuilder::Or({not_in, PredicateBuilder::IsNull(1, "value", FieldType::INT)})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, or_null, {1, 2})); + // The native filter still works when the filtering column is omitted from the projection. + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in, {2}, true)); + + const std::shared_ptr unsupported = PredicateBuilder::NotIn( + 1, "value", FieldType::INT, {Literal(int32_t{7}), Literal(FieldType::INT)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_partial, + PredicateBuilder::And({not_in, unsupported})); + // The residual must evaluate the NULL-containing list; native AND keeps only the safe child. + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, and_partial, {2})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_partial, + PredicateBuilder::Or({not_in, unsupported})); + ASSERT_FALSE(Pushed(or_partial, arrow::schema({arrow::field("rowid", arrow::int64()), + arrow::field("value", arrow::int32())}))); +} + +TEST(VortexPredicateConverterTest, NotInUnsafeLiteralsAreNotPushed) { + const std::shared_ptr schema = TestSchema(); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {}), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(FieldType::INT)}), schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, + {Literal(int32_t{1}), Literal(FieldType::INT)}), + schema)); + ASSERT_FALSE(Pushed( + PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int64_t{2})}), + schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "absent", FieldType::INT, {Literal(int32_t{1})}), + schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(int32_t{1})}), nullptr)); + ASSERT_FALSE(Pushed( + PredicateBuilder::NotIn(0, "bin", FieldType::BINARY, {Literal(FieldType::BINARY, "x", 1)}), + arrow::schema({arrow::field("bin", arrow::binary())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "f", FieldType::FLOAT, {Literal(1.5F)}), + arrow::schema({arrow::field("f", arrow::float32())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "d", FieldType::DOUBLE, {Literal(1.5)}), + arrow::schema({arrow::field("d", arrow::float64())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 10, 3))}), + arrow::schema({arrow::field("dec", arrow::decimal128(10, 2))}))); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "ts", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(1, 1))}), + arrow::schema({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MILLI))}))); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn( + 0, "ts", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(std::numeric_limits::max()))}), + arrow::schema({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::NANO))}))); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "ts", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(1000))}), + arrow::schema({arrow::field( + "ts", arrow::timestamp(arrow::TimeUnit::MICRO, "Asia/Shanghai"))}))); +} + +TEST(VortexPredicateConverterTest, InUnsupportedOrInconsistentLiteralsAreNotPushed) { + const std::shared_ptr schema = TestSchema(); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, {}), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(0, "i", FieldType::INT, {Literal(FieldType::INT)}), schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, + {Literal(int32_t{1}), Literal(FieldType::INT)}), + schema)); + ASSERT_FALSE(Pushed( + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int64_t{2})}), + schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(0, "absent", FieldType::INT, {Literal(int32_t{1})}), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1})}), nullptr)); + ASSERT_FALSE(Pushed( + PredicateBuilder::In(0, "bin", FieldType::BINARY, {Literal(FieldType::BINARY, "x", 1)}), + arrow::schema({arrow::field("bin", arrow::binary())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "f", FieldType::FLOAT, {Literal(1.5F)}), + arrow::schema({arrow::field("f", arrow::float32())}))); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "d", FieldType::DOUBLE, {Literal(1.5)}), + arrow::schema({arrow::field("d", arrow::float64())}))); + const std::shared_ptr decimal_schema = + arrow::schema({arrow::field("dec", arrow::decimal128(10, 2))}); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 10, 2)), + Literal(Decimal::FromUnscaledLong(123, 10, 3))}), + decimal_schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 11, 2))}), + decimal_schema)); +} + +TEST(VortexPredicateConverterTest, InTimestampsRequireMatchingUnitsAndLosslessValues) { + const std::shared_ptr schema = arrow::schema({ + arrow::field("seconds", arrow::timestamp(arrow::TimeUnit::SECOND)), + arrow::field("millis", arrow::timestamp(arrow::TimeUnit::MILLI)), + arrow::field("micros", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC")), + arrow::field("nanos", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("local", arrow::timestamp(arrow::TimeUnit::MICRO, "Asia/Shanghai")), + }); + for (const char* name : {"seconds", "millis", "micros", "nanos"}) { + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, name, FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(-1000)), + Literal(Timestamp::FromEpochMillis(2000))}), + schema)); + } + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "micros", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(-1, 1000))}), + schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "nanos", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(-1, 1))}), + schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "seconds", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(1))}), + schema)); + for (const char* name : {"millis", "micros"}) { + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, name, FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(1, 1))}), + schema)); + } + for (const char* name : {"micros", "nanos"}) { + ASSERT_FALSE( + Pushed(PredicateBuilder::In( + 0, name, FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(std::numeric_limits::max()))}), + schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In( + 0, name, FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(std::numeric_limits::min()))}), + schema)); + } + // The millisecond multiplication can fit while adding the fractional part overflows. + const int64_t maximum = std::numeric_limits::max(); + ASSERT_FALSE( + Pushed(PredicateBuilder::In( + 0, "micros", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis( + maximum / 1000, static_cast((maximum % 1000 + 1) * 1000)))}), + schema)); + ASSERT_FALSE(Pushed( + PredicateBuilder::In(0, "nanos", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis( + maximum / 1000000, static_cast(maximum % 1000000 + 1)))}), + schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "local", FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(1000))}), + schema)); +} + +TEST(VortexPredicateConverterTest, AndDropsUnsupportedChildrenButOrRejects) { + const std::shared_ptr schema = TestSchema(); + const std::shared_ptr supported = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(static_cast(1))); + const std::shared_ptr other = + PredicateBuilder::In(0, "i", FieldType::INT, + {Literal(static_cast(2)), Literal(static_cast(3))}); + const std::shared_ptr unsupported = PredicateBuilder::In( + 0, "i", FieldType::INT, {Literal(static_cast(1)), Literal(FieldType::INT)}); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_partial, + PredicateBuilder::And({supported, unsupported})); + ASSERT_TRUE(Pushed(and_partial, schema)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_none, + PredicateBuilder::And({unsupported, unsupported})); + ASSERT_FALSE(Pushed(and_none, schema)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_full, + PredicateBuilder::Or({supported, other})); + ASSERT_TRUE(Pushed(or_full, schema)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_partial, + PredicateBuilder::Or({supported, unsupported})); + ASSERT_FALSE(Pushed(or_partial, schema)); +} + +TEST(VortexPredicateConverterTest, MembershipListBudgetCountsBooleanNodesAndNullGuards) { + const std::shared_ptr schema = TestSchema(); + std::vector literals; + for (int32_t value = 0; value < 33; ++value) { + const std::string text = fmt::format("value-{}", value); + literals.emplace_back(FieldType::STRING, text.data(), text.size()); + } + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals) + : PredicateBuilder::In(1, "s", FieldType::STRING, literals); + // ListContains lowers to 33 EQ + 32 OR nodes. NOT IN adds NOT plus the 3-node null guard. + const uint32_t nodes = 65 + (negate ? 4 : 0); + ASSERT_FALSE(Pushed(predicate, schema)); + ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); + ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); + ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::FORCE, 0})); + ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::OFF, nodes})); + } + literals.erase(literals.begin() + 32, literals.end()); + ASSERT_TRUE(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, literals), schema)); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); + literals.erase(literals.begin() + 31, literals.end()); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); + literals.erase(literals.begin() + 30, literals.end()); + ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); + literals.erase(literals.begin() + 21, literals.end()); + ASSERT_TRUE(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, literals), schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); +} + +TEST(VortexPredicateConverterTest, MembershipDeduplicatesBeforeApplyingBudget) { + const std::shared_ptr schema = TestSchema(); + std::vector strings(4096, Literal(FieldType::STRING, "same", 4)); + strings.emplace_back(FieldType::STRING, "other", 5); + const std::shared_ptr in = PredicateBuilder::In(1, "s", FieldType::STRING, strings); + const std::shared_ptr not_in = + PredicateBuilder::NotIn(1, "s", FieldType::STRING, strings); + ASSERT_FALSE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 2})); + ASSERT_TRUE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 3})); + ASSERT_FALSE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 6})); + ASSERT_TRUE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 7})); + + // Raw list length is not the decision: 4096 duplicates do not consume native nodes. + for (int32_t unique_count : {21, 32, 33}) { + std::vector repeated(4096, Literal(FieldType::STRING, "same", 4)); + for (int32_t value = 1; value < unique_count; ++value) { + const std::string text = fmt::format("value-{}", value); + repeated.emplace_back(FieldType::STRING, text.data(), text.size()); + } + ASSERT_EQ(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, repeated), schema), + unique_count <= 32); + ASSERT_EQ(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, repeated), schema), + unique_count <= 30); + } + + const std::vector integers(4096, Literal(int32_t{7})); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, integers), schema, + {MembershipPushdownMode::AUTO, 1})); + // The integer complement is one NEQ plus the 3-node null guard. + ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, integers), schema, + {MembershipPushdownMode::AUTO, 4})); +} + +TEST(VortexPredicateConverterTest, MembershipKeepsHugeContinuousSetsButBoundsManyRanges) { + const std::shared_ptr schema = TestSchema(); + std::vector literals; + for (int32_t value = 0; value < 65536; ++value) { + literals.emplace_back(value); + } + const std::shared_ptr in = PredicateBuilder::In(0, "i", FieldType::INT, literals); + const std::shared_ptr not_in = + PredicateBuilder::NotIn(0, "i", FieldType::INT, literals); + ASSERT_TRUE(Pushed(in, schema)); + ASSERT_TRUE(Pushed(not_in, schema)); + ASSERT_FALSE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 2})); + ASSERT_TRUE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 3})); + ASSERT_FALSE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 5})); + ASSERT_TRUE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 6})); + + literals.clear(); + for (int32_t start = 0; start < 200; start += 10) { + for (int32_t value = start; value < start + 8; ++value) { + literals.emplace_back(value); + } + } + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + // 20 ranges, each 3 nodes, joined by 19 connectives; NOT IN adds the null guard. + const uint32_t nodes = 79 + (negate ? 3 : 0); + ASSERT_FALSE(Pushed(predicate, schema)); + ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); + ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); + ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::FORCE, 0})); + } +} + +TEST(VortexPredicateConverterTest, MembershipBudgetCountsMixedRangesAndSingletons) { + const std::shared_ptr schema = TestSchema(); + std::vector literals = {Literal(int32_t{20}), Literal(int32_t{50})}; + for (int32_t value = 1; value <= 10; ++value) { + literals.emplace_back(value); + literals.emplace_back(value + 29); + } + // Two ranges + two singletons + three connectives = 11 nodes. + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + const uint32_t nodes = 11 + (negate ? 3 : 0); + ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); + ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); + } +} + +TEST(VortexPredicateConverterTest, ForceMembershipStillRejectsUnsafeLiterals) { + const std::shared_ptr schema = TestSchema(); + const VortexPredicateOptions force{MembershipPushdownMode::FORCE, 0}; + std::vector literals; + for (int32_t value = 0; value < 1024; ++value) { + literals.emplace_back(value); + } + for (const Literal& unsafe : {Literal(FieldType::INT), Literal(int64_t{1024})}) { + std::vector invalid = literals; + invalid.push_back(unsafe); + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, invalid), schema, force)); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, invalid), schema, force)); + } + ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), + arrow::schema({arrow::field("i", arrow::int64())}), force)); + for (bool negate : {false, true}) { + const std::vector timestamps = { + Literal(Timestamp::FromEpochMillis(0)), + Literal(Timestamp::FromEpochMillis(std::numeric_limits::max()))}; + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "ts", FieldType::TIMESTAMP, timestamps) + : PredicateBuilder::In(0, "ts", FieldType::TIMESTAMP, timestamps); + ASSERT_FALSE(Pushed( + predicate, arrow::schema({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::NANO))}), + force)); + } +} + +TEST(VortexPredicateConverterTest, MembershipOffPreservesOtherNativePredicates) { + const std::shared_ptr schema = TestSchema(); + const VortexPredicateOptions off{MembershipPushdownMode::OFF, 0}; + const std::shared_ptr equal = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int32_t{1})); + const std::shared_ptr in = + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1})}); + ASSERT_FALSE(Pushed(in, schema, off)); + ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(int32_t{1})}), + schema, off)); + ASSERT_TRUE(Pushed(equal, schema, off)); + ASSERT_TRUE(Pushed(PredicateBuilder::IsNull(0, "i", FieldType::INT), schema, off)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({in, equal})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({in, equal})); + ASSERT_TRUE(Pushed(and_predicate, schema, off)); + ASSERT_FALSE(Pushed(or_predicate, schema, off)); +} + +TEST(VortexPredicateConverterTest, SparseMembershipModesReturnNativeOrCandidateRows) { + NativeInCase test_case = {"sparse_budget", arrow::int32(), FieldType::INT, "[0,null,63,2]", {}}; + for (int32_t value = 0; value <= 64; value += 2) { + test_case.literals.emplace_back(value); + } + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "value", FieldType::INT, test_case.literals) + : PredicateBuilder::In(1, "value", FieldType::INT, test_case.literals); + const std::vector expected = + negate ? std::vector{2} : std::vector{0, 3}; + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, true, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, predicate, expected, true, + {MembershipPushdownMode::FORCE, 0})); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, true, + {MembershipPushdownMode::OFF}, false)); + } +} + +TEST(VortexPredicateConverterTest, StringMembershipModesAndPartialAndOrUseSafeCandidates) { + NativeInCase test_case = { + "string_budget", arrow::utf8(), FieldType::STRING, R"(["s0",null,"absent","s32"])", {}}; + for (int32_t value = 0; value < 33; ++value) { + const std::string text = fmt::format("s{}", value); + test_case.literals.emplace_back(FieldType::STRING, text.data(), text.size()); + } + const std::shared_ptr in = + PredicateBuilder::In(1, "value", FieldType::STRING, test_case.literals); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, in, {0, 1, 2, 3}, true, {}, false)); + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, in, {0, 3}, true, {MembershipPushdownMode::FORCE, 0})); + const std::shared_ptr less_than_three = + PredicateBuilder::LessThan(0, "rowid", FieldType::BIGINT, Literal(int64_t{3})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({in, less_than_three})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({in, less_than_three})); + // AND keeps the other child, so the residual receives a safe superset, including null rows. + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, and_predicate, {0, 1, 2}, true)); + // Dropping only the IN child of OR would lose row 3; reject the entire native OR instead. + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, or_predicate, {0, 1, 2, 3}, true, {}, false)); +} + +TEST(VortexPredicateConverterTest, NullPredicateIsNotPushed) { + ASSERT_FALSE(Pushed(nullptr, TestSchema())); +} + +} // namespace +} // namespace paimon::vortex::test diff --git a/src/paimon/format/vortex/vortex_reader_builder.h b/src/paimon/format/vortex/vortex_reader_builder.h index 42d5cfd17..74f69ec9d 100644 --- a/src/paimon/format/vortex/vortex_reader_builder.h +++ b/src/paimon/format/vortex/vortex_reader_builder.h @@ -24,14 +24,21 @@ #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/reader_builder.h" #include "paimon/format/vortex/vortex_file_batch_reader.h" +#include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/memory/memory_pool.h" namespace paimon::vortex { class VortexReaderBuilder : public ReaderBuilder { public: - explicit VortexReaderBuilder(int32_t batch_size) - : batch_size_(batch_size), pool_(GetDefaultPool()), arrow_pool_(GetArrowPool(pool_)) {} + explicit VortexReaderBuilder( + int32_t batch_size, const VortexPredicateOptions& predicate_options = {}, + VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE) + : batch_size_(batch_size), + predicate_options_(predicate_options), + coalesce_mode_(coalesce_mode), + pool_(GetDefaultPool()), + arrow_pool_(GetArrowPool(pool_)) {} ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { pool_ = pool; @@ -48,11 +55,14 @@ class VortexReaderBuilder : public ReaderBuilder { if (pool_ == nullptr) { return Status::Invalid("Vortex reader memory pool is nullptr"); } - return VortexFileBatchReader::Create(input, batch_size_, arrow_pool_); + return VortexFileBatchReader::Create(input, batch_size_, arrow_pool_, predicate_options_, + coalesce_mode_); } private: int32_t batch_size_; + VortexPredicateOptions predicate_options_; + VortexIoCoalesceMode coalesce_mode_; std::shared_ptr pool_; std::shared_ptr arrow_pool_; }; diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index b7cc73b96..bed173fa4 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -3727,6 +3727,233 @@ TEST_F(ScanAndReadInteTest, TestMosaicJavaAndPythonCompatibility) { #endif #ifdef PAIMON_ENABLE_VORTEX +namespace { + +Result> ReadVortexSdkRows( + const std::string& table_path, const std::vector>& splits, + const std::vector& projection, const std::shared_ptr& predicate, + bool exact_filter, bool late_materializing) { + ReadContextBuilder builder(table_path); + builder.SetReadFieldNames(projection) + .SetPredicate(predicate) + .EnablePredicateFilter(exact_filter) + .EnableLateMaterializing(late_materializing) + .EnablePrefetch(true) + .SetPrefetchBatchCount(2) + .SetPrefetchMaxParallelNum(2) + .AddOption(Options::READ_BATCH_SIZE, "2"); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read, TableRead::Create(std::move(context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, read->CreateReader(splits)); + return ReadResultCollector::CollectResult(std::move(reader)); +} + +// Append splits need not have a stable file order. Compare full rows as a multiset. +void AssertVortexSdkRows(const std::shared_ptr& actual, + const std::shared_ptr& type, + const std::string& expected_json, bool exact_filter = true) { + auto expected = arrow::ipc::internal::json::ArrayFromJSON(type, expected_json).ValueOrDie(); + if (!actual) { + ASSERT_EQ(expected->length(), 0); + return; + } + ASSERT_TRUE(actual->type()->Equals(type)) << actual->type()->ToString(); + if (exact_filter) { + ASSERT_EQ(actual->length(), expected->length()) << actual->ToString(); + } else { + ASSERT_GE(actual->length(), expected->length()); + } + arrow::ArrayVector candidates; + for (const auto& chunk : actual->chunks()) { + ASSERT_TRUE(chunk->ValidateFull().ok()); + for (int64_t row = 0; row < chunk->length(); ++row) { + candidates.push_back(chunk->Slice(row, 1)); + } + } + std::vector used(candidates.size(), false); + for (int64_t row = 0; row < expected->length(); ++row) { + bool found = false; + for (size_t candidate = 0; candidate < candidates.size(); ++candidate) { + if (!used[candidate] && expected->Slice(row, 1)->Equals(candidates[candidate])) { + used[candidate] = true; + found = true; + break; + } + } + ASSERT_TRUE(found) << "Missing row: " << expected->Slice(row, 1)->ToString() + << "\nActual: " << actual->ToString(); + } +} + +} // namespace + +TEST_F(ScanAndReadInteTest, TestVortexSdkPredicatesWithResidualAndLateMaterialization) { + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("tag", arrow::utf8()), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "-1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([ + [0,"skip","zero"],[1,null,"one"],[2,"keep","two"],[3,"drop","three"], + [4,"keep",null],[5,"keep","five"],[6,"skip","six"],[7,"keep","seven"] + ])", + {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 0, std::nullopt)); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + const auto result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), fields[2], fields[0], fields[1]}); + auto equal = PredicateBuilder::Equal(1, "id", FieldType::INT, Literal(int32_t{4})); + auto greater = PredicateBuilder::GreaterThan(1, "id", FieldType::INT, Literal(int32_t{1})); + auto tag = + PredicateBuilder::Equal(2, "tag", FieldType::STRING, Literal(FieldType::STRING, "keep", 4)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({greater, tag})); + // Sparse IN(64) exceeds the auto expression budget and requires the SDK residual filter. + std::vector sparse = {Literal(int32_t{2}), Literal(int32_t{5})}; + for (int32_t index = 0; index < 62; ++index) { + sparse.emplace_back(1000 + index * 2); + } + auto membership = PredicateBuilder::In(1, "id", FieldType::INT, sparse); + const std::vector, std::string>> cases = { + {equal, R"([[0,null,4,"keep"]])"}, + {conjunction, R"([[0,"two",2,"keep"],[0,null,4,"keep"], + [0,"five",5,"keep"],[0,"seven",7,"keep"]])"}, + {membership, R"([[0,"two",2,"keep"],[0,"five",5,"keep"]])"}}; + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (const auto& [predicate, expected] : cases) { + SCOPED_TRACE(predicate->ToString()); + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, {"payload", "id", "tag"}, predicate, + exact_filter, late_materializing)); + // A non-exact reader may prune candidates, but it must retain every matching row. + ASSERT_NO_FATAL_FAILURE( + AssertVortexSdkRows(actual, result_type, expected, exact_filter)); + } + } + } +} + +TEST_F(ScanAndReadInteTest, TestVortexSdkPrimaryKeyMergeBeforeValueFiltering) { + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector fields = {arrow::field("pk", arrow::utf8()), + arrow::field("value", arrow::int32()), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {"pk"}, options, + /*is_streaming_mode=*/true)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr original, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([["a",10,"old-a"],["b",20,"old-b"],["c",30,"old-c"]])", {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(original), 0, std::nullopt)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr changes, + TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([["a",1,"new-a"],["b",20,"old-b"],["d",40,"new-d"]])", {}, 0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(helper->WriteAndCommit(std::move(changes), 1, std::nullopt)); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + const auto result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1], fields[2]}); + auto value = PredicateBuilder::GreaterOrEqual(1, "value", FieldType::INT, Literal(int32_t{10})); + auto key = PredicateBuilder::In( + 0, "pk", FieldType::STRING, + {Literal(FieldType::STRING, "a", 1), Literal(FieldType::STRING, "b", 1), + Literal(FieldType::STRING, "c", 1)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr both, PredicateBuilder::And({key, value})); + const std::vector, std::string>> cases = { + {nullptr, R"([[0,"a",1,"new-a"],[0,"c",30,"old-c"],[0,"d",40,"new-d"]])"}, + {value, R"([[0,"c",30,"old-c"],[0,"d",40,"new-d"]])"}, + {key, R"([[0,"a",1,"new-a"],[0,"c",30,"old-c"]])"}, + {both, R"([[0,"c",30,"old-c"]])"}}; + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (const auto& [predicate, expected] : cases) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, {"pk", "value", "payload"}, predicate, + /*exact_filter=*/true, late_materializing)); + ASSERT_NO_FATAL_FAILURE(AssertVortexSdkRows(actual, result_type, expected)); + } + } +} + +TEST_F(ScanAndReadInteTest, TestVortexSdkSchemaEvolutionWithMissingRowIndexColumn) { + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector original_fields = {arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "-1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(original_fields), {}, {}, + options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(original_fields), + R"([[0,"old-zero"],[1,"old-one"],[2,null]])", {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 0, std::nullopt)); + helper.reset(); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + const arrow::FieldVector fields = {original_fields[0], original_fields[1], + arrow::field("added", arrow::utf8()), + arrow::field("_row_idx", arrow::int64())}; + ASSERT_OK(TestHelper::WriteNextSchema(dir->GetFileSystem(), table_path, + {DataField(0, fields[0]), DataField(1, fields[1]), + DataField(2, fields[2]), DataField(3, fields[3])}, + /*highest_field_id=*/3, options)); + ASSERT_OK_AND_ASSIGN(helper, TestHelper::Create(table_path, options, false)); + ASSERT_OK_AND_ASSIGN( + batch, TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([[3,"new-three","present",303]])", {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 1, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + const auto result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1], fields[2], fields[3]}); + auto greater = PredicateBuilder::GreaterOrEqual(0, "id", FieldType::INT, Literal(int32_t{1})); + auto missing = PredicateBuilder::IsNull(3, "_row_idx", FieldType::BIGINT); + const std::vector, std::string>> cases = { + {greater, R"([[0,1,"old-one",null,null],[0,2,null,null,null], + [0,3,"new-three","present",303]])"}, + {missing, R"([[0,0,"old-zero",null,null],[0,1,"old-one",null,null], + [0,2,null,null,null]])"}}; + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (const auto& [predicate, expected] : cases) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, {"id", "payload", "added", "_row_idx"}, + predicate, /*exact_filter=*/true, late_materializing)); + ASSERT_NO_FATAL_FAILURE(AssertVortexSdkRows(actual, result_type, expected)); + } + } +} + TEST_F(ScanAndReadInteTest, TestVortexJavaCompatibility) { // Fixture written by Paimon Java 2.0.0 with file.format=vortex; see // test/test_data/vortex/append_java_compat.db/append_java_compat/README.md. This verifies that From 63ea6093ec7d7c7a8256b27f74f42ef95d5192cb Mon Sep 17 00:00:00 2001 From: zhouyc-ali <301200591+zhouyc-ali@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:20:27 +0000 Subject: [PATCH 2/4] fix(vortex): reduce dense selection memory and fix CI checks Retain the smaller side of bitmap selections and preserve physical row mapping. Add regression coverage for dense, sparse, and reconfigured selections. Use callback I/O in native-scan test fixtures to avoid the Rust async filesystem path reported by TSAN, without adding suppressions. Add Apache license headers to both Vortex patch files. --- cmake_modules/vortex-fmm-cache-shards.diff | 17 +++ cmake_modules/vortex-point-read.diff | 17 +++ .../vortex/vortex_file_batch_reader.cpp | 76 +++++++--- .../format/vortex/vortex_file_batch_reader.h | 7 +- .../format/vortex/vortex_file_format_test.cpp | 141 ++++++++++++++++++ .../vortex_predicate_converter_test.cpp | 35 ++++- 6 files changed, 272 insertions(+), 21 deletions(-) diff --git a/cmake_modules/vortex-fmm-cache-shards.diff b/cmake_modules/vortex-fmm-cache-shards.diff index 8ce134137..801e18e3a 100644 --- a/cmake_modules/vortex-fmm-cache-shards.diff +++ b/cmake_modules/vortex-fmm-cache-shards.diff @@ -1,3 +1,20 @@ +# 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. + diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs diff --git a/cmake_modules/vortex-point-read.diff b/cmake_modules/vortex-point-read.diff index 8b311e1cd..8c24d3e58 100644 --- a/cmake_modules/vortex-point-read.diff +++ b/cmake_modules/vortex-point-read.diff @@ -1,3 +1,20 @@ +# 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. + diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.cpp b/src/paimon/format/vortex/vortex_file_batch_reader.cpp index b5791fc49..fa7aa6324 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.cpp +++ b/src/paimon/format/vortex/vortex_file_batch_reader.cpp @@ -201,7 +201,7 @@ VortexFileBatchReader::~VortexFileBatchReader() { } Result VortexFileBatchReader::OpenNextPartitionStream() { - if (has_selection_ && selection_row_ids_.empty()) { + if (has_selection_ && selected_row_count_ == 0) { return false; } if (scan_ == nullptr) { @@ -239,12 +239,13 @@ Result VortexFileBatchReader::OpenNextPartitionStream() { options.projection = projection.get(); } } - // VX_SELECTION_INCLUDE_RANGE maps to IncludeByIndex, not a contiguous range. - // The scan copies the individual indices from selection_row_ids_. - if (has_selection_) { - options.selection.idx = selection_row_ids_.data(); - options.selection.idx_len = selection_row_ids_.size(); - options.selection.include = VX_SELECTION_INCLUDE_RANGE; + // These FFI modes take individual indices. A dense bitmap sends only excluded rows; + // an all-selected bitmap needs no indices and uses the default INCLUDE_ALL mode. + if (has_selection_ && !selection_indices_.empty()) { + options.selection.idx = selection_indices_.data(); + options.selection.idx_len = selection_indices_.size(); + options.selection.include = + selection_excludes_ ? VX_SELECTION_EXCLUDE_RANGE : VX_SELECTION_INCLUDE_RANGE; } // Filtering drops rows; include_row_idx preserves their original physical positions. // The scan clones the filter, so it can be freed after scan creation. @@ -414,17 +415,40 @@ Status VortexFileBatchReader::SetReadSchema( if (read_schema == nullptr) { return Status::Invalid("Vortex read schema is nullptr"); } - // Keep selected physical row ids to map output positions after selection pushdown. - std::vector selection_row_ids; + std::vector selection_indices; + bool selection_excludes = false; + uint64_t selected_row_count = total_rows_; if (selection_bitmap) { - selection_row_ids.reserve(static_cast(selection_bitmap->Cardinality())); - for (auto iter = selection_bitmap->Begin(); iter != selection_bitmap->End(); ++iter) { + const int32_t cardinality = selection_bitmap->Cardinality(); + if (cardinality < 0 || static_cast(cardinality) > total_rows_) { + return Status::Invalid("Vortex selection cardinality exceeds the file row count"); + } + selected_row_count = static_cast(cardinality); + selection_excludes = selected_row_count > total_rows_ / 2; + selection_indices.reserve(static_cast( + selection_excludes ? total_rows_ - selected_row_count : selected_row_count)); + uint64_t next_row = 0; + const auto end = selection_bitmap->End(); + for (auto iter = selection_bitmap->Begin(); iter != end; ++iter) { const auto row_id = static_cast(*iter); if (row_id >= total_rows_) { return Status::Invalid( fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows_)); } - selection_row_ids.push_back(row_id); + if (selection_excludes) { + // Enumerate gaps directly instead of first expanding all selected rows. + for (; next_row < row_id; ++next_row) { + selection_indices.push_back(next_row); + } + next_row = row_id + 1; + } else { + selection_indices.push_back(row_id); + } + } + if (selection_excludes) { + for (; next_row < total_rows_; ++next_row) { + selection_indices.push_back(next_row); + } } } // Leave the active scan and its row mapping unchanged if validation or import fails. @@ -434,7 +458,10 @@ Status VortexFileBatchReader::SetReadSchema( read_schema_ = std::move(imported_read_schema); predicate_filter_ = std::dynamic_pointer_cast(predicate); has_selection_ = selection_bitmap.has_value(); - selection_row_ids_ = std::move(selection_row_ids); + selection_excludes_ = selection_excludes; + selected_row_count_ = selected_row_count; + selection_indices_ = std::move(selection_indices); + metrics_->SetCounter("vortex.selection-index-count", selection_indices_.size()); ReleaseStream(); current_batch_.reset(); current_batch_offset_ = 0; @@ -468,13 +495,28 @@ Result VortexFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch .Value(static_cast(batch_row_id)); } const uint64_t emitted_offset = previous_first_row_ + batch_row_id; - // An ordered selection emits rows in the same order as selection_row_ids_. + // Ordered scans emit selected rows in physical order, regardless of which side is stored. if (has_selection_) { - if (emitted_offset >= selection_row_ids_.size()) { + if (emitted_offset >= selected_row_count_) { return Status::Invalid(fmt::format("emitted row offset {} exceeds selection size {}", - emitted_offset, selection_row_ids_.size())); + emitted_offset, selected_row_count_)); + } + if (!selection_excludes_) { + return selection_indices_[emitted_offset]; + } + // For excluded row i, excluded[i] - i is the number of selected rows before it. + // Count exclusions at or before this selected rank without expanding a row mapping. + size_t begin = 0; + size_t end = selection_indices_.size(); + while (begin < end) { + const size_t middle = begin + (end - begin) / 2; + if (selection_indices_[middle] - middle <= emitted_offset) { + begin = middle + 1; + } else { + end = middle; + } } - return selection_row_ids_[emitted_offset]; + return emitted_offset + begin; } return emitted_offset; } diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.h b/src/paimon/format/vortex/vortex_file_batch_reader.h index b0a1180e9..5175b10ad 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.h +++ b/src/paimon/format/vortex/vortex_file_batch_reader.h @@ -118,9 +118,12 @@ class VortexFileBatchReader : public FileBatchReader { uint64_t previous_first_row_ = std::numeric_limits::max(); uint64_t previous_batch_row_count_ = 0; - // Selected physical row ids in output order, used when selection is pushed without a filter. + // Keep the smaller side of the bitmap. Dense selections store excluded physical row ids, + // avoiding an O(file rows) index list both here and in the native scan's copy. bool has_selection_ = false; - std::vector selection_row_ids_; + bool selection_excludes_ = false; + uint64_t selected_row_count_ = 0; + std::vector selection_indices_; // A pushed filter returns physical row ids in a hidden _row_idx column. Retain the last // batch's column because emitted offsets no longer identify physical rows after filtering. diff --git a/src/paimon/format/vortex/vortex_file_format_test.cpp b/src/paimon/format/vortex/vortex_file_format_test.cpp index c1ddffccc..ae8791de8 100644 --- a/src/paimon/format/vortex/vortex_file_format_test.cpp +++ b/src/paimon/format/vortex/vortex_file_format_test.cpp @@ -19,6 +19,7 @@ // Round-trip unit tests for the Vortex file format: write, then read back via vortex-ffi. +#include #include #include #include @@ -61,6 +62,7 @@ #include "paimon/fs/file_system.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/reader/batch_reader.h" @@ -1164,6 +1166,145 @@ TEST_F(VortexFileFormatTest, SelectionPushdownReturnsSelectedRowsAndPhysicalRowI ASSERT_EQ(got_row_ids, (std::vector{1, 3, 5, 7})); } +TEST_F(VortexFileFormatTest, DenseSelectionKeepsOnlyExcludedIndicesAndPreservesNullValues) { + constexpr int32_t kRows = 4096; + constexpr int32_t kDeletedRow = 2048; + arrow::Int32Builder id_builder; + arrow::Int32Builder value_builder; + for (int32_t row = 0; row < kRows; ++row) { + ASSERT_TRUE(id_builder.Append(row).ok()); + if (row % 17 == 0) { + ASSERT_TRUE(value_builder.AppendNull().ok()); + } else { + ASSERT_TRUE(value_builder.Append(row * 10).ok()); + } + } + std::shared_ptr ids; + std::shared_ptr values; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(value_builder.Finish(&values).ok()); + const auto schema = arrow::schema( + {arrow::field("id", arrow::int32(), false), arrow::field("value", arrow::int32())}); + const auto data = arrow::StructArray::Make({ids, values}, schema->fields()).ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "dense-selection.vortex"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/1024)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateReader(path, /*batch_size=*/127)); + RoaringBitmap32 selection; + selection.AddRange(0, kRows); + selection.RemoveRange(kDeletedRow, kDeletedRow + 1); + ASSERT_OK(ConfigureReader(reader.get(), schema, nullptr, selection)); + ASSERT_OK_AND_ASSIGN(uint64_t index_count, + reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); + // Both the host and the native FFI copy contain one excluded index, not kRows - 1 indices. + ASSERT_EQ(index_count, uint64_t{1}); + uint64_t emitted = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + const auto& batch = checked_cast(*array); + const auto& batch_ids = checked_cast(*batch.field(0)); + const auto& batch_values = checked_cast(*batch.field(1)); + // Reverse requests also check that physical-id lookup does not depend on call order. + for (int64_t i = array->length(); i > 0; --i) { + const uint64_t output_row = emitted + static_cast(i - 1); + const uint64_t expected_row = output_row + (output_row >= kDeletedRow ? 1 : 0); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i - 1))); + ASSERT_EQ(physical_row, expected_row); + ASSERT_EQ(batch_ids.Value(i - 1), static_cast(expected_row)); + ASSERT_EQ(batch_values.IsNull(i - 1), expected_row % 17 == 0); + if (!batch_values.IsNull(i - 1)) { + ASSERT_EQ(batch_values.Value(i - 1), static_cast(expected_row * 10)); + } + } + emitted += static_cast(array->length()); + } + ASSERT_EQ(emitted, uint64_t{kRows - 1}); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK_AND_ASSIGN(index_count, + reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); + ASSERT_EQ(index_count, uint64_t{1}); +} + +TEST_F(VortexFileFormatTest, DenseSelectionMapsConsecutiveExcludedRowsAndSparseSelections) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const std::vector> selections = { + {1, 2, 3, 4, 5}, {0, 1, 2, 3, 4}, {0, 1, 4, 5}, {2, 3, 4, 5}, + {0, 1, 2, 3}, {0, 2, 3, 5}, {0, 5}}; + for (const auto& selected_rows : selections) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From(selected_rows))); + ASSERT_OK_AND_ASSIGN(uint64_t index_count, reader->GetReaderMetrics()->GetCounter( + "vortex.selection-index-count")); + ASSERT_EQ(index_count, std::min(selected_rows.size(), size_t{6} - selected_rows.size())); + std::vector actual; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + const size_t offset = actual.size(); + actual.resize(offset + static_cast(array->length())); + for (int64_t i = array->length(); i > 0; --i) { + ASSERT_OK_AND_ASSIGN( + uint64_t row, reader->GetPreviousBatchFileRowId(static_cast(i - 1))); + actual[offset + static_cast(i - 1)] = row; + } + } + ASSERT_EQ(actual, (std::vector(selected_rows.begin(), selected_rows.end()))); + } +} + +TEST_F(VortexFileFormatTest, FullEmptyAndDenseSelectionsCanBeReconfiguredWithoutExpandedMapping) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + RoaringBitmap32 full; + full.AddRange(0, 6); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, full)); + ASSERT_OK_AND_ASSIGN(uint64_t index_count, + reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); + ASSERT_EQ(index_count, uint64_t{0}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); + ASSERT_NE(first, nullptr); + ASSERT_EQ(first->length(), 2); + ASSERT_OK_AND_ASSIGN(uint64_t row, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(row, uint64_t{1}); + + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32())); + ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); + ASSERT_EQ(first, nullptr); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From({1, 2, 3, 4, 5}))); + ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); + ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row, uint64_t{1}); + ASSERT_NOK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From({0, 1, 2, 3, 6}))); + ASSERT_NOK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, + RoaringBitmap32::From({-1, 1, 2, 3, 4}))); + ASSERT_OK_AND_ASSIGN(index_count, + reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); + ASSERT_EQ(index_count, uint64_t{1}); + ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(row, uint64_t{2}); + ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); + ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row, uint64_t{3}); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema())); + ASSERT_OK_AND_ASSIGN(index_count, + reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); + ASSERT_EQ(index_count, uint64_t{0}); + ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); + ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row, uint64_t{0}); +} + TEST_F(VortexFileFormatTest, FilterPushdownReturnsMatchingRowsAndPhysicalRowIds) { constexpr int64_t kRows = 10; arrow::Int32Builder id_builder; diff --git a/src/paimon/format/vortex/vortex_predicate_converter_test.cpp b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp index 3759218ea..ce63c7e54 100644 --- a/src/paimon/format/vortex/vortex_predicate_converter_test.cpp +++ b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp @@ -134,9 +134,40 @@ void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptr output, + file_system.Create(path, /*overwrite=*/false)); + auto output_context = std::make_shared(output); + vx_callback_sink* sink = nullptr; + bool output_closed = false; + ScopeGuard release_sink([&sink, &output_context, &output, &output_closed]() { + if (!output_closed) { + output_context->MarkAborted(); + if (sink != nullptr) { + vx_callback_sink_abort(std::exchange(sink, nullptr)); + } + output_context->WaitForIdleCallbacks(); + [[maybe_unused]] Status close_status = output->Close(); + } + }); + // Use the production callback I/O path, keeping Rust's asynchronous filesystem adapter + // out of this predicate test. The array owns the borrowed dtype until the sink closes. + const vx_dtype* dtype = vx_array_dtype(array.get()); + ASSERT_NE(dtype, nullptr); + sink = vx_callback_sink_open( + session.get(), VortexOutputContext::MakeCallbacks(output_context), dtype, &error); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_NE(sink, nullptr); + vx_callback_sink_push(sink, array.get(), &error); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + // Closing consumes the sink on both success and failure. + vx_callback_sink_close(std::exchange(sink, nullptr), &error); + ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); + ASSERT_OK(output_context->GetCallbackStatus()); + output_closed = true; + ASSERT_OK(output->Close()); + } ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system.Open(path)); ASSERT_OK_AND_ASSIGN(int64_t length, input->Length()); auto input_context = std::make_shared(input); From 61bd6a49a91cb3e28bde799ff40d7832ff2f5e1c Mon Sep 17 00:00:00 2001 From: zhouyc-ali <301200591+zhouyc-ali@users.noreply.github.com> Date: Wed, 4 Nov 2026 09:22:57 +0000 Subject: [PATCH 3/4] fix(vortex): address predicate pushdown review feedback Make membership and I/O limits configurable, simplify reader state, and bound native Arrow materialization, including nested arrays. Preserve safe temporal and decimal fallbacks, propagate type errors, and extend native and SDK regressions without changing shared readers. --- cmake_modules/ApplyVortexReadPatch.cmake | 59 +- cmake_modules/vortex-fmm-cache-shards.diff | 131 -- cmake_modules/vortex.diff | 407 ++++- crates/vortex_callback_io/callback_io.rs | 1469 ++++++++++++++++- docs/source/api/file_format.rst | 56 + docs/source/user_guide/read.rst | 20 - src/paimon/format/vortex/vortex_defs.h | 62 + src/paimon/format/vortex/vortex_ffi.h | 27 +- .../vortex/vortex_file_batch_reader.cpp | 425 ++--- .../format/vortex/vortex_file_batch_reader.h | 84 +- .../format/vortex/vortex_file_format.cpp | 51 +- .../format/vortex/vortex_file_format_test.cpp | 1282 +++++++++++--- .../format/vortex/vortex_format_writer.cpp | 46 +- .../format/vortex/vortex_io_callbacks.h | 6 +- .../vortex/vortex_predicate_converter.cpp | 502 +++--- .../vortex/vortex_predicate_converter.h | 22 +- .../vortex_predicate_converter_test.cpp | 1187 +++++++++++-- .../format/vortex/vortex_reader_builder.h | 15 +- .../format/vortex/vortex_stats_extractor.cpp | 6 +- test/inte/realtime_write_inte_test.cpp | 164 ++ test/inte/scan_and_read_inte_test.cpp | 393 ++++- 21 files changed, 5213 insertions(+), 1201 deletions(-) delete mode 100644 cmake_modules/vortex-fmm-cache-shards.diff create mode 100644 src/paimon/format/vortex/vortex_defs.h diff --git a/cmake_modules/ApplyVortexReadPatch.cmake b/cmake_modules/ApplyVortexReadPatch.cmake index 90b3c23d8..46129c30d 100644 --- a/cmake_modules/ApplyVortexReadPatch.cmake +++ b/cmake_modules/ApplyVortexReadPatch.cmake @@ -17,14 +17,29 @@ # A reused .patched marker can predate the scan/Arrow ABI additions. Validate the base patch # before changing any reader source; unknown or older sources need a fresh pinned checkout. -set(_base_files vortex-ffi/src/lib.rs vortex-ffi/src/scan.rs vortex-ffi/src/dtype.rs - vortex-ffi/cinclude/vortex.h) +set(_base_files + vortex-ffi/src/lib.rs + vortex-ffi/src/scan.rs + vortex-ffi/src/dtype.rs + vortex-ffi/cinclude/vortex.h + vortex-ffi/Cargo.toml + Cargo.lock + vortex-array/src/executor.rs + vortex-buffer/src/trusted_len.rs + vortex-array/src/arrow/executor/list.rs) set(_base_after 3fdbba15c951286ea0e565e3e21d75de1d7d4065a358f16aba2d47750e7ace76 00aba4a0d7339fe8d961fb302572deae58adb9f06eb9b29313aa9503e8352d49 - af3c5d2652a37cd8c7e7d1e2c377a290c9a5e0dac158bd4ed8d0b6922552bec0 - 1ad352c8476bfa41972dd597261fb2753b6903f4e90e9d606b85f9b04ed4c895) -foreach(_index RANGE 0 3) + c6e76821d117533e9b2448bef3842ec566def1384911abae1559410619c71e21 + 1ad352c8476bfa41972dd597261fb2753b6903f4e90e9d606b85f9b04ed4c895 + b24d3f3a83933fe4c061b1d8fd8e2ff963f2cf468eecd5e8ef045a7c039fa328 + 84d1415d7fa71c080cbd24c198712cf227429cf9fb508698e6c2e58674401410 + 2a2a7ff5bd7941899fcb1d9ad4262263bc183920041b56bc9bd328a43d76c9e0 + 9779c7a90f839669dd204cd1c3e3fbcace9fd6ca19583ec0deb51d52b51783e7 + 56de0a4a7e7dca7e3619790a7058f9109d2b56a77eaa18ea1ec2987cc17d8dfe) +list(LENGTH _base_files _base_count) +math(EXPR _base_last "${_base_count} - 1") +foreach(_index RANGE 0 ${_base_last}) list(GET _base_files ${_index} _file) list(GET _base_after ${_index} _expected) file(SHA256 "${VORTEX_SOURCE_DIR}/${_file}" _actual) @@ -34,7 +49,8 @@ foreach(_index RANGE 0 3) endif() endforeach() -# Accept clean reader sources, the reader patch alone, or both reader patches. +# Accept clean reader sources or the reader patch alone. Also check the cache sources so a +# reused checkout cannot silently retain an obsolete shard override. set(_files vortex-layout/src/layouts/dict/reader.rs vortex-ffi/src/session.rs @@ -47,42 +63,28 @@ set(_before d77641a60e1d4daf84f58f30d5cc0f5f64bc09f30ff1d711862b66606ec9a772 e75aefbaa363748ab9e75cc21b1ec6d74701e2c973d1eb204e657c4e8076e7f6 2da8289cbd44c64bcbdf102cdc2be291448fc3aa77a0ca34d2f7efb7dfde2419) -set(_reader_patched +set(_after 4bfb5f4bf8f707027702e9157d2d40d1b64445692d766a3a2d9f685998af3806 093cf553ed52c01a31ff131b96f226157f37ffe67593136d3bd67832f85e362e d77641a60e1d4daf84f58f30d5cc0f5f64bc09f30ff1d711862b66606ec9a772 e75aefbaa363748ab9e75cc21b1ec6d74701e2c973d1eb204e657c4e8076e7f6 2da8289cbd44c64bcbdf102cdc2be291448fc3aa77a0ca34d2f7efb7dfde2419) -set(_after - 4bfb5f4bf8f707027702e9157d2d40d1b64445692d766a3a2d9f685998af3806 - 093cf553ed52c01a31ff131b96f226157f37ffe67593136d3bd67832f85e362e - 79be1013b9f04e683de40f8a9881e4d15e07945dd8d86b7760d3abf01d14ceb6 - 36d599cbea7c0bd1855d560c2e171e9b61117f4e69cfe3bac4dcd0b126cb549f - 83edf1951eb564ab86e28a6e5f1c9fdb49788fc42ba43d7d26db9fd8f08cf40d) -set(_cache_patch "${CMAKE_CURRENT_LIST_DIR}/vortex-fmm-cache-shards.diff") set(_all_clean TRUE) -set(_all_reader_patched TRUE) set(_all_after TRUE) foreach(_index RANGE 0 4) list(GET _files ${_index} _file) list(GET _before ${_index} _old) - list(GET _reader_patched ${_index} _intermediate) list(GET _after ${_index} _new) file(SHA256 "${VORTEX_SOURCE_DIR}/${_file}" _actual) if(NOT _actual STREQUAL _old) set(_all_clean FALSE) endif() - if(NOT _actual STREQUAL _intermediate) - set(_all_reader_patched FALSE) - endif() if(NOT _actual STREQUAL _new) set(_all_after FALSE) endif() endforeach() -if(NOT _all_clean - AND NOT _all_reader_patched - AND NOT _all_after) - message(FATAL_ERROR "Vortex read patch: unknown or partly updated source; preserve and inspect local changes" +if(NOT _all_clean AND NOT _all_after) + message(FATAL_ERROR "Vortex read patch: unknown, obsolete or partly updated source; preserve local changes and use a fresh pinned source checkout" ) endif() @@ -101,24 +103,15 @@ function(_run_vortex_read_patch _path _direction _dry_run) endif() endfunction() -# The patches affect disjoint files, so both dry-runs complete before any write. +# Complete the dry-run before changing any reader source. if(_all_clean) _run_vortex_read_patch("${VORTEX_READ_PATCH}" --forward TRUE) else() _run_vortex_read_patch("${VORTEX_READ_PATCH}" --reverse TRUE) endif() -if(_all_after) - _run_vortex_read_patch("${_cache_patch}" --reverse TRUE) -else() - _run_vortex_read_patch("${_cache_patch}" --forward TRUE) -endif() - if(_all_clean) _run_vortex_read_patch("${VORTEX_READ_PATCH}" --forward FALSE) endif() -if(NOT _all_after) - _run_vortex_read_patch("${_cache_patch}" --forward FALSE) -endif() foreach(_index RANGE 0 4) list(GET _files ${_index} _file) list(GET _after ${_index} _expected) diff --git a/cmake_modules/vortex-fmm-cache-shards.diff b/cmake_modules/vortex-fmm-cache-shards.diff deleted file mode 100644 index 801e18e3a..000000000 --- a/cmake_modules/vortex-fmm-cache-shards.diff +++ /dev/null @@ -1,131 +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. - -diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs ---- a/vortex-layout/src/layouts/struct_/reader.rs -+++ b/vortex-layout/src/layouts/struct_/reader.rs -@@ -115,7 +115,8 @@ - expanded_root_expr, - lazy_children, - field_lookup, -- partitioned_expr_cache: Default::default(), -+ // Each file evaluates few expressions; avoid CPU-count-sized shard allocations. -+ partitioned_expr_cache: DashMap::with_hasher_and_shard_amount(Default::default(), 16), - }) - } - -diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs ---- a/vortex-file/src/v2/file_stats_reader.rs -+++ b/vortex-file/src/v2/file_stats_reader.rs -@@ -74,7 +74,8 @@ - file_stats, - struct_fields, - session, -- prune_cache: Default::default(), -+ // File-level pruning reuses only a small set of expressions per reader. -+ prune_cache: DashMap::with_hasher_and_shard_amount(Default::default(), 16), - } - } - -diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs ---- a/vortex-layout/src/segments/shared.rs -+++ b/vortex-layout/src/segments/shared.rs -@@ -32,7 +32,8 @@ - pub fn new(inner: S) -> Self { - Self { - inner, -- in_flight: DashMap::default(), -+ // A per-file request set does not need a shard for every available CPU. -+ in_flight: DashMap::with_hasher_and_shard_amount(Default::default(), 16), - } - } - } -@@ -65,10 +66,16 @@ - - #[cfg(test)] - mod tests { -+ use std::sync::Barrier; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; -- -+ use std::thread; -+ -+ use futures::executor::block_on; -+ use futures::future::join_all; - use vortex_buffer::ByteBuffer; -+ use vortex_error::VortexResult; -+ use vortex_error::vortex_err; - - use super::*; - use crate::segments::SegmentSink; -@@ -119,5 +126,56 @@ - } - - #[tokio::test] -+ async fn test_shared_source_deduplicates_many_segments_across_threads() -> VortexResult<()> { -+ let source = CountingSegmentSource::default(); -+ let mut ids = Vec::new(); -+ let mut expected = Vec::new(); -+ for value in 0..64u8 { -+ let data = ByteBuffer::from(vec![value; 4]); -+ ids.push( -+ source -+ .segments -+ .write(SequenceId::root().downgrade(), vec![data.clone()]) -+ .await?, -+ ); -+ expected.push(data); -+ } -+ -+ let shared_source = SharedSegmentSource::new(source.clone()); -+ let barrier = Barrier::new(8); -+ thread::scope(|scope| -> VortexResult<()> { -+ let mut requests = Vec::new(); -+ for _ in 0..8 { -+ let shared_source = &shared_source; -+ let ids = &ids; -+ let expected = &expected; -+ let barrier = &barrier; -+ requests.push(scope.spawn(move || -> VortexResult<()> { -+ let pending = ids -+ .iter() -+ .map(|id| shared_source.request(*id)) -+ .collect::>(); -+ // Hold every request until all threads share the same in-flight futures. -+ barrier.wait(); -+ for (actual, expected) in block_on(join_all(pending)).into_iter().zip(expected) -+ { -+ assert_eq!(actual?.unwrap_host(), *expected); -+ } -+ Ok(()) -+ })); -+ } -+ for request in requests { -+ request -+ .join() -+ .map_err(|_| vortex_err!("segment requester panicked"))??; -+ } -+ Ok(()) -+ })?; -+ -+ assert_eq!(source.request_count.load(Ordering::SeqCst), ids.len()); -+ Ok(()) -+ } -+ -+ #[tokio::test] - async fn test_shared_source_handles_dropped_futures() { - let source = CountingSegmentSource::default(); diff --git a/cmake_modules/vortex.diff b/cmake_modules/vortex.diff index d1b088f8e..1733405c4 100644 --- a/cmake_modules/vortex.diff +++ b/cmake_modules/vortex.diff @@ -1,16 +1,20 @@ diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs +index f4ce314c07373fa5b89f9225da1e6162ce956443..af526729ba3672eb1d507e54ebf14c750c83f8e1 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs -@@ -8,4 +8,5 @@ +@@ -6,6 +6,7 @@ + //! Native interface to Vortex arrays, types, files and streams. + mod array; +mod callback_io; mod array_iterator; mod binary; mod data_source; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs +index 921c4f506211b299751f38645fdbde6155c9ad86..8c4560d2d5c88b8479a4a8b47b692924860d415b 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs -@@ -25,8 +25,11 @@ +@@ -25,8 +25,11 @@ use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; @@ -22,7 +26,7 @@ diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs use vortex::layout::scan::arrow::RecordBatchIteratorAdapter; use vortex::scan::DataSourceScan; use vortex::scan::Partition; -@@ -112,4 +115,7 @@ +@@ -112,6 +115,9 @@ pub struct vx_scan_options { pub limit: u64, /// If true, return in storage order. pub ordered: bool, @@ -30,7 +34,11 @@ diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs + /// row index, stable across `filter` and `selection`. The host strips it before use. + pub include_row_idx: bool, } -@@ -182,5 +188,12 @@ + + #[repr(C)] +@@ -180,7 +186,14 @@ fn scan_request(opts: *const vx_scan_options) -> VortexResult { + let limit = (opts.limit != 0).then_some(opts.limit); + Ok(ScanRequest { - projection, + // paimon-cpp: when include_row_idx is set, append a hidden `_row_idx` column so the host can @@ -44,31 +52,38 @@ diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs filter, row_range, selection, -@@ -356,3 +369,5 @@ +@@ -353,7 +366,9 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( + let array_stream = partition.execute()?; + let dtype = array_stream.dtype(); + - let schema = dtype.to_arrow_schema()?; + // paimon-cpp: export standard (non-view) string/binary types so the C++ reader skips its + // per-row view->standard rebuild; execute_arrow casts to this target internally. + let schema = crate::callback_io::normalize_arrow_schema_views(dtype.to_arrow_schema()?); let schema = Arc::new(schema); let target = Field::new_struct("", schema.fields().clone(), false); + diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs +index c74fa71b6edf43710d446a6e97a5ca936428ad67..d06a40f63b6c53fd359c0eecfeae9cd352c60bd7 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs -@@ -344,7 +344,9 @@ +@@ -344,7 +344,10 @@ pub unsafe extern "C-unwind" fn vx_dtype_to_arrow_schema( ) -> c_int { try_or(err, 1, || { let dtype = vx_dtype::as_ref(dtype); - let arrow_schema = dtype.to_arrow_schema()?; + // paimon-cpp: report standard (non-view) string/binary types so the exported schema matches + // the data that vx_partition_scan_arrow now produces (see callback_io normalize). -+ let arrow_schema = crate::callback_io::normalize_arrow_schema_views(dtype.to_arrow_schema()?); ++ let arrow_schema = ++ crate::callback_io::normalize_arrow_schema_views(dtype.to_arrow_schema()?); let arrow_schema = FFI_ArrowSchema::try_from(&arrow_schema)?; unsafe { ptr::write(schema, arrow_schema) }; Ok(0) diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h +index 26c7b409cd982abce17fbb41a131814b7613018c..75da718b9ae99bdf4420893374af862ddb501782 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h -@@ -622,4 +622,9 @@ +@@ -622,6 +622,11 @@ typedef struct { * If true, return in storage order. */ bool ordered; @@ -78,3 +93,379 @@ diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h + */ + bool include_row_idx; } vx_scan_options; + + #ifdef __cplusplus +diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml +index 6407f97fe46303f6dc9bd6b7f1f0f44392f5e534..426a840263dc6513187c7bd6ae8f24433b4a440f 100644 +--- a/vortex-ffi/Cargo.toml ++++ b/vortex-ffi/Cargo.toml +@@ -25,6 +25,7 @@ arrow-schema = { workspace = true } + async-fs = { workspace = true } + futures = { workspace = true } + itertools = { workspace = true } ++jiff = { workspace = true } + mimalloc = { workspace = true, optional = true } + object_store = { workspace = true, features = ["aws", "azure", "gcp"] } + paste = { workspace = true } +@@ -35,6 +36,7 @@ url = { workspace = true, features = [] } + vortex = { workspace = true, features = ["object_store"] } + + [dev-dependencies] ++parking_lot = { workspace = true } + rand = { workspace = true } + tempfile = { workspace = true } + vortex-array = { workspace = true, features = ["_test-harness"] } +diff --git a/Cargo.lock b/Cargo.lock +index fd66ecc567062fe122290bdb541bb3fa4420babe..731e2f40c163e9ef5b1feae672aceeccb8760308 100644 +--- a/Cargo.lock ++++ b/Cargo.lock +@@ -9747,8 +9747,10 @@ dependencies = [ + "cbindgen", + "futures", + "itertools 0.14.0", ++ "jiff", + "mimalloc", + "object_store 0.13.2", ++ "parking_lot", + "paste", + "prost 0.14.4", + "rand 0.10.1", +diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs +index d6070ac1a4db34380f8aca8b463e8d513f5d1757..071591448281f19bd770f02e74ad0d839acc1801 100644 +--- a/vortex-array/src/executor.rs ++++ b/vortex-array/src/executor.rs +@@ -26,7 +26,6 @@ use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_error::vortex_ensure; + use vortex_error::vortex_panic; +-use vortex_session::Ref; + use vortex_session::SessionExt; + use vortex_session::VortexSession; + +@@ -424,14 +423,9 @@ impl Executable for ArrayRef { + } + } + +- let tmp_session = ctx.session().clone(); +- let kernels = tmp_session.get_opt::(); +- + for (slot_idx, slot) in array.slots().iter().enumerate() { + let Some(child) = slot else { continue }; +- if let Some(executed_parent) = +- execute_parent_for_child(&array, child, slot_idx, kernels.as_ref(), ctx)? +- { ++ if let Some(executed_parent) = execute_parent_for_child(&array, child, slot_idx, ctx)? { + ctx.log(format_args!( + "execute_parent: slot[{}]({}) rewrote {} -> {}", + slot_idx, +@@ -542,13 +536,15 @@ fn execute_parent_for_child( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, +- kernels: Option<&Ref>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { +- if let Some(kernels) = kernels +- && let Some(plugins) = +- kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) +- { ++ // Session refs hold a DashMap shard lock. Release it before invoking kernels or encodings, ++ // which may lazily initialize another session variable on the same shard. ++ let plugins = ctx ++ .session() ++ .get_opt::() ++ .and_then(|kernels| kernels.find_execute_parent(parent.encoding_id(), child.encoding_id())); ++ if let Some(plugins) = plugins { + for plugin in plugins.as_ref() { + if let Some(result) = plugin(child, parent, slot_idx, ctx)? { + return Ok(Some(result)); +@@ -561,14 +557,9 @@ fn execute_parent_for_child( + + /// Try execute_parent on each occupied slot of the array. + fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { +- let tmp_session = ctx.session().clone(); +- let kernels = tmp_session.get_opt::(); +- + for (slot_idx, slot) in array.slots().iter().enumerate() { + let Some(child) = slot else { continue }; +- if let Some(executed_parent) = +- execute_parent_for_child(array, child, slot_idx, kernels.as_ref(), ctx)? +- { ++ if let Some(executed_parent) = execute_parent_for_child(array, child, slot_idx, ctx)? { + ctx.log(format_args!( + "execute_parent: slot[{}]({}) rewrote {} -> {}", + slot_idx, +@@ -818,3 +809,57 @@ impl VortexSessionExecute for VortexSession { + ExecutionCtx::new(self.clone()) + } + } ++ ++#[cfg(test)] ++mod tests { ++ use std::sync::atomic::AtomicUsize; ++ use std::sync::atomic::Ordering; ++ ++ use vortex_buffer::buffer; ++ ++ use super::*; ++ use crate::arrays::ScalarFnArray; ++ use crate::assert_arrays_eq; ++ use crate::scalar_fn::TypedScalarFnInstance; ++ use crate::scalar_fn::fns::binary::Binary; ++ use crate::scalar_fn::fns::operators::Operator; ++ ++ static REENTRANT_KERNEL_CALLS: AtomicUsize = AtomicUsize::new(0); ++ ++ fn reentrant_kernel( ++ _child: &ArrayRef, ++ _parent: &ArrayRef, ++ _child_idx: usize, ++ ctx: &mut ExecutionCtx, ++ ) -> VortexResult> { ++ // Reacquiring this exact SessionVar for writing must not inherit the registry's read lock. ++ // Unlike relying on two TypeIds sharing a shard, this fails deterministically if it does. ++ drop(ctx.session().get_mut::()); ++ REENTRANT_KERNEL_CALLS.fetch_add(1, Ordering::SeqCst); ++ Ok(Some(buffer![11_i32, 22].into_array())) ++ } ++ ++ #[test] ++ fn execute_parent_kernel_can_reenter_session() -> VortexResult<()> { ++ let lhs = buffer![1_i32, 2].into_array(); ++ let rhs = buffer![10_i32, 20].into_array(); ++ let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased(); ++ let parent = ScalarFnArray::try_new(scalar_fn, vec![lhs.clone(), rhs])?.into_array(); ++ let session = VortexSession::empty().with::(); ++ session.get::().register_execute_parent( ++ parent.encoding_id(), ++ lhs.encoding_id(), ++ &[reentrant_kernel], ++ ); ++ let mut ctx = ExecutionCtx::new(session); ++ let expected = buffer![11_i32, 22].into_array(); ++ ++ // Cover both the single-step executor and the matcher-driven execution loop. ++ let single_step = parent.clone().execute::(&mut ctx)?; ++ assert_arrays_eq!(single_step, expected); ++ let canonical = parent.execute::(&mut ctx)?.into_array(); ++ assert_arrays_eq!(canonical, expected); ++ assert_eq!(REENTRANT_KERNEL_CALLS.load(Ordering::SeqCst), 2); ++ Ok(()) ++ } ++} +diff --git a/vortex-buffer/src/trusted_len.rs b/vortex-buffer/src/trusted_len.rs +index 13cf25d0546428dfc719e077b696d352242a0b9e..f8ce4fc04d68085c5aba75494e98a447f13826b4 100644 +--- a/vortex-buffer/src/trusted_len.rs ++++ b/vortex-buffer/src/trusted_len.rs +@@ -32,6 +32,11 @@ impl Iterator for TrustedLenAdapter { + + #[inline] + fn next(&mut self) -> Option { ++ #[cfg(not(debug_assertions))] ++ { ++ self.inner.next() ++ } ++ #[cfg(debug_assertions)] + match self.inner.next() { + None => { + #[cfg(debug_assertions)] +diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs +index e8bf95e14cdd3faed26d9fcc464da8276b523001..a002d7926eded2941863f94968329345e9d075eb 100644 +--- a/vortex-array/src/arrow/executor/list.rs ++++ b/vortex-array/src/arrow/executor/list.rs +@@ -12,26 +12,33 @@ use vortex_buffer::BufferMut; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_ensure; ++use vortex_error::vortex_err; + + use crate::ArrayRef; + use crate::Canonical; + use crate::ExecutionCtx; ++use crate::IntoArray; + use crate::arrays::Chunked; ++use crate::arrays::ChunkedArray; + use crate::arrays::List; + use crate::arrays::ListArray; + use crate::arrays::ListView; + use crate::arrays::ListViewArray; ++use crate::arrays::PrimitiveArray; + use crate::arrays::chunked::ChunkedArrayExt; + use crate::arrays::list::ListArrayExt; + use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewDataParts; + use crate::arrays::listview::ListViewRebuildMode; ++use crate::arrays::listview::list_view_from_list; + use crate::arrow::executor::validity::to_arrow_null_buffer; + use crate::arrow::session::ArrowSessionExt; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::dtype::Nullability; ++use crate::dtype::PType; ++use crate::validity::Validity; + + #[allow(rustdoc::broken_intra_doc_links)] + /// Convert a Vortex VarBinArray into an Arrow [`GenericListArray`](arrow_array:array::GenericListArray). +@@ -42,7 +49,15 @@ pub(super) fn to_arrow_list( + ) -> VortexResult { + // If the Vortex array is already in List format, we can directly convert it. + if let Some(array) = array.as_opt::() { +- return list_to_list::(&array.into_owned(), elements_field, ctx); ++ let array = array.into_owned(); ++ // NULL list slots may own non-empty child ranges. They are not logically referenced ++ // and must not be expanded even when they lie between two valid lists. ++ let array = if array.validity()?.execute_no_nulls(array.len(), ctx)? { ++ array ++ } else { ++ compact_list_view_elements(&list_view_from_list(array, ctx)?, ctx)? ++ }; ++ return list_to_list::(&array, elements_field, ctx); + } + + // Converting each chunk individually, then using the fast concat logic from arrow +@@ -56,30 +71,16 @@ pub(super) fn to_arrow_list( + return Ok(arrow_select::concat::concat(&refs)?); + } + +- // If the Vortex array is a ListViewArray, rebuild to ZCTL if needed and convert. ++ // If the Vortex array is a ListViewArray, select only its referenced elements. + let array = match array.try_downcast::() { +- Ok(array) => { +- let zctl = if array.is_zero_copy_to_list() { +- array +- } else { +- array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? +- }; +- return list_view_zctl::(zctl, elements_field, ctx); +- } ++ Ok(array) => return list_view_to_list::(array, elements_field, ctx), + Err(a) => a, + }; + +- // Otherwise, we execute the array to become a ListViewArray, then rebuild to ZCTL. +- // Note: arrow_cast::cast supports ListView → List (apache/arrow-rs#8735), but it +- // unconditionally uses take. Our rebuild uses a heuristic that picks list-by-list +- // for large lists, which avoids materializing a large index buffer. ++ // Otherwise, execute to ListView but keep shared children lazy until the referenced ranges ++ // have been selected. Rebuilding non-ZCTL views first can canonicalize unreferenced children. + let list_view = array.execute::(ctx)?; +- let zctl = if list_view.is_zero_copy_to_list() { +- list_view +- } else { +- list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? +- }; +- list_view_zctl::(zctl, elements_field, ctx) ++ list_view_to_list::(list_view, elements_field, ctx) + } + + #[allow(rustdoc::broken_intra_doc_links)] +@@ -89,6 +90,11 @@ fn list_to_list( + elements_field: &FieldRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { ++ // Slicing a Vortex List retains its complete elements array. Trim before exporting children ++ // so a bounded parent batch cannot materialize unreferenced string/binary values. Resetting ++ // before narrowing offsets also avoids overflow from a large unused leading prefix. ++ let array = array.reset_offsets(false, ctx)?; ++ + // We must cast the offsets to the required offset type. + let offsets = array + .offsets() +@@ -119,16 +125,22 @@ fn list_to_list( + ))) + } + +-fn list_view_zctl( ++fn list_view_to_list( + array: ListViewArray, + elements_field: &FieldRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { +- assert!(array.is_zero_copy_to_list()); ++ if !array.is_zero_copy_to_list() || !array.validity()?.execute_no_nulls(array.len(), ctx)? { ++ return list_to_list::( ++ &compact_list_view_elements(&array, ctx)?, ++ elements_field, ++ ctx, ++ ); ++ } + + if array.is_empty() { + let elements = ctx.session().clone().arrow().execute_arrow( +- array.elements().clone(), ++ array.elements().slice(0..0)?, + Some(elements_field.as_ref()), + ctx, + )?; +@@ -140,6 +152,10 @@ fn list_view_zctl( + ))); + } + ++ // ZCTL views have contiguous referenced elements, but a sliced view can still retain the ++ // complete native chunk. Trim and rebase before converting children or narrowing offsets. ++ let array = array.rebuild(ListViewRebuildMode::TrimElements, ctx)?; ++ + let ListViewDataParts { + elements, + offsets, +@@ -200,6 +216,59 @@ fn list_view_zctl( + ))) + } + ++// Assemble only referenced slices in output order. In particular, a take of arbitrary encoded ++// children may execute the entire values array before selecting indices; a ListView rebuild may ++// do the same. Keeping the ranges as chunks lets recursive Arrow export materialize each slice. ++fn compact_list_view_elements( ++ array: &ListViewArray, ++ ctx: &mut ExecutionCtx, ++) -> VortexResult { ++ let validity = array.validity()?; ++ let valid = validity.execute_mask(array.len(), ctx)?; ++ let offsets = array ++ .offsets() ++ .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? ++ .execute::(ctx)? ++ .to_buffer::(); ++ let sizes = array ++ .sizes() ++ .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? ++ .execute::(ctx)? ++ .to_buffer::(); ++ ++ let mut chunks = Vec::with_capacity(array.len()); ++ let mut list_offsets = BufferMut::::with_capacity(array.len() + 1); ++ let mut n_elements = 0u64; ++ list_offsets.push(n_elements); ++ for index in 0..array.len() { ++ if valid.value(index) && sizes[index] != 0 { ++ let start = usize::try_from(offsets[index]) ++ .map_err(|error| vortex_err!("List offset out of range: {error}"))?; ++ let size = usize::try_from(sizes[index]) ++ .map_err(|error| vortex_err!("List size out of range: {error}"))?; ++ let end = start ++ .checked_add(size) ++ .ok_or_else(|| vortex_err!("List element range overflow"))?; ++ n_elements = n_elements ++ .checked_add(sizes[index]) ++ .ok_or_else(|| vortex_err!("List output element count overflow"))?; ++ chunks.push(array.elements().slice(start..end)?); ++ } ++ list_offsets.push(n_elements); ++ } ++ ++ let elements = if chunks.is_empty() { ++ Canonical::empty(array.elements().dtype()).into_array() ++ } else { ++ ChunkedArray::try_new(chunks, array.elements().dtype().clone())?.into_array() ++ }; ++ ListArray::try_new( ++ elements, ++ PrimitiveArray::new(list_offsets.freeze(), Validity::NonNullable).into_array(), ++ validity, ++ ) ++} ++ + #[cfg(test)] + mod tests { + use std::sync::Arc; diff --git a/crates/vortex_callback_io/callback_io.rs b/crates/vortex_callback_io/callback_io.rs index 2c49e19ba..bb5f5d7f5 100644 --- a/crates/vortex_callback_io/callback_io.rs +++ b/crates/vortex_callback_io/callback_io.rs @@ -14,6 +14,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![allow(non_camel_case_types)] + //! Callback-based I/O for the Vortex C API. //! //! The stock C API can only read a whole file already materialized in memory @@ -48,10 +50,17 @@ use std::ffi::CStr; use std::ffi::c_char; +use std::ffi::c_int; use std::ffi::c_void; use std::io; +use std::ptr::write; use std::sync::Arc; +use arrow_array::RecordBatch; +use arrow_array::cast::AsArray; +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_schema::ArrowError; +use arrow_schema::Field; use futures::FutureExt; use futures::SinkExt; use futures::TryStreamExt; @@ -59,6 +68,8 @@ use futures::channel::mpsc; use futures::channel::mpsc::Sender; use futures::future::BoxFuture; use vortex::array::ArrayRef; +use vortex::array::VortexSessionExecute; +use vortex::array::arrow::ArrowSessionExt; use vortex::array::buffer::BufferHandle; use vortex::array::stream::ArrayStreamAdapter; use vortex::buffer::Alignment; @@ -70,6 +81,7 @@ use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::extension::datetime::Date; +use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex::file::OpenOptionsSessionExt; @@ -84,6 +96,7 @@ use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Handle; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::layout::scan::arrow::RecordBatchIteratorAdapter; use vortex::scalar::PValue; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; @@ -96,6 +109,8 @@ use crate::error::try_or; use crate::error::try_or_default; use crate::error::vx_error; use crate::scalar::vx_scalar; +use crate::scan::VxPartitionScan; +use crate::scan::vx_partition; use crate::session::vx_session; /// Fill `length` bytes starting at `offset` into `dst`. @@ -299,42 +314,20 @@ unsafe fn data_source_new_callback( Ok(vx_data_source::new(file.data_source()?)) } -/// Create a data source that reads through host callbacks. -/// -/// `size` is the total length of the file in bytes; the host knows it up front, so it is passed -/// here instead of being fetched through another callback. +/// Create a data source that reads through host callbacks with explicit I/O coalescing. /// -/// The callbacks (and the context they carry) are owned by the returned data source and released -/// when it is freed, including when this call fails. +/// `size` is the total file length in bytes. `max_size` must be positive, and `max_gap` must +/// not exceed it. A smaller gap avoids reading unrelated bytes between sparse requests. /// -/// On error, returns NULL and sets "err". +/// The callbacks and their context are owned by the returned data source and released when +/// it is freed. They are consumed and released even when this call fails, including invalid +/// coalescing options. On error, returns NULL and sets `err`. /// /// # Safety /// /// `session` must be a valid `vx_session`. `callbacks.read_at_fn` must be non-null, thread-safe, -/// and valid for as long as the data source lives. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_new_callback( - session: *const vx_session, - callbacks: vx_input_callbacks, - size: u64, - err: *mut *mut vx_error, -) -> *const vx_data_source { - try_or(err, std::ptr::null(), || unsafe { - data_source_new_callback(session, callbacks, size, CoalesceConfig::object_storage()) - }) -} - -/// Create a callback data source with an explicit I/O coalescing policy. -/// -/// The original entry point retains its object-storage defaults. Low-latency hosts can choose a -/// smaller gap to avoid reading unrelated bytes between sparse requests. `max_size` must be -/// positive and `max_gap` must not exceed it. The callbacks are consumed even on invalid options. -/// -/// # Safety -/// -/// Same requirements as [`vx_data_source_new_callback`]. `err` must be null or point to a null -/// `*mut vx_error`. +/// and valid for as long as the data source lives. `callbacks.ctx` must be valid for the supplied +/// callbacks until `release_fn` is called. `err` must be null or point to a null `*mut vx_error`. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_new_callback_with_coalesce( session: *const vx_session, @@ -571,6 +564,131 @@ pub unsafe extern "C-unwind" fn vx_callback_sink_abort(sink: *mut vx_callback_si drop(unsafe { Box::from_raw(sink) }); } +/// Retain the compressed input chunk until its row slices have been consumed. In particular, +/// never materialize all of a dictionary's repeated values before applying the host's row limit. +struct BoundedArrayIterator { + input: I, + current: Option, + offset: usize, + batch_size: usize, + finished: bool, +} + +impl>> BoundedArrayIterator { + fn new(input: I, batch_size: usize) -> VortexResult { + vortex_ensure!(batch_size > 0, "Vortex Arrow batch size must be positive"); + Ok(Self { + input, + current: None, + offset: 0, + batch_size, + finished: false, + }) + } +} + +impl>> Iterator for BoundedArrayIterator { + type Item = VortexResult; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + loop { + if let Some(chunk) = &self.current { + // Subtract first so even a caller-provided usize::MAX row limit cannot overflow. + let length = self.batch_size.min(chunk.len() - self.offset); + let end = self.offset + length; + let sliced = chunk.slice(self.offset..end); + self.offset = end; + if end == chunk.len() || sliced.is_err() { + self.current = None; + self.offset = 0; + } + if sliced.is_err() { + self.finished = true; + } + return Some(sliced); + } + match self.input.next() { + Some(Ok(chunk)) if chunk.is_empty() => continue, + Some(Ok(chunk)) => self.current = Some(chunk), + Some(Err(error)) => { + self.finished = true; + return Some(Err(error)); + } + None => { + self.finished = true; + return None; + } + } + } + } +} + +/// Scan a partition to standard Arrow arrays, slicing compressed chunks before conversion. +/// +/// Each exported batch contains at most `batch_size` rows. Conversion is lazy: requesting the +/// next Arrow batch converts only its slice, not every slice of the current native chunk. +/// This is a row limit, not a byte limit: each batch must still fit Arrow's Utf8/Binary offsets. +/// +/// A non-null `partition` is consumed on success and failure, including a zero batch size. The +/// caller must release a successful `stream` through the Arrow C stream release callback. +/// +/// # Safety +/// +/// `partition` must be null or a valid, unconsumed partition handle. `session` must be null or +/// a valid session handle. `stream` must be null or point to writable, uninitialized/released storage. +/// On failure the stream storage is unchanged; error ownership follows the other Vortex FFI calls. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_partition_scan_arrow_with_batch_size( + session: *const vx_session, + partition: *mut vx_partition, + batch_size: usize, + stream: *mut FFI_ArrowArrayStream, + err: *mut *mut vx_error, +) -> c_int { + try_or(err, 1, || { + vortex_ensure!(!partition.is_null(), "Vortex partition is null"); + // Take ownership before validating the other arguments, so callers never need to guess + // whether an invalid request has already consumed the partition. + let partition = match *vx_partition::into_box(partition) { + VxPartitionScan::Pending(partition) => partition, + _ => vortex_bail!("Vortex partition is already being consumed"), + }; + vortex_ensure!(batch_size > 0, "Vortex Arrow batch size must be positive"); + vortex_ensure!(!session.is_null(), "Vortex session is null"); + vortex_ensure!(!stream.is_null(), "Vortex Arrow stream output is null"); + let array_stream = partition.execute()?; + let schema = Arc::new(normalize_arrow_schema_views( + array_stream.dtype().to_arrow_schema()?, + )); + let target = Field::new_struct("", schema.fields().clone(), false); + let session = vx_session::as_ref(session).clone(); + let mut chunks = + BoundedArrayIterator::new(RUNTIME.block_on_stream(array_stream), batch_size)?; + let mut failed = false; + let iter = std::iter::from_fn(move || { + if failed { + return None; + } + let chunk = chunks.next()?; + let result = (|| -> VortexResult { + let mut ctx = session.create_execution_ctx(); + let arrow = session + .arrow() + .execute_arrow(chunk?, Some(&target), &mut ctx)?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + })(); + failed = result.is_err(); + Some(result.map_err(|error| ArrowError::ExternalError(Box::new(error)))) + }); + let reader = RecordBatchIteratorAdapter::new(iter, schema); + unsafe { write(stream, FFI_ArrowArrayStream::new(Box::new(reader))) }; + Ok(0) + }) +} + /// Export standard Arrow string/binary types recursively so the declared schema matches the /// buffer layout expected by Paimon's C++ reader. pub(crate) fn normalize_arrow_schema_views(schema: arrow_schema::Schema) -> arrow_schema::Schema { @@ -580,7 +698,7 @@ pub(crate) fn normalize_arrow_schema_views(schema: arrow_schema::Schema) -> arro .cloned() .map(normalize_field_views) .collect::(); - arrow_schema::Schema::new_with_metadata(fields, schema.metadata.clone()) + arrow_schema::Schema::new_with_metadata(fields, schema.metadata) } fn normalize_field_views(field: arrow_schema::FieldRef) -> arrow_schema::FieldRef { @@ -600,10 +718,10 @@ fn normalize_dtype_views(dtype: &arrow_schema::DataType) -> arrow_schema::DataTy DataType::Struct(fields) => { DataType::Struct(fields.iter().cloned().map(normalize_field_views).collect()) } - DataType::List(field) => DataType::List(normalize_field_views(field.clone())), - DataType::LargeList(field) => DataType::LargeList(normalize_field_views(field.clone())), + DataType::List(field) => DataType::List(normalize_field_views(Arc::clone(field))), + DataType::LargeList(field) => DataType::LargeList(normalize_field_views(Arc::clone(field))), DataType::FixedSizeList(field, size) => { - DataType::FixedSizeList(normalize_field_views(field.clone()), *size) + DataType::FixedSizeList(normalize_field_views(Arc::clone(field)), *size) } other => other.clone(), } @@ -629,7 +747,8 @@ pub enum vx_time_unit { /// /// `timezone` is a null-terminated UTF-8 IANA time-zone name, or null for a local (wall-clock) /// timestamp. Vortex's own C ABI has no timestamp/date scalar constructor, so paimon-cpp adds these -/// two here to push down temporal predicates. +/// two here to push down temporal predicates. Returns null without an error when a valid value +/// exceeds the native temporal range; other failures return null and populate `err`. /// /// # Safety /// @@ -658,6 +777,19 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_timestamp( .map_err(|e| vortex_err!("timezone is not valid utf8: {e}"))?; Some(Arc::from(name)) }; + // Check timezone errors independently of the value, including values outside the native + // range. A malformed timezone must not be mistaken for an unsupported constant. + if tz.is_some() { + TemporalMetadata::Timestamp(&TimeUnit::Seconds, &tz).to_jiff(0)?; + } + // Timestamp::unpack_native uses Jiff Timestamp even without a timezone. Check both its + // Span setters and epoch range without introducing a timezone-database lookup. + let Ok(span) = time_unit.to_jiff_span(value) else { + return Ok(std::ptr::null_mut()); + }; + if jiff::Timestamp::UNIX_EPOCH.checked_add(span).is_err() { + return Ok(std::ptr::null_mut()); + } let dtype = DType::Extension( Timestamp::new_with_tz(time_unit, tz, Nullability::from(is_nullable)).erased(), ); @@ -669,6 +801,8 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_timestamp( } /// Create a date scalar literal: `days` days since the Unix epoch (1970-01-01). +/// Returns null without an error when the value exceeds the native civil-date range; other +/// failures return null and populate `err`. /// /// # Safety /// @@ -680,6 +814,14 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_date( err: *mut *mut vx_error, ) -> *mut vx_scalar { try_or(err, std::ptr::null_mut(), || { + // The fixed Days unit and epoch are valid; only a native range limit can fail here. + // Date's storage accepts all i32 values, but formatting needs a representable civil date. + if TemporalMetadata::Date(&TimeUnit::Days) + .to_jiff(i64::from(days)) + .is_err() + { + return Ok(std::ptr::null_mut()); + } let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullability::from(is_nullable)).erased()); Ok(vx_scalar::new(Scalar::try_new( @@ -695,18 +837,49 @@ mod tests { use std::fs::File; use std::os::unix::fs::FileExt; use std::ptr; - use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use arrow_array::Array; + use arrow_array::ffi_stream::ArrowArrayStreamReader; + use futures::StreamExt; + use parking_lot::Mutex; use tempfile::NamedTempFile; + use vortex::VortexSessionDefault; + use vortex::array::Canonical; + use vortex::array::ExecutionCtx; use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::ConstantArray; + use vortex::array::arrays::Dict; + use vortex::array::arrays::DictArray; + use vortex::array::arrays::ListArray; + use vortex::array::arrays::ListViewArray; use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::Struct; + use vortex::array::arrays::StructArray; + use vortex::array::arrays::VarBinViewArray; + use vortex::array::arrays::scalar_fn::ScalarFnFactoryExt; + use vortex::array::arrays::struct_::StructArrayExt; + use vortex::array::arrow::ArrowExport; + use vortex::array::arrow::ArrowExportVTable; + use vortex::array::arrow::ArrowSession; + use vortex::array::expr::stats::Precision; + use vortex::array::scalar_fn::Arity; + use vortex::array::scalar_fn::ChildName; + use vortex::array::scalar_fn::EmptyOptions; + use vortex::array::scalar_fn::ExecutionArgs; + use vortex::array::scalar_fn::ScalarFnId; + use vortex::array::scalar_fn::ScalarFnVTable; + use vortex::array::stream::SendableArrayStream; use vortex::array::validity::Validity; use vortex::buffer::buffer; use vortex::dtype::DType; use vortex::dtype::PType; + use vortex::scan::Partition; + use vortex::session::VortexSession; + use vortex::session::registry::Id; use super::*; use crate::array::vx_array; @@ -716,6 +889,8 @@ mod tests { use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_free; use crate::error::vx_error_free; + use crate::error::vx_error_get_message; + use crate::scalar::vx_scalar_free; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; use crate::session::vx_session_free; @@ -723,11 +898,931 @@ mod tests { use crate::sink::vx_array_sink_close; use crate::sink::vx_array_sink_open_file; use crate::sink::vx_array_sink_push; + use crate::string::vx_string; + + fn dictionary_view_chunk() -> ArrayRef { + let codes = PrimitiveArray::new(buffer![0u32, 1, 0, 2, 0, 0, 1], Validity::NonNullable) + .into_array(); + let text = "repeated long dictionary value ".repeat(128); + let bytes = vec![0xfdu8; 4096]; + let text_values = + VarBinViewArray::from_iter_nullable_str([Some(text.as_str()), None, Some("")]); + let byte_values = VarBinViewArray::from_iter_nullable_bin([ + Some(bytes.as_slice()), + None, + Some(b"".as_slice()), + ]); + StructArray::new( + ["text", "bytes"].into(), + vec![ + DictArray::new(codes.clone(), text_values.into_array()).into_array(), + DictArray::new(codes, byte_values.into_array()).into_array(), + ], + 7, + Validity::NonNullable, + ) + .into_array() + } + + // Test-only exporter intercepts the exact input handed to the child Arrow conversion. + // Rejecting a too-large input proves that trimming occurs before expansion, rather than + // merely checking the shape of an Arrow child that has already been materialized. + #[derive(Debug)] + struct ChildExportBound { + max_elements: usize, + lengths: Arc>>, + } + + impl ArrowExportVTable for ChildExportBound { + fn arrow_ext_id(&self) -> Id { + Id::new("paimon.test.bounded-child") + } + + fn vortex_id(&self) -> Id { + self.arrow_ext_id() + } + + fn to_arrow_field( + &self, + _name: &str, + _dtype: &DType, + _session: &ArrowSession, + ) -> VortexResult> { + Ok(None) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.lengths.lock().push(array.len()); + vortex_ensure!( + array.len() <= self.max_elements, + "child export received {} values, expected at most {}", + array.len(), + self.max_elements, + ); + Ok(ArrowExport::Unsupported(array)) + } + } + + fn bounded_child_field(binary: bool) -> Field { + Field::new( + "item", + if binary { + arrow_schema::DataType::Binary + } else { + arrow_schema::DataType::Utf8 + }, + true, + ) + .with_metadata( + [( + arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_owned(), + "paimon.test.bounded-child".to_owned(), + )] + .into(), + ) + } + + fn bounded_child_session(max_elements: usize) -> (VortexSession, Arc>>) { + let session = VortexSession::default(); + let lengths = Arc::new(Mutex::new(Vec::new())); + session + .arrow() + .register_exporter(Arc::new(ChildExportBound { + max_elements, + lengths: Arc::clone(&lengths), + })); + (session, lengths) + } + + fn nested_dictionary_elements(binary: bool) -> ArrayRef { + let codes = PrimitiveArray::new(buffer![0u32, 1, 0, 2, 1, 0, 2, 0], Validity::NonNullable) + .into_array(); + let text = "nested dictionary value ".repeat(128); + let bytes = vec![0xfdu8; 4096]; + let values = if binary { + VarBinViewArray::from_iter_nullable_bin([ + Some(bytes.as_slice()), + None, + Some(b"".as_slice()), + ]) + .into_array() + } else { + VarBinViewArray::from_iter_nullable_str([Some(text.as_str()), None, Some("")]) + .into_array() + }; + DictArray::new(codes, values).into_array() + } + + fn nested_list_fixture(binary: bool, view: bool) -> ArrayRef { + let elements = nested_dictionary_elements(binary); + let validity = Validity::from_iter([true, true, false, true, true]); + if view { + let offsets = PrimitiveArray::new(buffer![0i64, 2, 2, 4, 6], Validity::NonNullable); + let sizes = PrimitiveArray::new(buffer![2i64, 0, 2, 2, 2], Validity::NonNullable); + // SAFETY: each view ends exactly at the next offset, and the last view ends at 8. + unsafe { + ListViewArray::new_unchecked( + elements, + offsets.into_array(), + sizes.into_array(), + validity, + ) + .with_zero_copy_to_list(true) + } + .into_array() + } else { + ListArray::new( + elements, + PrimitiveArray::new(buffer![0i64, 2, 2, 4, 6, 8], Validity::NonNullable) + .into_array(), + validity, + ) + .into_array() + } + } + + fn check_nested_batches( + binary: bool, + view: bool, + target: &Field, + ) -> VortexResult<()> { + let (session, lengths) = bounded_child_session(2); + let mut ctx = session.create_execution_ctx(); + let input = nested_list_fixture(binary, view); + let mut row_count = 0; + for chunk in BoundedArrayIterator::new(std::iter::once(Ok(input)), 2)? { + let array = session + .arrow() + .execute_arrow(chunk?, Some(target), &mut ctx)?; + let list = array.as_list::(); + assert_eq!(list.value_offsets()[0].as_usize(), 0); + for row in 0..list.len() { + let source_row = row_count + row; + assert_eq!(list.is_null(row), source_row == 2); + if source_row == 1 { + assert_eq!(list.value_length(row).as_usize(), 0); + } else if source_row != 2 { + let child = list.value(row); + assert_eq!(child.len(), 2); + for element in 0..2 { + let index = [0, 0, 0, 4, 6][source_row] + element; + let code = [0, 1, 0, 2, 1, 0, 2, 0][index]; + assert_eq!(child.is_null(element), code == 1); + if code != 1 { + if binary { + let value = child.as_binary::().value(element); + assert_eq!( + value, + if code == 0 { &[0xfdu8; 4096][..] } else { &[] } + ); + } else { + let value = child.as_string::().value(element); + assert_eq!( + value, + if code == 0 { + "nested dictionary value ".repeat(128) + } else { + String::new() + }, + ); + } + } + } + } + } + row_count += list.len(); + } + assert_eq!(row_count, 5); + assert_eq!(*lengths.lock(), [2, 2, 2]); + Ok(()) + } #[test] - fn scalar_new_timestamp_and_date_build_extension_scalars() { - use crate::scalar::vx_scalar_free; + fn bounded_nested_lists_limit_child_export_before_materialization() -> VortexResult<()> { + check_null_list_spans_are_not_materialized(false)?; + for binary in [false, true] { + let child = Arc::new(bounded_child_field(binary)); + let target = Field::new("", arrow_schema::DataType::List(Arc::clone(&child)), true); + check_nested_batches::(binary, false, &target)?; + let target = Field::new("", arrow_schema::DataType::LargeList(child), true); + check_nested_batches::(binary, false, &target)?; + } + Ok(()) + } + + #[test] + fn bounded_nested_list_views_limit_child_export_before_materialization() -> VortexResult<()> { + check_non_zctl_canonical_input_bound()?; + check_null_list_spans_are_not_materialized(true)?; + for binary in [false, true] { + let child = Arc::new(bounded_child_field(binary)); + let target = Field::new("", arrow_schema::DataType::List(Arc::clone(&child)), true); + check_nested_batches::(binary, true, &target)?; + let target = Field::new("", arrow_schema::DataType::LargeList(child), true); + check_nested_batches::(binary, true, &target)?; + } + Ok(()) + } + + // Unlike an Arrow exporter, this function runs when a lazy element array is first executed + // to canonical form. ScalarFn's existing slice rule pushes each selected range through to + // the child, so the guard also detects execution hidden inside a ListView rebuild or take. + #[derive(Clone)] + struct CanonicalInputBound { + max_elements: usize, + lengths: Arc>>, + } + + impl ScalarFnVTable for CanonicalInputBound { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + Id::new("paimon.test.bounded-canonical-input") + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _index: usize) -> ChildName { + ChildName::new_arc("input".into()) + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(args[0].clone()) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.lengths.lock().push(args.row_count()); + vortex_ensure!( + args.row_count() <= self.max_elements, + "canonical execution received {} elements, expected at most {}", + args.row_count(), + self.max_elements, + ); + Ok(args.get(0)?.execute::(ctx)?.into_array()) + } + } + fn check_null_list_spans_are_not_materialized(view: bool) -> VortexResult<()> { + for validity_kind in 0..3 { + let (session, arrow_lengths) = bounded_child_session(2); + let canonical_lengths = Arc::new(Mutex::new(Vec::new())); + let mut ctx = session.create_execution_ctx(); + let elements = CanonicalInputBound { + max_elements: 2, + lengths: Arc::clone(&canonical_lengths), + } + .try_new_array( + 2048, + EmptyOptions, + vec![ + VarBinViewArray::from_iter_str((0..2048).map(|index| index.to_string())) + .into_array(), + ], + )?; + let validity = match validity_kind { + 0 => Validity::from_iter([true, false, true]), + 1 => Validity::AllInvalid, + _ => Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), + }; + let input = if view { + // SAFETY: offsets/sizes describe consecutive, non-overlapping ranges, even + // though the NULL middle slot owns most of the elements array. + unsafe { + ListViewArray::new_unchecked( + elements, + PrimitiveArray::new(buffer![0i64, 2, 2046], Validity::NonNullable) + .into_array(), + PrimitiveArray::new(buffer![2i64, 2044, 2], Validity::NonNullable) + .into_array(), + validity, + ) + .with_zero_copy_to_list(true) + } + .into_array() + } else { + ListArray::new( + elements, + PrimitiveArray::new(buffer![0i64, 2, 2046, 2048], Validity::NonNullable) + .into_array(), + validity, + ) + .into_array() + }; + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(bounded_child_field(false))), + true, + ); + let all_null = validity_kind != 0; + let mut source_row = 0; + for chunk in BoundedArrayIterator::new(std::iter::once(Ok(input)), 1)? { + let array = session + .arrow() + .execute_arrow(chunk?, Some(&target), &mut ctx)?; + let list = array.as_list::(); + let is_null = all_null || source_row == 1; + assert_eq!(list.len(), 1); + assert_eq!(list.is_null(0), is_null); + assert_eq!(list.value_length(0), if is_null { 0 } else { 2 }); + assert_eq!(list.values().len(), if is_null { 0 } else { 2 }); + if !is_null { + let values = list.values().as_string::(); + let start = if source_row == 0 { 0 } else { 2046 }; + assert_eq!(values.value(0), start.to_string()); + assert_eq!(values.value(1), (start + 1).to_string()); + } + source_row += 1; + } + assert_eq!(source_row, 3); + if all_null { + assert!(canonical_lengths.lock().is_empty()); + assert_eq!(*arrow_lengths.lock(), [0, 0, 0]); + } else { + assert_eq!(*canonical_lengths.lock(), [2, 2]); + assert_eq!(*arrow_lengths.lock(), [2, 0, 2]); + } + } + Ok(()) + } + + fn check_non_zctl_canonical_input_bound() -> VortexResult<()> { + // Two items trigger the old take path; 129 trigger the old list-by-list path. The + // shared array stays small so even the unpatched negative control uses little memory. + for size in [2usize, 129] { + let size_i32 = i32::try_from(size) + .map_err(|error| vortex_err!("Test list size out of i32 range: {error}"))?; + let (session, arrow_lengths) = bounded_child_session(2 * size); + let canonical_lengths = Arc::new(Mutex::new(Vec::new())); + let mut ctx = session.create_execution_ctx(); + let elements = CanonicalInputBound { + max_elements: size, + lengths: Arc::clone(&canonical_lengths), + } + .try_new_array( + 2048, + EmptyOptions, + vec![ + ConstantArray::new( + Scalar::utf8("a long nested value ".repeat(128), Nullability::NonNullable), + 2048, + ) + .into_array(), + ], + )?; + // Overlap, repeated ranges and a large internal gap all occur among valid rows. + // The final null row also has a nonzero size that must never be materialized. + let input = ListViewArray::new( + elements, + PrimitiveArray::new(buffer![1i64, 2, 1, 1900, 0], Validity::NonNullable) + .into_array(), + PrimitiveArray::new( + buffer![size as i64, size as i64, size as i64, size as i64, 2048], + Validity::NonNullable, + ) + .into_array(), + Validity::from_iter([true, true, true, true, false]), + ) + .into_array(); + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(bounded_child_field(false))), + true, + ); + let mut source_row = 0; + for chunk in BoundedArrayIterator::new(std::iter::once(Ok(input)), 2)? { + let array = session + .arrow() + .execute_arrow(chunk?, Some(&target), &mut ctx)?; + let list = array.as_list::(); + for row in 0..list.len() { + assert_eq!(list.is_null(row), source_row == 4); + assert_eq!( + list.value_length(row), + if source_row == 4 { 0 } else { size_i32 }, + ); + if source_row != 4 { + let values = list.value(row); + for value in values.as_string::().iter() { + assert_eq!(value, Some("a long nested value ".repeat(128).as_str())); + } + } + source_row += 1; + } + } + assert_eq!(source_row, 5); + assert_eq!(*canonical_lengths.lock(), [size, size, size, size]); + assert_eq!(*arrow_lengths.lock(), [2 * size, 2 * size, 0]); + } + Ok(()) + } + + #[test] + fn bounded_empty_nested_lists_export_no_shared_children() -> VortexResult<()> { + for view in [false, true] { + for binary in [false, true] { + let (session, lengths) = bounded_child_session(0); + let mut ctx = session.create_execution_ctx(); + let input = nested_list_fixture(binary, view).slice(2..2)?; + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(bounded_child_field(binary))), + true, + ); + let array = session + .arrow() + .execute_arrow(input, Some(&target), &mut ctx)?; + let list = array.as_list::(); + assert_eq!(list.len(), 0); + assert_eq!(list.values().len(), 0); + assert_eq!(*lengths.lock(), [0]); + } + } + Ok(()) + } + + fn check_null_nested_list_batch( + view: bool, + with_struct: bool, + batch_size: usize, + validity_kind: u8, + ) -> VortexResult<()> { + let (session, arrow_lengths) = bounded_child_session(2); + let canonical_lengths = Arc::new(Mutex::new(Vec::new())); + let mut ctx = session.create_execution_ctx(); + let values = CanonicalInputBound { + max_elements: 2, + lengths: Arc::clone(&canonical_lengths), + } + .try_new_array( + 2048, + EmptyOptions, + vec![ + VarBinViewArray::from_iter_str((0..2048).map(|index| index.to_string())) + .into_array(), + ], + )?; + let inner = ListArray::new( + values, + PrimitiveArray::new( + (0..=2048i64).collect::>(), + Validity::NonNullable, + ) + .into_array(), + Validity::NonNullable, + ) + .into_array(); + let inner_type = arrow_schema::DataType::List(Arc::new(bounded_child_field(false))); + let (inner, inner_type) = if with_struct { + ( + StructArray::new(["items"].into(), vec![inner], 2048, Validity::NonNullable) + .into_array(), + arrow_schema::DataType::Struct(vec![Field::new("items", inner_type, false)].into()), + ) + } else { + (inner, inner_type) + }; + let validity = match validity_kind { + 0 => Validity::from_iter([true, false, true]), + 1 => Validity::AllInvalid, + _ => Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), + }; + let input = if view { + // SAFETY: the three views refer to contiguous ranges, including the + // unreferenced large middle range owned by a NULL outer list. + unsafe { + ListViewArray::new_unchecked( + inner, + PrimitiveArray::new(buffer![0i64, 2, 2046], Validity::NonNullable).into_array(), + PrimitiveArray::new(buffer![2i64, 2044, 2], Validity::NonNullable).into_array(), + validity, + ) + .with_zero_copy_to_list(true) + } + .into_array() + } else { + ListArray::new( + inner, + PrimitiveArray::new(buffer![0i64, 2, 2046, 2048], Validity::NonNullable) + .into_array(), + validity, + ) + .into_array() + }; + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(Field::new("item", inner_type, false))), + true, + ); + let all_null = validity_kind != 0; + let mut source_row = 0; + for chunk in BoundedArrayIterator::new(std::iter::once(Ok(input)), batch_size)? { + let array = session + .arrow() + .execute_arrow(chunk?, Some(&target), &mut ctx)?; + let outer = array.as_list::(); + assert!(outer.len() <= batch_size); + for row in 0..outer.len() { + let is_null = all_null || source_row == 1; + assert_eq!(outer.is_null(row), is_null); + assert_eq!(outer.value_length(row), if is_null { 0 } else { 2 }); + if !is_null { + let value = outer.value(row); + let inner = if with_struct { + Arc::clone(value.as_struct().column(0)) + } else { + value + }; + let inner = inner.as_list::(); + assert_eq!(inner.len(), 2); + let start = if source_row == 0 { 0 } else { 2046 }; + for index in 0..2 { + let values = inner.value(index); + assert_eq!(values.len(), 1); + assert_eq!( + values.as_string::().value(0), + (start + index).to_string(), + ); + } + } + source_row += 1; + } + } + assert_eq!(source_row, 3); + if all_null { + assert!(canonical_lengths.lock().is_empty()); + assert!(arrow_lengths.lock().iter().all(|length| *length == 0)); + } else { + assert_eq!(canonical_lengths.lock().iter().sum::(), 4); + assert_eq!(arrow_lengths.lock().iter().sum::(), 4); + } + Ok(()) + } + + fn check_null_nested_list_batches(view: bool) -> VortexResult<()> { + for with_struct in [false, true] { + for batch_size in [1, 2] { + for validity_kind in 0..3 { + check_null_nested_list_batch(view, with_struct, batch_size, validity_kind)?; + } + } + } + Ok(()) + } + + #[test] + fn bounded_null_nested_lists_skip_unreferenced_children() -> VortexResult<()> { + check_null_nested_list_batches(false) + } + + #[test] + fn bounded_null_nested_list_views_skip_unreferenced_children() -> VortexResult<()> { + check_null_nested_list_batches(true) + } + + #[test] + fn bounded_struct_and_list_nesting_recursively_limits_child_export() -> VortexResult<()> { + for with_struct in [false, true] { + let (session, lengths) = bounded_child_session(2); + let mut ctx = session.create_execution_ctx(); + let inner = ListArray::new( + nested_dictionary_elements(false).slice(0..6)?, + PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4, 5, 6], Validity::NonNullable) + .into_array(), + Validity::NonNullable, + ) + .into_array(); + let inner_type = arrow_schema::DataType::List(Arc::new(bounded_child_field(false))); + let (inner, inner_type) = if with_struct { + ( + StructArray::new(["items"].into(), vec![inner], 6, Validity::NonNullable) + .into_array(), + arrow_schema::DataType::Struct( + vec![Field::new("items", inner_type, false)].into(), + ), + ) + } else { + (inner, inner_type) + }; + let outer = ListArray::new( + inner, + PrimitiveArray::new(buffer![0i64, 2, 4, 6], Validity::NonNullable).into_array(), + Validity::NonNullable, + ) + .into_array(); + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(Field::new("item", inner_type, false))), + false, + ); + for chunk in BoundedArrayIterator::new(std::iter::once(Ok(outer)), 1)? { + let array = session + .arrow() + .execute_arrow(chunk?, Some(&target), &mut ctx)?; + let list = array.as_list::(); + assert_eq!(list.len(), 1); + assert_eq!(list.values().len(), 2); + let inner = if with_struct { + Arc::clone(list.values().as_struct().column(0)) + } else { + Arc::clone(list.values()) + }; + assert_eq!(inner.as_list::().values().len(), 2); + } + assert_eq!(*lengths.lock(), [2, 2, 2]); + } + Ok(()) + } + + #[test] + fn bounded_nested_offsets_rebase_before_narrowing_without_large_allocation() -> VortexResult<()> + { + // ConstantArray stores one scalar even though the logical unused prefix exceeds i32. + // This exercises the overflow boundary without allocating a multi-GiB array. + let start = i32::MAX as usize + 17; + for view in [false, true] { + let (session, lengths) = bounded_child_session(2); + let mut ctx = session.create_execution_ctx(); + let elements = + ConstantArray::new(Scalar::utf8("x", Nullability::NonNullable), start + 2) + .into_array(); + let input = if view { + // SAFETY: the single view refers to exactly the final two logical elements. + unsafe { + ListViewArray::new_unchecked( + elements, + PrimitiveArray::new(buffer![start as i64], Validity::NonNullable) + .into_array(), + PrimitiveArray::new(buffer![2i64], Validity::NonNullable).into_array(), + Validity::NonNullable, + ) + .with_zero_copy_to_list(true) + } + .into_array() + } else { + ListArray::new( + elements, + PrimitiveArray::new( + buffer![start as i64, (start + 2) as i64], + Validity::NonNullable, + ) + .into_array(), + Validity::NonNullable, + ) + .into_array() + }; + let target = Field::new( + "", + arrow_schema::DataType::List(Arc::new(bounded_child_field(false))), + false, + ); + let array = session + .arrow() + .execute_arrow(input, Some(&target), &mut ctx)?; + let list = array.as_list::(); + assert_eq!(list.value_offsets(), &[0, 2]); + assert_eq!(list.values().as_string::().value(0), "x"); + assert_eq!(list.values().as_string::().value(1), "x"); + assert_eq!(*lengths.lock(), [2]); + } + Ok(()) + } + + #[derive(Default)] + struct PartitionCounts { + polls: AtomicUsize, + executions: AtomicUsize, + drops: AtomicUsize, + } + + struct ArrayPartition { + chunks: Vec, + counts: Arc, + } + + impl Drop for ArrayPartition { + fn drop(&mut self) { + self.counts.drops.fetch_add(1, Ordering::SeqCst); + } + } + + impl Partition for ArrayPartition { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn index(&self) -> usize { + 0 + } + fn row_count(&self) -> Precision { + Precision::exact( + self.chunks + .iter() + .map(|chunk| chunk.len() as u64) + .sum::(), + ) + } + fn byte_size(&self) -> Precision { + Precision::Absent + } + fn execute(self: Box) -> VortexResult { + self.counts.executions.fetch_add(1, Ordering::SeqCst); + let dtype = self.chunks[0].dtype().clone(); + let counts = Arc::clone(&self.counts); + let stream = + futures::stream::iter(self.chunks.clone().into_iter().map(Ok)).inspect(move |_| { + counts.polls.fetch_add(1, Ordering::SeqCst); + }); + Ok(Box::pin(ArrayStreamAdapter::new(dtype, stream))) + } + } + + fn array_partition(chunks: Vec, counts: Arc) -> *mut vx_partition { + vx_partition::new(VxPartitionScan::Pending(Box::new(ArrayPartition { + chunks, + counts, + }))) + } + + #[test] + fn bounded_arrays_preserve_encoding_skip_empty_and_stop_after_error() { + let chunk = dictionary_view_chunk(); + let empty = chunk.slice(0..0).unwrap(); + let polls = AtomicUsize::new(0); + let input = vec![ + Ok(empty), + Ok(chunk.clone()), + Err(vortex_err!("input failure")), + Ok(chunk.clone()), + ] + .into_iter() + .inspect(|_| { + polls.fetch_add(1, Ordering::SeqCst); + }); + let mut slices = BoundedArrayIterator::new(input, 3).unwrap(); + assert_eq!(polls.load(Ordering::SeqCst), 0); + let first = slices.next().unwrap().unwrap(); + assert_eq!(first.len(), 3); + assert!(first.as_::().unmasked_field(0).is::()); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!(slices.next().unwrap().unwrap().len(), 3); + assert_eq!(slices.next().unwrap().unwrap().len(), 1); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert!(slices.next().unwrap().is_err()); + assert!(slices.next().is_none()); + assert_eq!(polls.load(Ordering::SeqCst), 3); + + let mut whole = BoundedArrayIterator::new(std::iter::once(Ok(chunk)), usize::MAX).unwrap(); + assert_eq!(whole.next().unwrap().unwrap().len(), 7); + assert!(whole.next().is_none()); + assert!( + BoundedArrayIterator::new(std::iter::empty::>(), 0).is_err() + ); + assert!( + BoundedArrayIterator::new(std::iter::empty::>(), 3) + .unwrap() + .next() + .is_none() + ); + } + + #[test] + #[cfg_attr(miri, ignore)] + fn partition_arrow_row_limit_applies_before_native_export() { + // The original API is a negative control: C++ slicing cannot change its native batch. + for bounded in [false, true] { + let chunk = dictionary_view_chunk(); + let counts = Arc::new(PartitionCounts::default()); + let chunks = if bounded { + vec![ + chunk.slice(0..0).unwrap(), + chunk.clone(), + chunk.slice(0..0).unwrap(), + ] + } else { + vec![chunk] + }; + unsafe { + let session = vx_session_new(); + let partition = array_partition(chunks, Arc::clone(&counts)); + let mut output = FFI_ArrowArrayStream::empty(); + let mut error = ptr::null_mut(); + let status = if bounded { + vx_partition_scan_arrow_with_batch_size( + session, + partition, + 3, + &raw mut output, + &raw mut error, + ) + } else { + crate::scan::vx_partition_scan_arrow( + session, + partition, + &raw mut output, + &raw mut error, + ) + }; + assert_eq!(status, 0); + assert!(error.is_null()); + assert_eq!(counts.drops.load(Ordering::SeqCst), 1); + assert_eq!(counts.executions.load(Ordering::SeqCst), 1); + assert_eq!(counts.polls.load(Ordering::SeqCst), 0); + let mut reader = ArrowArrayStreamReader::try_new(output).unwrap(); + let mut rows = 0; + let mut lengths = Vec::new(); + let expected_codes = [0, 1, 0, 2, 0, 0, 1]; + for batch in reader.by_ref() { + let batch = batch.unwrap(); + lengths.push(batch.num_rows()); + let text = batch.column(0).as_string::(); + let bytes = batch.column(1).as_binary::(); + for row in 0..batch.num_rows() { + let code = expected_codes[rows + row]; + assert_eq!(text.is_null(row), code == 1); + assert_eq!(bytes.is_null(row), code == 1); + if code != 1 { + assert_eq!( + text.value(row), + if code == 0 { + "repeated long dictionary value ".repeat(128) + } else { + String::new() + } + ); + assert_eq!( + bytes.value(row), + if code == 0 { + vec![0xfd; 4096] + } else { + Vec::new() + } + ); + } + } + rows += batch.num_rows(); + assert_eq!( + counts.polls.load(Ordering::SeqCst), + if bounded { 2 } else { 1 } + ); + } + assert_eq!(rows, 7); + assert_eq!(lengths, if bounded { vec![3, 3, 1] } else { vec![7] }); + assert_eq!( + counts.polls.load(Ordering::SeqCst), + if bounded { 3 } else { 1 } + ); + drop(reader); + vx_session_free(session); + } + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn rejected_partition_arrow_requests_consume_partition_without_polling() { + unsafe { + let session = vx_session_new(); + for invalid in 0..3 { + let counts = Arc::new(PartitionCounts::default()); + let partition = array_partition(vec![dictionary_view_chunk()], Arc::clone(&counts)); + let mut output = FFI_ArrowArrayStream::empty(); + let mut error = ptr::null_mut(); + let status = vx_partition_scan_arrow_with_batch_size( + if invalid == 1 { ptr::null() } else { session }, + partition, + if invalid == 0 { 0 } else { 3 }, + if invalid == 2 { + ptr::null_mut() + } else { + &raw mut output + }, + &raw mut error, + ); + assert_eq!(status, 1); + assert!(!error.is_null()); + assert!(output.release.is_none()); + assert_eq!(counts.drops.load(Ordering::SeqCst), 1); + assert_eq!(counts.executions.load(Ordering::SeqCst), 0); + assert_eq!(counts.polls.load(Ordering::SeqCst), 0); + vx_error_free(error); + } + vx_session_free(session); + } + } + + #[test] + fn scalar_new_timestamp_and_date_build_extension_scalars() { unsafe { let mut error = ptr::null_mut(); let ts = vx_scalar_new_timestamp( @@ -768,12 +1863,216 @@ mod tests { } } + #[test] + fn scalar_new_timestamp_declines_span_overflow_without_unwinding() { + for (value, unit) in [ + (i64::MIN, vx_time_unit::VX_TIME_UNIT_SECOND), + (i64::MAX, vx_time_unit::VX_TIME_UNIT_SECOND), + (i64::MIN, vx_time_unit::VX_TIME_UNIT_MILLISECOND), + (i64::MAX, vx_time_unit::VX_TIME_UNIT_MILLISECOND), + (i64::MIN, vx_time_unit::VX_TIME_UNIT_MICROSECOND), + (i64::MAX, vx_time_unit::VX_TIME_UNIT_MICROSECOND), + (i64::MIN, vx_time_unit::VX_TIME_UNIT_NANOSECOND), + ] { + for timezone in [ptr::null(), c"UTC".as_ptr()] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = + vx_scalar_new_timestamp(value, unit, timezone, false, &raw mut error); + assert!(scalar.is_null()); + assert!(error.is_null()); + assert!( + vx_scalar_new_timestamp(value, unit, timezone, true, ptr::null_mut()) + .is_null() + ); + } + } + } + } + + #[test] + fn scalar_new_timestamp_declines_native_epoch_range_limits() { + for (value, unit) in [ + (253_402_207_201, vx_time_unit::VX_TIME_UNIT_SECOND), + (-377_705_023_202, vx_time_unit::VX_TIME_UNIT_SECOND), + (253_402_207_201_000, vx_time_unit::VX_TIME_UNIT_MILLISECOND), + ( + 253_402_207_201_000_000, + vx_time_unit::VX_TIME_UNIT_MICROSECOND, + ), + // Paimon permits the end of year 9999, beyond Jiff Timestamp's upper bound. + ( + 253_402_300_799_999_999, + vx_time_unit::VX_TIME_UNIT_MICROSECOND, + ), + ] { + for timezone in [ptr::null(), c"UTC".as_ptr()] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = + vx_scalar_new_timestamp(value, unit, timezone, false, &raw mut error); + assert!(scalar.is_null()); + assert!(error.is_null()); + } + } + } + } + + #[test] + fn scalar_new_timestamp_preserves_timezone_errors_outside_native_range() { + for timezone in [c"Not/A/Timezone", c"\xff"] { + for value in [0, i64::MAX] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = vx_scalar_new_timestamp( + value, + vx_time_unit::VX_TIME_UNIT_MICROSECOND, + timezone.as_ptr(), + false, + &raw mut error, + ); + assert!(scalar.is_null()); + assert!(!error.is_null()); + assert!(!vx_string::as_str(vx_error_get_message(error)).is_empty()); + vx_error_free(error); + } + } + } + } + + #[test] + fn scalar_new_timestamp_preserves_valid_storage_and_metadata() { + let cases = [ + (vx_time_unit::VX_TIME_UNIT_SECOND, TimeUnit::Seconds, -1), + (vx_time_unit::VX_TIME_UNIT_SECOND, TimeUnit::Seconds, 0), + ( + vx_time_unit::VX_TIME_UNIT_MILLISECOND, + TimeUnit::Milliseconds, + 1, + ), + ( + vx_time_unit::VX_TIME_UNIT_MICROSECOND, + TimeUnit::Microseconds, + -1, + ), + ( + vx_time_unit::VX_TIME_UNIT_MICROSECOND, + TimeUnit::Microseconds, + 253_402_207_200_999_999, + ), + ( + vx_time_unit::VX_TIME_UNIT_NANOSECOND, + TimeUnit::Nanoseconds, + i64::MIN + 1, + ), + ( + vx_time_unit::VX_TIME_UNIT_NANOSECOND, + TimeUnit::Nanoseconds, + i64::MAX, + ), + ]; + for (unit, expected_unit, value) in cases { + for timezone in [ptr::null(), c"UTC".as_ptr()] { + for nullable in [false, true] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = vx_scalar_new_timestamp( + value, + unit, + timezone, + nullable, + &raw mut error, + ); + assert!( + error.is_null(), + "unit={unit:?}, value={value}, timezone={}, nullable={nullable}: {}", + if timezone.is_null() { "none" } else { "UTC" }, + if error.is_null() { + "no native error" + } else { + vx_string::as_str(vx_error_get_message(error)) + } + ); + assert!(!scalar.is_null()); + let scalar_ref = vx_scalar::as_ref(scalar); + assert_eq!( + scalar_ref + .value() + .unwrap() + .as_primitive() + .cast::() + .unwrap(), + value + ); + assert_eq!(scalar_ref.dtype().is_nullable(), nullable); + let DType::Extension(ext) = scalar_ref.dtype() else { + panic!("timestamp scalar must carry an extension dtype"); + }; + let metadata = ext.metadata_opt::().unwrap(); + assert_eq!(metadata.unit, expected_unit); + assert_eq!( + metadata.tz.as_deref(), + if timezone.is_null() { + None + } else { + Some("UTC") + } + ); + assert!(!scalar_ref.as_extension().to_string().is_empty()); + vx_scalar_free(scalar); + } + } + } + } + } + + #[test] + fn scalar_new_date_declines_out_of_range_civil_dates() { + for days in [i32::MIN, i32::MAX] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = vx_scalar_new_date(days, false, &raw mut error); + assert!(scalar.is_null()); + assert!(error.is_null()); + assert!(vx_scalar_new_date(days, true, ptr::null_mut()).is_null()); + } + } + } + + #[test] + fn scalar_new_date_preserves_supported_dates() { + for days in [-719_528, -1, 0, 19_000, 2_932_896] { + for nullable in [false, true] { + unsafe { + let mut error = ptr::null_mut(); + let scalar = vx_scalar_new_date(days, nullable, &raw mut error); + assert!(error.is_null()); + assert!(!scalar.is_null()); + let scalar_ref = vx_scalar::as_ref(scalar); + assert_eq!( + scalar_ref + .value() + .unwrap() + .as_primitive() + .cast::() + .unwrap(), + days + ); + assert_eq!(scalar_ref.dtype().is_nullable(), nullable); + assert!(!scalar_ref.as_extension().to_string().is_empty()); + vx_scalar_free(scalar); + } + } + } + } + #[test] fn normalize_arrow_schema_views_rewrites_view_types_recursively() { + use std::sync::Arc; + use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; - use std::sync::Arc; let schema = Schema::new(vec![ Field::new("plain_utf8", DataType::Utf8, false), @@ -873,7 +2172,7 @@ mod tests { dst: *mut u8, length: usize, ) -> i32 { - // SAFETY: `ctx` is the `FileContext` handed to `vx_data_source_new_callback`, and `dst` + // SAFETY: `ctx` is the `FileContext` handed to the callback data source, and `dst` // covers `length` bytes, per the callback contract. let context = unsafe { &*ctx.cast::() }; let buffer = unsafe { std::slice::from_raw_parts_mut(dst, length) }; @@ -900,7 +2199,7 @@ mod tests { unsafe { let session = vx_session_new(); - let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype.clone())); + let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = ptr::null_mut(); let c_path = std::ffi::CString::new(path.clone()).unwrap(); @@ -916,7 +2215,10 @@ mod tests { vx_array_free(vx_array_ptr); vx_dtype_free(vx_dtype_ptr); - for local in [false, true] { + for config in [ + CoalesceConfig::object_storage(), + CoalesceConfig::new(64 << 10, 4 << 20), + ] { let file = File::open(&path).unwrap(); let size = file.metadata().unwrap().len(); let context = Box::into_raw(Box::new(FileContext { file })); @@ -926,18 +2228,14 @@ mod tests { release_fn: Some(file_release), }; - let data_source = if local { - vx_data_source_new_callback_with_coalesce( - session, - callbacks, - size, - 64 << 10, - 4 << 20, - &raw mut error, - ) - } else { - vx_data_source_new_callback(session, callbacks, size, &raw mut error) - }; + let data_source = vx_data_source_new_callback_with_coalesce( + session, + callbacks, + size, + config.distance, + config.max_size, + &raw mut error, + ); assert!(error.is_null()); assert!(!data_source.is_null()); @@ -970,7 +2268,15 @@ mod tests { // SAFETY: target owns the boxed Arc, and CallbackReadAt validates the range first. let context = unsafe { &*ctx.cast::>() }; context.reads.fetch_add(1, Ordering::SeqCst); - let bytes = &context.bytes[offset as usize..offset as usize + length]; + let Ok(start) = usize::try_from(offset) else { + return -1; + }; + let Some(end) = start.checked_add(length) else { + return -1; + }; + let Some(bytes) = context.bytes.get(start..end) else { + return -1; + }; unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, length) }; 0 } @@ -1012,10 +2318,9 @@ mod tests { .unwrap() .to_host() .await; - assert_eq!( - buffer.as_ref(), - &context.bytes[offset as usize..offset as usize + length] - ); + let start = usize::try_from(offset).unwrap(); + let end = start.checked_add(length).unwrap(); + assert_eq!(buffer.as_ref(), &context.bytes[start..end]); } assert!(reader.read_at(199_999, 2, Alignment::none()).await.is_err()); assert!( @@ -1035,11 +2340,11 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn cancelled_in_flight_read_releases_context_after_callback_returns() { - use std::sync::Condvar; use std::sync::mpsc; use std::time::Duration; use futures::channel::oneshot; + use parking_lot::Condvar; struct BlockingContext { entered: Mutex>>, @@ -1059,17 +2364,14 @@ mod tests { length: usize, ) -> i32 { let context = unsafe { &*ctx.cast::>() }; - if let Some(entered) = context.entered.lock().unwrap().take() { + if let Some(entered) = context.entered.lock().take() { let _ = entered.send(()); } - let (_proceed, timeout) = context - .wake - .wait_timeout_while( - context.proceed.lock().unwrap(), - Duration::from_secs(20), - |go| !*go, - ) - .unwrap(); + let mut proceed = context.proceed.lock(); + let timeout = + context + .wake + .wait_while_for(&mut proceed, |go| !*go, Duration::from_secs(20)); if !timeout.timed_out() && context.status == 0 { unsafe { ptr::write_bytes(dst, 0x5a, length) }; } @@ -1122,7 +2424,7 @@ mod tests { drop(reader); assert_eq!(context.releases.load(Ordering::SeqCst), 0); assert!(!context.completed.load(Ordering::SeqCst)); - *context.proceed.lock().unwrap() = true; + *context.proceed.lock() = true; context.wake.notify_all(); RUNTIME .handle() @@ -1231,7 +2533,15 @@ mod tests { release_fn: Some(mark_released), }; - let data_source = vx_data_source_new_callback(session, callbacks, 0, &raw mut error); + let config = CoalesceConfig::object_storage(); + let data_source = vx_data_source_new_callback_with_coalesce( + session, + callbacks, + 0, + config.distance, + config.max_size, + &raw mut error, + ); assert!(data_source.is_null()); assert!(!error.is_null()); vx_error_free(error); @@ -1259,7 +2569,7 @@ mod tests { // `length` bytes, per the callback contract. let holder = unsafe { &*ctx.cast::>() }; let bytes = unsafe { std::slice::from_raw_parts(src, length) }; - holder.data.lock().unwrap().extend_from_slice(bytes); + holder.data.lock().extend_from_slice(bytes); 0 } @@ -1271,8 +2581,10 @@ mod tests { ) -> i32 { // SAFETY: as above; `dst` covers `length` writable bytes. let holder = unsafe { &*ctx.cast::>() }; - let data = holder.data.lock().unwrap(); - let start = offset as usize; + let data = holder.data.lock(); + let Ok(start) = usize::try_from(offset) else { + return -1; + }; let Some(end) = start.checked_add(length) else { return -1; }; @@ -1296,7 +2608,7 @@ mod tests { unsafe { let session = vx_session_new(); - let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype.clone())); + let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = ptr::null_mut(); let out_ctx = Box::into_raw(Box::new(Arc::clone(&shared))); @@ -1323,11 +2635,12 @@ mod tests { assert!(error.is_null()); vx_array_free(vx_array_ptr); - let size = shared.data.lock().unwrap().len() as u64; + let size = shared.data.lock().len() as u64; assert!(size > 0, "the sink must have written through the callback"); let in_ctx = Box::into_raw(Box::new(Arc::clone(&shared))); - let data_source = vx_data_source_new_callback( + let config = CoalesceConfig::object_storage(); + let data_source = vx_data_source_new_callback_with_coalesce( session, vx_input_callbacks { ctx: in_ctx.cast::(), @@ -1335,6 +2648,8 @@ mod tests { release_fn: Some(buffer_release), }, size, + config.distance, + config.max_size, &raw mut error, ); assert!(error.is_null()); diff --git a/docs/source/api/file_format.rst b/docs/source/api/file_format.rst index f90e07f22..9ae2ca7d4 100644 --- a/docs/source/api/file_format.rst +++ b/docs/source/api/file_format.rst @@ -55,3 +55,59 @@ Interface .. doxygenclass:: paimon::WrittenFileStatsProvider :members: :undoc-members: + + +.. _cpp-api-vortex-read-options: + +Vortex Read Options +=================== + +Pass Vortex read options through ``ReadContextBuilder::AddOption`` at runtime. +Read options override table options for each key; a key absent from both uses +its default. Changes take effect when a new reader is created and do not change +existing readers. + +Predicate Pushdown +------------------ + +``vortex.read.in-pushdown.max-values`` sets the maximum list length eligible for +native IN/NOT IN pushdown. It accepts a decimal integer from ``0`` to +``4294967295`` (``UINT32_MAX``), with a default of ``32``. Negative, malformed or +out-of-range values are rejected. The limit applies independently to each IN or +NOT IN predicate and counts the original list entries, including duplicates. + +A positive value permits lists of up to that length, subject to type and value +safety checks. ``0`` disables only IN/NOT IN pushdown. A list exceeding the limit +is not pushed down and is never truncated. Other eligible predicates may still +be pushed down; an OR expression requires every branch to be convertible. + +The Vortex adapter does not deduplicate lists or compress consecutive integer +values into ranges. The query engine is responsible for these optimizations. +To allow up to 128 entries per IN/NOT IN predicate: + +.. code-block:: cpp + + read_context_builder.AddOption("vortex.read.in-pushdown.max-values", "128"); + +Native pushdown may return candidate rows requiring further filtering. For exact +results, retain the full residual predicate: enable ``EnablePredicateFilter(true)`` +with ``TableRead``, or apply equivalent filtering in the caller. This also applies +when an IN/NOT IN predicate exceeds the configured limit. + +I/O Coalescing +-------------- + +``vortex.read.io-coalesce.max-gap`` and ``vortex.read.io-coalesce.max-size`` +accept decimal byte counts without unit suffixes, up to +``9223372036854775807`` (``INT64_MAX``). Their defaults are ``1048576`` +(1 MiB) and ``16777216`` (16 MiB), respectively. ``max-gap`` must be non-negative, +``max-size`` must be positive, and ``max-gap`` must not exceed ``max-size``. + +``max-size`` is a threshold for merging requested ranges. An individual original +request may exceed it, so it does not impose a hard read-size or memory limit. +For example, configure a 64 KiB gap and a 4 MiB merging threshold: + +.. code-block:: cpp + + read_context_builder.AddOption("vortex.read.io-coalesce.max-gap", "65536"); + read_context_builder.AddOption("vortex.read.io-coalesce.max-size", "4194304"); diff --git a/docs/source/user_guide/read.rst b/docs/source/user_guide/read.rst index ff1a9aa00..a94776353 100644 --- a/docs/source/user_guide/read.rst +++ b/docs/source/user_guide/read.rst @@ -41,26 +41,6 @@ cross-language task scheduling and interaction (e.g., Java and C++), substantial across the two language ecosystems. -Vortex Read Options -------------------- - -For Vortex files, pass these options through ``ReadContextBuilder::AddOption``: - -- ``vortex.read.membership-pushdown`` accepts ``auto`` (default), ``force`` or ``off``. - ``auto`` limits each IN/NOT IN expression to 64 comparison and boolean nodes after - deduplication and integer-range compression. ``force`` bypasses this budget while - retaining type and value safety checks. ``off`` disables only IN/NOT IN pushdown. -- ``vortex.read.io-coalesce`` accepts ``object-storage`` (default) or ``local``. - The default retains Vortex's object-storage coalescing policy. ``local`` uses a - 64 KiB maximum gap and a 4 MiB coalescing limit for low-latency storage. - Choose the mode explicitly; it is not selected from the file path. - -Native pushdown may return candidate rows requiring further filtering. For exact -results, retain the full residual predicate: enable ``EnablePredicateFilter(true)`` -with ``TableRead``, or apply equivalent filtering in the caller. ``force`` does not -bypass this requirement. - - Schema Evolution ----------------------- Scope and Compatibility diff --git a/src/paimon/format/vortex/vortex_defs.h b/src/paimon/format/vortex/vortex_defs.h new file mode 100644 index 000000000..04f5c104e --- /dev/null +++ b/src/paimon/format/vortex/vortex_defs.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 + +namespace paimon::vortex { + +inline constexpr char kVortexReadInPushdownMaxValues[] = "vortex.read.in-pushdown.max-values"; +inline constexpr uint32_t kDefaultVortexReadInPushdownMaxValues = 32; + +/// Limits applied independently to each IN or NOT IN predicate. +struct VortexPredicateOptions { + /// Maximum number of input literals, including duplicates; zero disables membership pushdown. + uint32_t max_in_values = kDefaultVortexReadInPushdownMaxValues; +}; + +inline constexpr char kVortexReadIoCoalesceMaxGap[] = "vortex.read.io-coalesce.max-gap"; +inline constexpr char kVortexReadIoCoalesceMaxSize[] = "vortex.read.io-coalesce.max-size"; + +// Preserve Vortex's object-storage defaults when no explicit values are supplied. +inline constexpr uint64_t kDefaultVortexReadIoCoalesceMaxGap = 1ULL << 20; +inline constexpr uint64_t kDefaultVortexReadIoCoalesceMaxSize = 16ULL << 20; + +/// I/O range coalescing limits, measured in bytes. +struct VortexIoCoalesceOptions { + /// Maximum gap between ranges considered for merging. + uint64_t max_gap = kDefaultVortexReadIoCoalesceMaxGap; + /// Maximum merged span; a single original request may exceed this threshold. + uint64_t max_size = kDefaultVortexReadIoCoalesceMaxSize; +}; + +/// Names for bitmap-selection counters set after a successful read-schema update. +class VortexMetrics { + public: + VortexMetrics() = delete; + ~VortexMetrics() = delete; + + /// Rows selected by the bitmap, before predicate filtering; no bitmap selects all rows. + static inline const char kSelectionRowCount[] = "vortex.selection-row-count"; + /// Stored selected or excluded row ids, using whichever side is smaller. + static inline const char kSelectionBufferRowCount[] = "vortex.selection-buffer-row-count"; +}; + +} // namespace paimon::vortex diff --git a/src/paimon/format/vortex/vortex_ffi.h b/src/paimon/format/vortex/vortex_ffi.h index f14e71071..1d1943e78 100644 --- a/src/paimon/format/vortex/vortex_ffi.h +++ b/src/paimon/format/vortex/vortex_ffi.h @@ -75,19 +75,22 @@ struct vx_input_callbacks { vx_release_fn release_fn; }; -/// Create a data source that reads through `callbacks`. `size` is the total file length in bytes. -/// Returns null and sets `err` on failure. -const vx_data_source* vx_data_source_new_callback(const vx_session* session, - vx_input_callbacks callbacks, uint64_t size, - vx_error** err); - -/// Create a callback data source with explicit I/O coalescing. max_size must be positive and -/// max_gap must not exceed max_size. The callback context is released exactly once, even on error. +/// Create a callback data source with explicit I/O coalescing. `size` is the total file length +/// in bytes. max_size must be positive and max_gap must not exceed max_size. Returns null and +/// sets `err` on failure. The callback context is released exactly once, even on error. const vx_data_source* vx_data_source_new_callback_with_coalesce(const vx_session* session, vx_input_callbacks callbacks, uint64_t size, uint64_t max_gap, uint64_t max_size, vx_error** err); +/// Export standard Arrow batches with at most `batch_size` rows, slicing encoded arrays before +/// materializing string/binary data. `batch_size` must be positive. Consumes `partition` on both +/// success and error; do not free or reuse it afterwards. Returns 0 on success, or 1 and sets +/// `err`. +int vx_partition_scan_arrow_with_batch_size(const vx_session* session, vx_partition* partition, + size_t batch_size, FFI_ArrowArrayStream* stream, + vx_error** err); + /// Append `length` bytes from `src` to the host sink. Returns 0 on success and non-zero on failure; /// a partial write must be reported as a failure. Writes for one sink are sequential, never /// concurrent. @@ -134,11 +137,13 @@ enum vx_time_unit { }; /// Create a timestamp literal in `unit` since the Unix epoch. `timezone` is an IANA zone name, -/// or null for wall-clock time. Returns null and sets `err` on failure; free with vx_scalar_free. +/// or null for wall-clock time. Returns null without an error when a valid value exceeds the +/// native temporal range; other failures return null and set `err`. Free with vx_scalar_free. vx_scalar* vx_scalar_new_timestamp(int64_t value, vx_time_unit unit, const char* timezone, bool is_nullable, vx_error** err); -/// Create a date scalar literal: `days` days since the Unix epoch (1970-01-01). Returns null and -/// sets `err` on failure; free the result with `vx_scalar_free`. +/// Create a date scalar literal: `days` days since the Unix epoch (1970-01-01). Returns null +/// without an error outside the native civil-date range; other failures return null and set `err`. +/// Free the result with `vx_scalar_free`. vx_scalar* vx_scalar_new_date(int32_t days, bool is_nullable, vx_error** err); } diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.cpp b/src/paimon/format/vortex/vortex_file_batch_reader.cpp index fa7aa6324..a8b284398 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.cpp +++ b/src/paimon/format/vortex/vortex_file_batch_reader.cpp @@ -36,6 +36,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/math.h" +#include "paimon/core/schema/arrow_schema_validator.h" #include "paimon/core/utils/nested_projection_utils.h" #include "paimon/format/vortex/vortex_ffi_util.h" #include "paimon/format/vortex/vortex_predicate_converter.h" @@ -49,6 +50,41 @@ namespace { // Must match the physical row index column appended by vortex-ffi scan.rs. constexpr char kRowIdxColumn[] = "_row_idx"; +// Validate field references without changing read-side types or metadata. Type conversion +// remains the responsibility of AlignToReadSchema. +Status ValidateReadFieldReferences(const std::shared_ptr& file_type, + const std::shared_ptr& read_type) { + if (file_type->id() != read_type->id()) { + return Status::OK(); + } + switch (read_type->id()) { + case arrow::Type::STRUCT: { + PAIMON_ASSIGN_OR_RAISE(std::vector projection, + ArrowUtils::CreateProjection(arrow::schema(file_type->fields()), + read_type->fields())); + for (int32_t i = 0; i < read_type->num_fields(); ++i) { + PAIMON_RETURN_NOT_OK(ValidateReadFieldReferences( + file_type->field(projection[i])->type(), read_type->field(i)->type())); + } + return Status::OK(); + } + case arrow::Type::LIST: + case arrow::Type::FIXED_SIZE_LIST: + // Container field names (for example item/element) need not match. + return ValidateReadFieldReferences(file_type->field(0)->type(), + read_type->field(0)->type()); + case arrow::Type::MAP: { + const auto& file_map = checked_cast(*file_type); + const auto& read_map = checked_cast(*read_type); + PAIMON_RETURN_NOT_OK( + ValidateReadFieldReferences(file_map.key_type(), read_map.key_type())); + return ValidateReadFieldReferences(file_map.item_type(), read_map.item_type()); + } + default: + return Status::OK(); + } +} + // Vortex reports the file's zone and paimon asks for the reader's; the instants are the same. std::shared_ptr AlignTimestampTimezones( const std::shared_ptr& data, @@ -113,12 +149,136 @@ std::shared_ptr AlignTimestampTimezones( } // namespace +Result VortexFileBatchReader::SelectionState::Create( + const std::optional& bitmap, uint64_t total_rows) { + SelectionState state; + state.has_selection_ = bitmap.has_value(); + state.selected_row_count_ = total_rows; + if (bitmap) { + const int32_t cardinality = bitmap->Cardinality(); + if (cardinality < 0 || static_cast(cardinality) > total_rows) { + return Status::Invalid("Vortex selection cardinality exceeds the file row count"); + } + state.selected_row_count_ = static_cast(cardinality); + state.excludes_ = state.selected_row_count_ > total_rows / 2; + state.indices_.reserve(static_cast( + state.excludes_ ? total_rows - state.selected_row_count_ : state.selected_row_count_)); + uint64_t next_row = 0; + const auto end = bitmap->End(); + for (auto iter = bitmap->Begin(); iter != end; ++iter) { + const auto row_id = static_cast(*iter); + if (row_id >= total_rows) { + return Status::Invalid( + fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows)); + } + if (state.excludes_) { + // Enumerate gaps directly instead of first expanding all selected rows. + for (; next_row < row_id; ++next_row) { + state.indices_.push_back(next_row); + } + next_row = row_id + 1; + } else { + state.indices_.push_back(row_id); + } + } + if (state.excludes_) { + for (; next_row < total_rows; ++next_row) { + state.indices_.push_back(next_row); + } + } + } + return state; +} + +void VortexFileBatchReader::SelectionState::ApplyToScan(vx_scan_options* options) const { + // A dense bitmap sends excluded rows; all-selected uses the default INCLUDE_ALL mode. + if (has_selection_ && !indices_.empty()) { + options->selection.idx = indices_.data(); + options->selection.idx_len = indices_.size(); + options->selection.include = + excludes_ ? VX_SELECTION_EXCLUDE_RANGE : VX_SELECTION_INCLUDE_RANGE; + } +} + +Result VortexFileBatchReader::SelectionState::GetPhysicalRowId( + uint64_t selected_rank) const { + // Ordered scans emit selected rows in physical order, regardless of which side is stored. + if (has_selection_) { + if (selected_rank >= selected_row_count_) { + return Status::Invalid(fmt::format("emitted row offset {} exceeds selection size {}", + selected_rank, selected_row_count_)); + } + if (!excludes_) { + return indices_[selected_rank]; + } + // For excluded row i, excluded[i] - i is the number of selected rows before it. + // Count exclusions at or before this selected rank without expanding a row mapping. + size_t begin = 0; + size_t end = indices_.size(); + while (begin < end) { + const size_t middle = begin + (end - begin) / 2; + if (indices_[middle] - middle <= selected_rank) { + begin = middle + 1; + } else { + end = middle; + } + } + return selected_rank + begin; + } + return selected_rank; +} + +Status VortexFileBatchReader::RowPositionState::CaptureRowIndices( + const std::shared_ptr& array) { + if (has_filter_) { + // Capture physical row indices before AlignToReadSchema removes the hidden column. + const auto& struct_type = checked_cast(*array->type()); + const int32_t row_idx_pos = struct_type.num_fields() - 1; + if (row_idx_pos < 0 || struct_type.field(row_idx_pos)->name() != kRowIdxColumn) { + return Status::Invalid("Vortex scan did not append the _row_idx column"); + } + previous_batch_row_idx_ = + checked_cast(*array).field(row_idx_pos); + } + return Status::OK(); +} + +void VortexFileBatchReader::RowPositionState::RecordBatch(uint64_t row_count) { + previous_first_row_ = rows_emitted_; + previous_batch_row_count_ = row_count; + rows_emitted_ += row_count; +} + +Result VortexFileBatchReader::RowPositionState::GetPhysicalRowId( + uint64_t batch_row_id, const SelectionState& selection) const { + if (previous_batch_row_count_ == 0) { + return Status::Invalid(previous_first_row_ == std::numeric_limits::max() + ? "no Vortex batch has been read yet" + : "last Vortex 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_)); + } + // The hidden column accounts for both filtering and selection; emitted offsets do not. + if (has_filter_) { + if (previous_batch_row_idx_ == nullptr || + batch_row_id >= static_cast(previous_batch_row_idx_->length())) { + return Status::Invalid( + fmt::format("Vortex filter row index unavailable for row {}", batch_row_id)); + } + return checked_cast(*previous_batch_row_idx_) + .Value(static_cast(batch_row_id)); + } + return selection.GetPhysicalRowId(previous_first_row_ + batch_row_id); +} + VortexFileBatchReader::VortexFileBatchReader( const std::shared_ptr& input, int32_t batch_size, std::shared_ptr input_context, VxSessionPtr session, VxDataSourcePtr data_source, VxScanPtr scan, const std::shared_ptr& file_schema, - uint64_t total_rows, const std::shared_ptr& arrow_pool, - const VortexPredicateOptions& predicate_options) + uint64_t total_rows, const VortexPredicateOptions& predicate_options, + const std::shared_ptr& arrow_pool) : input_(input), batch_size_(batch_size), input_context_(std::move(input_context)), @@ -133,8 +293,9 @@ VortexFileBatchReader::VortexFileBatchReader( Result> VortexFileBatchReader::Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& arrow_pool, - const VortexPredicateOptions& predicate_options, VortexIoCoalesceMode coalesce_mode) { + const VortexPredicateOptions& predicate_options, + const VortexIoCoalesceOptions& coalesce_options, + const std::shared_ptr& arrow_pool) { if (input == nullptr || arrow_pool == nullptr || batch_size <= 0) { return Status::Invalid( "Vortex reader requires non-null input and Arrow memory pool, and positive batch size"); @@ -149,17 +310,10 @@ Result> VortexFileBatchReader::Create( return Status::IOError("failed to create Vortex session"); } vx_error* error = nullptr; - const bool local_coalesce = coalesce_mode == VortexIoCoalesceMode::LOCAL; - // Smaller gaps limit read amplification for explicitly selected low-latency sources. - constexpr uint64_t kLocalMaxGap = 64ULL << 10; - constexpr uint64_t kLocalMaxSize = 4ULL << 20; VxDataSourcePtr data_source( - local_coalesce - ? vx_data_source_new_callback_with_coalesce( - session.get(), VortexInputContext::MakeCallbacks(input_context), length, - kLocalMaxGap, kLocalMaxSize, &error) - : vx_data_source_new_callback( - session.get(), VortexInputContext::MakeCallbacks(input_context), length, &error), + vx_data_source_new_callback_with_coalesce( + session.get(), VortexInputContext::MakeCallbacks(input_context), length, + coalesce_options.max_gap, coalesce_options.max_size, &error), vx_data_source_free); if (data_source == nullptr) { return VortexCallbackError("open Vortex callback data source", error, @@ -193,84 +347,67 @@ Result> VortexFileBatchReader::Create( // scan may be consumed only once. return std::unique_ptr(new VortexFileBatchReader( input, batch_size, std::move(input_context), std::move(session), std::move(data_source), - VxScanPtr(nullptr, vx_scan_free), file_schema, total_rows, arrow_pool, predicate_options)); + VxScanPtr(nullptr, vx_scan_free), file_schema, total_rows, predicate_options, arrow_pool)); } VortexFileBatchReader::~VortexFileBatchReader() { CloseInternal(); } +Status VortexFileBatchReader::CreateScan() { + vx_error* error = nullptr; + // Scan in storage order: Vortex defaults to ordered=false and may emit chunks out of order, + // which would misalign the physical row positions NextBatch assigns. + vx_scan_options options{}; + options.ordered = true; + // The scan clones the projection, so these expressions need only outlive scan creation. + VxExpressionPtr root(nullptr, vx_expression_free); + VxExpressionPtr projection(nullptr, vx_expression_free); + std::vector name_ptrs; + if (read_schema_ != nullptr && read_schema_->num_fields() > 0) { + name_ptrs.reserve(read_schema_->num_fields()); + // The schema owns these names; Vortex copies them when creating the projection. + for (const std::shared_ptr& field : read_schema_->fields()) { + name_ptrs.push_back(field->name().c_str()); + } + root = VxExpressionPtr(vx_expression_root(), vx_expression_free); + projection = + VxExpressionPtr(vx_expression_select(name_ptrs.data(), name_ptrs.size(), root.get()), + vx_expression_free); + options.projection = projection.get(); + } + selection_.ApplyToScan(&options); + // Filtering drops rows; include_row_idx preserves their original physical positions. + // The scan clones the filter, so it can be freed after scan creation. + VxExpressionPtr filter(nullptr, vx_expression_free); + // SetReadSchema validates requested fields against the file schema, so checking the + // file alone also protects the hidden index from collisions with requested fields. + const bool has_row_idx_column = !file_schema_->GetAllFieldIndices(kRowIdxColumn).empty(); + if (!has_row_idx_column) { + PAIMON_ASSIGN_OR_RAISE(filter, VortexPredicateConverter::Convert( + predicate_filter_, file_schema_, predicate_options_)); + } + if (filter != nullptr) { + options.filter = filter.get(); + options.include_row_idx = true; + row_position_.EnableFilterRowIds(); + } + VxScanPtr scan(vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), + vx_scan_free); + if (scan == nullptr) { + return VortexCallbackError("create Vortex scan", error, + input_context_->GetCallbackStatus()); + } + scan_ = std::move(scan); + return Status::OK(); +} + Result VortexFileBatchReader::OpenNextPartitionStream() { - if (has_selection_ && selected_row_count_ == 0) { + if (selection_.IsEmpty()) { return false; } if (scan_ == nullptr) { - vx_error* error = nullptr; - // Scan in storage order: Vortex defaults to ordered=false and may emit chunks out of order, - // which would misalign the physical row positions NextBatch assigns. - vx_scan_options options{}; - options.ordered = true; - // Missing fields are null-filled by AlignToReadSchema and cannot be selected here. - // The scan clones the projection, so these expressions need only outlive scan creation. - VxExpressionPtr root(nullptr, vx_expression_free); - VxExpressionPtr projection(nullptr, vx_expression_free); - std::vector names; - std::vector name_ptrs; - if (read_schema_ != nullptr) { - names.reserve(read_schema_->num_fields()); - // Native merge requires distinct fields; output alignment restores repeated columns. - std::vector selected_columns(file_schema_->num_fields(), 0); - for (const std::shared_ptr& field : read_schema_->fields()) { - const int32_t field_index = file_schema_->GetFieldIndex(field->name()); - if (field_index >= 0 && selected_columns[field_index] == 0) { - names.push_back(field->name()); - selected_columns[field_index] = 1; - } - } - if (!names.empty()) { - name_ptrs.reserve(names.size()); - for (const std::string& name : names) { - name_ptrs.push_back(name.c_str()); - } - root = VxExpressionPtr(vx_expression_root(), vx_expression_free); - projection = VxExpressionPtr( - vx_expression_select(name_ptrs.data(), name_ptrs.size(), root.get()), - vx_expression_free); - options.projection = projection.get(); - } - } - // These FFI modes take individual indices. A dense bitmap sends only excluded rows; - // an all-selected bitmap needs no indices and uses the default INCLUDE_ALL mode. - if (has_selection_ && !selection_indices_.empty()) { - options.selection.idx = selection_indices_.data(); - options.selection.idx_len = selection_indices_.size(); - options.selection.include = - selection_excludes_ ? VX_SELECTION_EXCLUDE_RANGE : VX_SELECTION_INCLUDE_RANGE; - } - // Filtering drops rows; include_row_idx preserves their original physical positions. - // The scan clones the filter, so it can be freed after scan creation. - VxExpressionPtr filter(nullptr, vx_expression_free); - // The hidden index must not collide with a file field or a requested missing field. - const bool has_row_idx_column = - file_schema_->GetFieldIndex(kRowIdxColumn) >= 0 || - (read_schema_ != nullptr && !read_schema_->GetAllFieldIndices(kRowIdxColumn).empty()); - if (!has_row_idx_column) { - filter = VortexPredicateConverter::Convert(predicate_filter_, file_schema_, - predicate_options_); - } - if (filter != nullptr) { - options.filter = filter.get(); - options.include_row_idx = true; - has_filter_ = true; - } - VxScanPtr scan( - vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), - vx_scan_free); - if (scan == nullptr) { - return VortexCallbackError("create Vortex scan", error, - input_context_->GetCallbackStatus()); - } - scan_ = std::move(scan); + PAIMON_RETURN_NOT_OK(CreateScan()); } vx_error* error = nullptr; VxPartitionPtr partition(vx_scan_next_partition(scan_.get(), &error), vx_partition_free); @@ -281,10 +418,12 @@ Result VortexFileBatchReader::OpenNextPartitionStream() { if (partition == nullptr) { return false; // Scan exhausted. } - // vx_partition_scan_arrow consumes the partition (even on error), so release ownership here. + // Arrow stream export consumes the partition (even on error), so release ownership here. ::ArrowArrayStream stream{}; error = nullptr; - if (vx_partition_scan_arrow(session_.get(), partition.release(), &stream, &error) != 0) { + if (vx_partition_scan_arrow_with_batch_size(session_.get(), partition.release(), + static_cast(batch_size_), &stream, + &error) != 0) { return VortexCallbackError("scan Vortex partition to Arrow", error, input_context_->GetCallbackStatus()); } @@ -364,7 +503,7 @@ Result VortexFileBatchReader::NextBatch() { current_batch_offset_ = 0; } if (current_batch_ == nullptr) { - previous_batch_row_count_ = 0; + row_position_.SetEof(); return BatchReader::MakeEofBatch(); } @@ -374,22 +513,11 @@ Result VortexFileBatchReader::NextBatch() { current_batch_->Slice(current_batch_offset_, row_count); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_array, ArrowUtils::NormalizeArrayOffsets(sliced_array, arrow_pool_.get())); - if (has_filter_) { - // Capture physical row indices before AlignToReadSchema removes the hidden column. - const auto& struct_type = checked_cast(*normalized_array->type()); - const int32_t row_idx_pos = struct_type.num_fields() - 1; - if (row_idx_pos < 0 || struct_type.field(row_idx_pos)->name() != kRowIdxColumn) { - return Status::Invalid("Vortex scan did not append the _row_idx column"); - } - previous_batch_row_idx_ = - checked_cast(*normalized_array).field(row_idx_pos); - } + PAIMON_RETURN_NOT_OK(row_position_.CaptureRowIndices(normalized_array)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr projected_array, AlignToReadSchema(normalized_array)); - previous_first_row_ = rows_emitted_; - previous_batch_row_count_ = static_cast(row_count); - rows_emitted_ += static_cast(row_count); + row_position_.RecordBatch(static_cast(row_count)); current_batch_offset_ += row_count; auto ffi_array = std::make_unique<::ArrowArray>(); @@ -415,110 +543,31 @@ Status VortexFileBatchReader::SetReadSchema( if (read_schema == nullptr) { return Status::Invalid("Vortex read schema is nullptr"); } - std::vector selection_indices; - bool selection_excludes = false; - uint64_t selected_row_count = total_rows_; - if (selection_bitmap) { - const int32_t cardinality = selection_bitmap->Cardinality(); - if (cardinality < 0 || static_cast(cardinality) > total_rows_) { - return Status::Invalid("Vortex selection cardinality exceeds the file row count"); - } - selected_row_count = static_cast(cardinality); - selection_excludes = selected_row_count > total_rows_ / 2; - selection_indices.reserve(static_cast( - selection_excludes ? total_rows_ - selected_row_count : selected_row_count)); - uint64_t next_row = 0; - const auto end = selection_bitmap->End(); - for (auto iter = selection_bitmap->Begin(); iter != end; ++iter) { - const auto row_id = static_cast(*iter); - if (row_id >= total_rows_) { - return Status::Invalid( - fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows_)); - } - if (selection_excludes) { - // Enumerate gaps directly instead of first expanding all selected rows. - for (; next_row < row_id; ++next_row) { - selection_indices.push_back(next_row); - } - next_row = row_id + 1; - } else { - selection_indices.push_back(row_id); - } - } - if (selection_excludes) { - for (; next_row < total_rows_; ++next_row) { - selection_indices.push_back(next_row); - } - } - } + PAIMON_ASSIGN_OR_RAISE(SelectionState selection, + SelectionState::Create(selection_bitmap, total_rows_)); // Leave the active scan and its row mapping unchanged if validation or import fails. PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_read_schema, arrow::ImportSchema(read_schema)); + PAIMON_RETURN_NOT_OK( + ArrowSchemaValidator::ValidateNoRedundantFields(imported_read_schema->fields())); + PAIMON_RETURN_NOT_OK(ValidateReadFieldReferences( + arrow::struct_(file_schema_->fields()), arrow::struct_(imported_read_schema->fields()))); // Recreate the single-use scan on the next read with the new projection and predicates. read_schema_ = std::move(imported_read_schema); predicate_filter_ = std::dynamic_pointer_cast(predicate); - has_selection_ = selection_bitmap.has_value(); - selection_excludes_ = selection_excludes; - selected_row_count_ = selected_row_count; - selection_indices_ = std::move(selection_indices); - metrics_->SetCounter("vortex.selection-index-count", selection_indices_.size()); + selection_ = std::move(selection); + metrics_->SetCounter(VortexMetrics::kSelectionRowCount, selection_.SelectedRowCount()); + metrics_->SetCounter(VortexMetrics::kSelectionBufferRowCount, selection_.IndexCount()); ReleaseStream(); current_batch_.reset(); current_batch_offset_ = 0; scan_.reset(); - rows_emitted_ = 0; - previous_first_row_ = std::numeric_limits::max(); - previous_batch_row_count_ = 0; - has_filter_ = false; - previous_batch_row_idx_.reset(); + row_position_.Reset(); return Status::OK(); } Result VortexFileBatchReader::GetPreviousBatchFileRowId(uint64_t batch_row_id) const { - if (previous_batch_row_count_ == 0) { - return Status::Invalid(previous_first_row_ == std::numeric_limits::max() - ? "no Vortex batch has been read yet" - : "last Vortex 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_)); - } - // The hidden column accounts for both filtering and selection; emitted offsets do not. - if (has_filter_) { - if (previous_batch_row_idx_ == nullptr || - batch_row_id >= static_cast(previous_batch_row_idx_->length())) { - return Status::Invalid( - fmt::format("Vortex filter row index unavailable for row {}", batch_row_id)); - } - return checked_cast(*previous_batch_row_idx_) - .Value(static_cast(batch_row_id)); - } - const uint64_t emitted_offset = previous_first_row_ + batch_row_id; - // Ordered scans emit selected rows in physical order, regardless of which side is stored. - if (has_selection_) { - if (emitted_offset >= selected_row_count_) { - return Status::Invalid(fmt::format("emitted row offset {} exceeds selection size {}", - emitted_offset, selected_row_count_)); - } - if (!selection_excludes_) { - return selection_indices_[emitted_offset]; - } - // For excluded row i, excluded[i] - i is the number of selected rows before it. - // Count exclusions at or before this selected rank without expanding a row mapping. - size_t begin = 0; - size_t end = selection_indices_.size(); - while (begin < end) { - const size_t middle = begin + (end - begin) / 2; - if (selection_indices_[middle] - middle <= emitted_offset) { - begin = middle + 1; - } else { - end = middle; - } - } - return emitted_offset + begin; - } - return emitted_offset; + return row_position_.GetPhysicalRowId(batch_row_id, selection_); } Result VortexFileBatchReader::GetNumberOfRows() const { diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.h b/src/paimon/format/vortex/vortex_file_batch_reader.h index 5175b10ad..8a8b122cb 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.h +++ b/src/paimon/format/vortex/vortex_file_batch_reader.h @@ -25,6 +25,7 @@ #include #include "arrow/type_fwd.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_ffi.h" #include "paimon/format/vortex/vortex_ffi_util.h" #include "paimon/format/vortex/vortex_io_callbacks.h" @@ -52,14 +53,16 @@ class VortexFileBatchReader : public FileBatchReader { public: static Result> Create( const std::shared_ptr& input, int32_t batch_size, - const std::shared_ptr& arrow_pool, - const VortexPredicateOptions& predicate_options = {}, - VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE); + const VortexPredicateOptions& predicate_options, + const VortexIoCoalesceOptions& coalesce_options, + const std::shared_ptr& arrow_pool); ~VortexFileBatchReader() override; Result NextBatch() override; Result> GetFileSchema() const override; + /// Read fields must be unique and present in the file, including nested projections. + /// Schema validation failures leave the active scan and row positions unchanged. 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; @@ -71,17 +74,69 @@ class VortexFileBatchReader : public FileBatchReader { } private: + // Bitmap selection owns the smaller side and maps selected ranks to physical rows. + class SelectionState { + public: + static Result Create(const std::optional& bitmap, + uint64_t total_rows); + + bool IsEmpty() const { + return has_selection_ && selected_row_count_ == 0; + } + uint64_t SelectedRowCount() const { + return selected_row_count_; + } + uint64_t IndexCount() const { + return indices_.size(); + } + void ApplyToScan(vx_scan_options* options) const; + Result GetPhysicalRowId(uint64_t selected_rank) const; + + private: + bool has_selection_ = false; + bool excludes_ = false; + uint64_t selected_row_count_ = 0; + std::vector indices_; + }; + + // Track output positions and retain native row ids when a pushed filter removes rows. + class RowPositionState { + public: + void Reset() { + *this = RowPositionState{}; + } + void EnableFilterRowIds() { + has_filter_ = true; + } + void SetEof() { + previous_batch_row_count_ = 0; + } + Status CaptureRowIndices(const std::shared_ptr& array); + void RecordBatch(uint64_t row_count); + Result GetPhysicalRowId(uint64_t batch_row_id, + const SelectionState& selection) const; + + private: + uint64_t rows_emitted_ = 0; + uint64_t previous_first_row_ = std::numeric_limits::max(); + uint64_t previous_batch_row_count_ = 0; + bool has_filter_ = false; + std::shared_ptr previous_batch_row_idx_; + }; + VortexFileBatchReader(const std::shared_ptr& input, int32_t batch_size, std::shared_ptr input_context, VxSessionPtr session, VxDataSourcePtr data_source, VxScanPtr scan, const std::shared_ptr& file_schema, uint64_t total_rows, - const std::shared_ptr& arrow_pool, - const VortexPredicateOptions& predicate_options); + const VortexPredicateOptions& predicate_options, + const std::shared_ptr& arrow_pool); /// Pull the next Arrow array from the current partition stream, opening the next partition's /// stream when the current one is exhausted. Returns nullptr at end of scan. Result> ReadNextArray(); - /// Open the next partition and fill `current_stream_` via `vx_partition_scan_arrow`. + /// Create the single-use scan with the current projection, selection and predicate. + Status CreateScan(); + /// Open the next partition and fill `current_stream_` with bounded Arrow batches. /// Returns false when there are no more partitions. Result OpenNextPartitionStream(); /// Reshape a scanned batch to `read_schema_`, selecting columns by name and recursing into @@ -114,21 +169,8 @@ class VortexFileBatchReader : public FileBatchReader { std::shared_ptr current_batch_; int64_t current_batch_offset_ = 0; - uint64_t rows_emitted_ = 0; - uint64_t previous_first_row_ = std::numeric_limits::max(); - uint64_t previous_batch_row_count_ = 0; - - // Keep the smaller side of the bitmap. Dense selections store excluded physical row ids, - // avoiding an O(file rows) index list both here and in the native scan's copy. - bool has_selection_ = false; - bool selection_excludes_ = false; - uint64_t selected_row_count_ = 0; - std::vector selection_indices_; - - // A pushed filter returns physical row ids in a hidden _row_idx column. Retain the last - // batch's column because emitted offsets no longer identify physical rows after filtering. - bool has_filter_ = false; - std::shared_ptr previous_batch_row_idx_; + SelectionState selection_; + RowPositionState row_position_; std::shared_ptr arrow_pool_; VortexPredicateOptions predicate_options_; diff --git a/src/paimon/format/vortex/vortex_file_format.cpp b/src/paimon/format/vortex/vortex_file_format.cpp index 0cec183b3..4d098dfdd 100644 --- a/src/paimon/format/vortex/vortex_file_format.cpp +++ b/src/paimon/format/vortex/vortex_file_format.cpp @@ -22,6 +22,8 @@ #include "arrow/c/bridge.h" #include "fmt/format.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/format/vortex/vortex_reader_builder.h" #include "paimon/format/vortex/vortex_stats_extractor.h" @@ -31,33 +33,32 @@ namespace paimon::vortex { Result> VortexFileFormat::CreateReaderBuilder( int32_t batch_size) const { - VortexPredicateOptions predicate_options; - const auto membership_mode = options_.find("vortex.read.membership-pushdown"); - if (membership_mode != options_.end()) { - if (membership_mode->second == "auto") { - predicate_options.membership_pushdown = MembershipPushdownMode::AUTO; - } else if (membership_mode->second == "force") { - predicate_options.membership_pushdown = MembershipPushdownMode::FORCE; - } else if (membership_mode->second == "off") { - predicate_options.membership_pushdown = MembershipPushdownMode::OFF; - } else { - return Status::Invalid( - fmt::format("vortex.read.membership-pushdown must be auto, force or off, got: {}", - membership_mode->second)); - } + PAIMON_ASSIGN_OR_RAISE(uint32_t max_in_values, OptionsUtils::GetValueFromMap( + options_, kVortexReadInPushdownMaxValues, + kDefaultVortexReadInPushdownMaxValues)); + const VortexPredicateOptions predicate_options{max_in_values}; + PAIMON_ASSIGN_OR_RAISE(int64_t max_gap, + OptionsUtils::GetValueFromMap( + options_, kVortexReadIoCoalesceMaxGap, + static_cast(kDefaultVortexReadIoCoalesceMaxGap))); + PAIMON_ASSIGN_OR_RAISE(int64_t max_size, + OptionsUtils::GetValueFromMap( + options_, kVortexReadIoCoalesceMaxSize, + static_cast(kDefaultVortexReadIoCoalesceMaxSize))); + if (max_gap < 0) { + return Status::Invalid(fmt::format("{} must be non-negative", kVortexReadIoCoalesceMaxGap)); } - VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE; - const auto coalesce = options_.find("vortex.read.io-coalesce"); - if (coalesce != options_.end()) { - if (coalesce->second == "local") { - coalesce_mode = VortexIoCoalesceMode::LOCAL; - } else if (coalesce->second != "object-storage") { - return Status::Invalid( - fmt::format("vortex.read.io-coalesce must be object-storage or local, got: {}", - coalesce->second)); - } + if (max_size <= 0) { + return Status::Invalid(fmt::format("{} must be positive", kVortexReadIoCoalesceMaxSize)); } - return std::make_unique(batch_size, predicate_options, coalesce_mode); + if (max_gap > max_size) { + return Status::Invalid(fmt::format("{} must not exceed {}, got {} > {}", + kVortexReadIoCoalesceMaxGap, + kVortexReadIoCoalesceMaxSize, max_gap, max_size)); + } + return std::make_unique( + batch_size, predicate_options, + VortexIoCoalesceOptions{static_cast(max_gap), static_cast(max_size)}); } Result> VortexFileFormat::CreateWriterBuilder( diff --git a/src/paimon/format/vortex/vortex_file_format_test.cpp b/src/paimon/format/vortex/vortex_file_format_test.cpp index ae8791de8..b3f6345ee 100644 --- a/src/paimon/format/vortex/vortex_file_format_test.cpp +++ b/src/paimon/format/vortex/vortex_file_format_test.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include @@ -57,6 +58,7 @@ #include "paimon/format/format_stats_extractor.h" #include "paimon/format/format_writer.h" #include "paimon/format/reader_builder.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_io_callbacks.h" #include "paimon/format/writer_builder.h" #include "paimon/fs/file_system.h" @@ -172,14 +174,14 @@ class VortexFileFormatTest : public ::testing::Test { Result> CreatePredicateReader( const std::string& path, const std::shared_ptr& schema, - const std::shared_ptr& predicate, const std::string& mode, - const std::string& io_coalesce_mode = "") const { + const std::shared_ptr& predicate, const std::string& max_values, + const std::map& io_coalesce_options = {}) const { std::map options = {{"file.format", "vortex"}}; - if (!mode.empty()) { - options.emplace("vortex.read.membership-pushdown", mode); + if (!max_values.empty()) { + options.emplace(kVortexReadInPushdownMaxValues, max_values); } - if (!io_coalesce_mode.empty()) { - options.emplace("vortex.read.io-coalesce", io_coalesce_mode); + for (const auto& [key, value] : io_coalesce_options) { + options[key] = value; } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr format, FileFormatFactory::Get("vortex", options)); @@ -218,6 +220,50 @@ class VortexFileFormatTest : public ::testing::Test { return row_ids; } + // Verify both native candidates and the real SDK residual predicate against an independent + // expected result. Every returned field is also checked against the source physical row. + void AssertNativeAndResidualRows(std::unique_ptr reader, + const std::shared_ptr& predicate, + const std::shared_ptr& source, + const std::vector& expected_candidates, + const std::vector& expected_matches) const { + const auto residual = std::dynamic_pointer_cast(predicate); + ASSERT_NE(residual, nullptr); + const auto& source_rows = checked_cast(*source); + std::vector candidates; + std::vector matches; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + const auto& rows = checked_cast(*array); + ASSERT_EQ(rows.num_fields(), source_rows.num_fields()); + ASSERT_OK_AND_ASSIGN(std::vector mask, residual->Test(*array, arrow_pool_.get())); + ASSERT_EQ(mask.size(), static_cast(array->length())); + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t row_id, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_LT(row_id, static_cast(source->length())); + candidates.push_back(row_id); + for (int32_t field = 0; field < rows.num_fields(); ++field) { + const auto actual = rows.field(field)->GetScalar(i); + const auto expected = + source_rows.field(field)->GetScalar(static_cast(row_id)); + ASSERT_TRUE(actual.ok()); + ASSERT_TRUE(expected.ok()); + ASSERT_TRUE(actual.ValueOrDie()->Equals(*expected.ValueOrDie())); + } + if (mask[i] != 0) { + matches.push_back(row_id); + } + } + } + ASSERT_EQ(candidates, expected_candidates); + ASSERT_EQ(matches, expected_matches); + } + static std::shared_ptr ReaderTestSchema() { return arrow::schema( {arrow::field("id", arrow::int32(), false), arrow::field("name", arrow::utf8())}); @@ -257,6 +303,16 @@ class VortexFileFormatTest : public ::testing::Test { return reader->SetReadSchema(&ffi_schema, predicate, selection); } + void AssertSelectionMetrics(const FileBatchReader* reader, uint64_t expected_rows, + uint64_t expected_buffer_rows) const { + const auto metrics = reader->GetReaderMetrics(); + ASSERT_OK_AND_ASSIGN(uint64_t rows, metrics->GetCounter(VortexMetrics::kSelectionRowCount)); + ASSERT_EQ(rows, expected_rows); + ASSERT_OK_AND_ASSIGN(uint64_t buffer_rows, + metrics->GetCounter(VortexMetrics::kSelectionBufferRowCount)); + ASSERT_EQ(buffer_rows, expected_buffer_rows); + } + Result> ReadNextBatch(FileBatchReader* reader) const { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); if (BatchReader::IsEofBatch(batch)) { @@ -316,6 +372,59 @@ class RecordingInputStream : public InputStream { int close_count_ = 0; }; +// Record callback I/O without serializing the underlying reads. +class ReadAtTrackingInputStream : public InputStream { + public: + explicit ReadAtTrackingInputStream(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 { + RecordRead(offset, size); + return input_->Read(buffer, size, offset); + } + void ReadAsync(char* buffer, int64_t size, int64_t offset, + std::function&& callback) override { + RecordRead(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(); + } + void ClearReadAtRanges() { + std::lock_guard lock(mutex_); + read_at_ranges_.clear(); + } + std::vector> GetReadAtRanges() const { + std::lock_guard lock(mutex_); + return read_at_ranges_; + } + + private: + void RecordRead(int64_t offset, int64_t size) { + std::lock_guard lock(mutex_); + read_at_ranges_.emplace_back(offset, size); + } + + std::shared_ptr input_; + mutable std::mutex mutex_; + std::vector> read_at_ranges_; +}; + } // namespace // The stream must stay open while a read callback still owns the context, and close exactly once @@ -473,6 +582,167 @@ TEST_F(VortexFileFormatTest, ReadRetypesStoredTimestampTimezone) { ASSERT_TRUE(values.IsNull(1)); } +// Java's LTZ schema and predicate literals use UTC. Normalize only the timezone metadata; +// changing the integer timestamps here would change the instants being stored. +TEST_F(VortexFileFormatTest, WriterNormalizesTimestampTimezonesAcrossPrecisions) { + const std::vector units = { + arrow::TimeUnit::SECOND, arrow::TimeUnit::MILLI, arrow::TimeUnit::MICRO, + arrow::TimeUnit::NANO}; + const std::string json = + R"([[-1000001,-1000001,-1000001],[-1,-1,-1],[null,null,0],[0,0,0],[1,1,1],[123456789,123456789,123456789]])"; + const std::vector zones = {"Asia/Shanghai", "America/New_York"}; + for (const std::string& zone : zones) { + for (arrow::TimeUnit::type unit : units) { + SCOPED_TRACE(fmt::format("zone={}, unit={}", zone, static_cast(unit))); + const auto make_schema = [unit](const std::string& timezone) { + const auto metadata = arrow::key_value_metadata({"timestamp-tag"}, {"retained"}); + return arrow::schema( + {arrow::field("plain", arrow::timestamp(unit), true, metadata), + arrow::field("ltz", arrow::timestamp(unit, timezone), true, metadata), + arrow::field("utc", arrow::timestamp(unit, "UTC"), false, metadata)}, + metadata); + }; + const auto schema = make_schema(zone); + const auto original_schema = make_schema(zone); + const auto utc_schema = make_schema("UTC"); + std::shared_ptr written = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + std::shared_ptr expected_utc = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(utc_schema->fields()), + json) + .ValueOrDie(); + const std::string path = + PathUtil::JoinPath(directory_->Str(), fmt::format("timezone-{}-{}.vortex", + zone.substr(zone.find('/') + 1), + static_cast(unit))); + ASSERT_OK(WriteFile(path, schema, written, /*batch_size=*/2)); + ASSERT_TRUE(schema->Equals(*original_schema, /*check_metadata=*/true)); + ASSERT_TRUE(written->type()->Equals(arrow::struct_(original_schema->fields()), + /*check_metadata=*/true)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> file_schema, + reader->GetFileSchema()); + std::shared_ptr stored_schema = + arrow::ImportSchema(file_schema.get()).ValueOrDie(); + ASSERT_TRUE(stored_schema->Equals(*utc_schema)) << stored_schema->ToString(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_utc, + ReadFile(path, utc_schema, /*batch_size=*/2, + /*expected_row_count=*/6)); + ASSERT_TRUE(actual_utc->Equals(expected_utc)); + ASSERT_TRUE(actual_utc->type()->Equals(expected_utc->type(), /*check_metadata=*/true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_local, + ReadFile(path, schema, /*batch_size=*/3, + /*expected_row_count=*/6)); + ASSERT_TRUE(actual_local->Equals(written)); + ASSERT_TRUE(actual_local->type()->Equals(written->type(), /*check_metadata=*/true)); + } + } +} + +TEST_F(VortexFileFormatTest, WriterNormalizesNestedTimestampTimezones) { + const auto make_schema = [](bool utc) { + const auto metadata = arrow::key_value_metadata({"nested-tag"}, {"retained"}); + const auto local_millis = + arrow::timestamp(arrow::TimeUnit::MILLI, utc ? "UTC" : "Asia/Shanghai"); + const auto local_nanos = + arrow::timestamp(arrow::TimeUnit::NANO, utc ? "UTC" : "America/New_York"); + const auto local_seconds = + arrow::timestamp(arrow::TimeUnit::SECOND, utc ? "UTC" : "America/New_York"); + return arrow::schema( + {arrow::field( + "r", + arrow::struct_( + {arrow::field("ts", local_millis, false, metadata), + arrow::field("plain", arrow::timestamp(arrow::TimeUnit::MICRO), true, + metadata), + arrow::field("l", + arrow::list(arrow::field("item", local_nanos, true, metadata)), + true, metadata)}), + false, metadata), + arrow::field( + "f", + arrow::fixed_size_list(arrow::field("item", local_seconds, true, metadata), 2), + true, metadata)}, + metadata); + }; + const auto schema = make_schema(/*utc=*/false); + const auto original_schema = make_schema(/*utc=*/false); + const auto utc_schema = make_schema(/*utc=*/true); + const std::string json = R"([ + [{"ts":-1,"plain":-1000001,"l":[-1,null,123456789]},[-1,0]], + [{"ts":-2,"plain":null,"l":null},null], + [{"ts":0,"plain":null,"l":[]},[1,null]], + [{"ts":1,"plain":1,"l":[0]},[null,-1000001]] + ])"; + std::shared_ptr written = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + std::shared_ptr expected_utc = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(utc_schema->fields()), json) + .ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "nested-timezones.vortex"); + ASSERT_OK(WriteFile(path, schema, written, /*batch_size=*/2)); + ASSERT_TRUE(schema->Equals(*original_schema, /*check_metadata=*/true)); + ASSERT_TRUE(written->type()->Equals(arrow::struct_(original_schema->fields()), + /*check_metadata=*/true)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> file_schema, reader->GetFileSchema()); + std::shared_ptr stored_schema = + arrow::ImportSchema(file_schema.get()).ValueOrDie(); + ASSERT_TRUE(stored_schema->Equals(*utc_schema)) << stored_schema->ToString(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_utc, + ReadFile(path, utc_schema, /*batch_size=*/3, /*expected_row_count=*/4)); + ASSERT_TRUE(actual_utc->Equals(expected_utc)); + ASSERT_TRUE(actual_utc->type()->Equals(expected_utc->type(), /*check_metadata=*/true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual_local, + ReadFile(path, schema, /*batch_size=*/1, /*expected_row_count=*/4)); + ASSERT_TRUE(actual_local->Equals(written)); + ASSERT_TRUE(actual_local->type()->Equals(written->type(), /*check_metadata=*/true)); +} + +// Use the native file reader so an SDK residual filter cannot hide declined LTZ pushdown. +TEST_F(VortexFileFormatTest, NormalizedTimestampTimezoneSupportsNativePredicates) { + const auto schema = arrow::schema( + {arrow::field("ltz", arrow::timestamp(arrow::TimeUnit::MILLI, "Asia/Shanghai"))}); + std::shared_ptr written = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), + R"([[-1001],[-1],[null],[0],[1],[1000]])") + .ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "timezone-predicates.vortex"); + ASSERT_OK(WriteFile(path, schema, written, /*batch_size=*/2)); + const Literal minus_one(Timestamp::FromEpochMillis(-1)); + const Literal zero(Timestamp::FromEpochMillis(0)); + const Literal one(Timestamp::FromEpochMillis(1)); + const Literal thousand(Timestamp::FromEpochMillis(1000)); + struct PredicateCase { + std::string name; + std::shared_ptr predicate; + std::vector row_ids; + }; + const std::vector cases = { + {"equal", PredicateBuilder::Equal(0, "ltz", FieldType::TIMESTAMP, minus_one), {1}}, + {"greater", PredicateBuilder::GreaterThan(0, "ltz", FieldType::TIMESTAMP, zero), {4, 5}}, + {"between", + PredicateBuilder::Between(0, "ltz", FieldType::TIMESTAMP, minus_one, one), + {1, 3, 4}}, + {"in", PredicateBuilder::In(0, "ltz", FieldType::TIMESTAMP, {minus_one, thousand}), {1, 5}}, + {"not-in", + PredicateBuilder::NotIn(0, "ltz", FieldType::TIMESTAMP, {minus_one, thousand}), + {0, 3, 4}}, + {"is-null", PredicateBuilder::IsNull(0, "ltz", FieldType::TIMESTAMP), {2}}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreatePredicateReader(path, schema, test_case.predicate, + /*max_values=*/"")); + AssertNativeAndResidualRows(std::move(reader), test_case.predicate, written, + test_case.row_ids, test_case.row_ids); + } +} + TEST_F(VortexFileFormatTest, WriteThenRead) { std::string path = PathUtil::JoinPath(directory_->Str(), "data.vortex"); arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false), @@ -541,6 +811,89 @@ TEST_F(VortexFileFormatTest, WriteThenReadListOfString) { AssertReadWithBatchSizes(path, schema, expected, {1, 2, 3, 5}); } +// NULL parent lists must export correctly even when a bounded batch has no child ranges. +TEST_F(VortexFileFormatTest, NestedListNullAndEmptyBatchesRoundTrip) { + struct NestedListCase { + std::string name; + std::shared_ptr type; + std::vector> values; + }; + const std::vector cases = { + {"list-list-int", + arrow::list(arrow::list(arrow::int32())), + {{"null-only", R"([[null]])"}, + {"all-null", R"([[null],[null]])"}, + {"empty-only", R"([[[]]])"}, + {"empty-inner", R"([[[[]]]])"}, + {"nullable-inner", R"([[[null,[],[null],[1,null,2]]]])"}, + {"mixed", R"([[null],[[]],[[[]]],[[null]],[[[1,2],[],null,[3]]],[null]])"}}}, + {"list-list-string", + arrow::list(arrow::list(arrow::utf8())), + {{"null-only", R"([[null]])"}, + {"all-null", R"([[null],[null]])"}, + {"empty-only", R"([[[]]])"}, + {"empty-inner", R"([[[[]]]])"}, + {"nullable-inner", R"([[[null,[],[null],["",null,"one"]]]])"}, + {"mixed", + R"([[null],[[]],[[[]]],[[null]],[[["one","a-long-string-beyond-the-inline-view-limit"],[],null,["three",null]]],[null]])"}}}, + {"list-struct-list-string", + arrow::list(arrow::struct_({arrow::field("tags", arrow::list(arrow::utf8()))})), + {{"null-only", R"([[null]])"}, + {"all-null", R"([[null],[null]])"}, + {"empty-only", R"([[[]]])"}, + {"empty-inner", R"([[[{"tags":[]}]]])"}, + {"nullable-inner", R"([[[null,{"tags":null},{"tags":[]},{"tags":[null,""]}]]])"}, + {"mixed", + R"([[null],[[]],[[null]],[[{"tags":[]}]],[[{"tags":["one","a-long-string-beyond-the-inline-view-limit",null]},{"tags":null},null]],[null]])"}}}}; + for (const NestedListCase& test_case : cases) { + const auto schema = arrow::schema({arrow::field("nested", test_case.type)}); + for (const auto& [name, json] : test_case.values) { + SCOPED_TRACE(fmt::format("type={}, case={}", test_case.name, name)); + const auto source_result = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json); + ASSERT_TRUE(source_result.ok()) << source_result.status().ToString(); + const std::shared_ptr source = source_result.ValueOrDie(); + ASSERT_TRUE(source->ValidateFull().ok()); + const std::string path = PathUtil::JoinPath( + directory_->Str(), fmt::format("{}-{}.vortex", test_case.name, name)); + // Keep mixed rows in one input batch so the reader slices shared child elements. + ASSERT_OK(WriteFile(path, schema, source, /*batch_size=*/16)); + for (int32_t batch_size : {1, 2}) { + SCOPED_TRACE(fmt::format("batch_size={}", batch_size)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateReader(path, batch_size)); + ASSERT_OK(ConfigureReader(reader.get(), schema)); + ASSERT_OK_AND_ASSIGN(uint64_t total_rows, reader->GetNumberOfRows()); + ASSERT_EQ(total_rows, static_cast(source->length())); + int64_t rows_seen = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, + ReadNextBatch(reader.get())); + if (batch == nullptr) { + break; + } + ASSERT_TRUE(batch->ValidateFull().ok()); + ASSERT_GT(batch->length(), 0); + ASSERT_LE(batch->length(), batch_size); + ASSERT_LE(rows_seen + batch->length(), source->length()); + ASSERT_TRUE(batch->type()->Equals(arrow::struct_(schema->fields()))); + ASSERT_TRUE(batch->Equals(source->Slice(rows_seen, batch->length()))) + << batch->ToString(); + for (int64_t i = 0; i < batch->length(); ++i) { + ASSERT_OK_AND_ASSIGN( + uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_EQ(physical_row, static_cast(rows_seen + i)); + } + rows_seen += batch->length(); + } + ASSERT_EQ(rows_seen, source->length()); + ASSERT_NOK(reader->GetPreviousBatchFileRowId(0)); + } + } + } +} + // Regression for the FIXED_SIZE_LIST branch of NormalizeViewArray: without recursion into the // elements, a fixed_size_list column comes back as fixed_size_list (Vortex // hardcodes Utf8 -> Utf8View), mismatching the normalized read schema and silently corrupting the @@ -814,26 +1167,14 @@ TEST_F(VortexFileFormatTest, ProjectedReadRejectsIncompatibleLeafType) { ASSERT_TRUE(rejected.status().IsInvalid()) << rejected.status().ToString(); } -// A read field the file does not carry is null-filled instead of rejected: that is the -// schema-evolution answer the rest of the read path gives, so every file yields the read schema. -TEST_F(VortexFileFormatTest, ProjectedReadNullFillsFieldMissingFromFile) { - std::string path = PathUtil::JoinPath(directory_->Str(), "missing-field.vortex"); - arrow::FieldVector fields = {arrow::field("id", arrow::int32(), false)}; - std::shared_ptr schema = arrow::schema(fields); - std::shared_ptr written = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([[1],[2],[3]])") - .ValueOrDie(); - ASSERT_OK(WriteFile(path, schema, written, /*batch_size=*/2)); - - arrow::FieldVector projected_fields = {fields[0], arrow::field("extra", arrow::utf8())}; - std::shared_ptr projected = arrow::schema(projected_fields); - std::shared_ptr expected = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(projected_fields), - R"([[1,null],[2,null],[3,null]])") - .ValueOrDie(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, - ReadFile(path, projected, /*batch_size=*/2, /*expected_row_count=*/3)); - ASSERT_TRUE(actual->Equals(expected)) << actual->ToString() << "\nvs\n" << expected->ToString(); +TEST_F(VortexFileFormatTest, ProjectedReadRejectsFieldMissingFromFile) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto projection = + arrow::schema({arrow::field("id", arrow::int32()), arrow::field("extra", arrow::utf8())}); + const Status status = ConfigureReader(reader.get(), projection); + ASSERT_TRUE(status.IsInvalid()) << status.ToString(); + ASSERT_NE(status.message().find("extra"), std::string::npos); } // Regression guard for ordered scans: 200 sequential rows span many Vortex chunks, and a small read @@ -1194,10 +1535,8 @@ TEST_F(VortexFileFormatTest, DenseSelectionKeepsOnlyExcludedIndicesAndPreservesN selection.AddRange(0, kRows); selection.RemoveRange(kDeletedRow, kDeletedRow + 1); ASSERT_OK(ConfigureReader(reader.get(), schema, nullptr, selection)); - ASSERT_OK_AND_ASSIGN(uint64_t index_count, - reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); - // Both the host and the native FFI copy contain one excluded index, not kRows - 1 indices. - ASSERT_EQ(index_count, uint64_t{1}); + // Both the host and the native FFI copy contain only one excluded row id. + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), kRows - 1, 1)); uint64_t emitted = 0; while (true) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); @@ -1225,9 +1564,7 @@ TEST_F(VortexFileFormatTest, DenseSelectionKeepsOnlyExcludedIndicesAndPreservesN } ASSERT_EQ(emitted, uint64_t{kRows - 1}); ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); - ASSERT_OK_AND_ASSIGN(index_count, - reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); - ASSERT_EQ(index_count, uint64_t{1}); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), kRows - 1, 1)); } TEST_F(VortexFileFormatTest, DenseSelectionMapsConsecutiveExcludedRowsAndSparseSelections) { @@ -1239,9 +1576,9 @@ TEST_F(VortexFileFormatTest, DenseSelectionMapsConsecutiveExcludedRowsAndSparseS ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From(selected_rows))); - ASSERT_OK_AND_ASSIGN(uint64_t index_count, reader->GetReaderMetrics()->GetCounter( - "vortex.selection-index-count")); - ASSERT_EQ(index_count, std::min(selected_rows.size(), size_t{6} - selected_rows.size())); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics( + reader.get(), selected_rows.size(), + std::min(selected_rows.size(), size_t{6} - selected_rows.size()))); std::vector actual; while (true) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); @@ -1266,9 +1603,7 @@ TEST_F(VortexFileFormatTest, FullEmptyAndDenseSelectionsCanBeReconfiguredWithout RoaringBitmap32 full; full.AddRange(0, 6); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, full)); - ASSERT_OK_AND_ASSIGN(uint64_t index_count, - reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); - ASSERT_EQ(index_count, uint64_t{0}); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 6, 0)); ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); ASSERT_NE(first, nullptr); ASSERT_EQ(first->length(), 2); @@ -1276,30 +1611,29 @@ TEST_F(VortexFileFormatTest, FullEmptyAndDenseSelectionsCanBeReconfiguredWithout ASSERT_EQ(row, uint64_t{1}); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32())); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 0, 0)); ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); ASSERT_EQ(first, nullptr); ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({1, 2, 3, 4, 5}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 5, 1)); ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); ASSERT_EQ(row, uint64_t{1}); ASSERT_NOK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({0, 1, 2, 3, 6}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 5, 1)); ASSERT_NOK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({-1, 1, 2, 3, 4}))); - ASSERT_OK_AND_ASSIGN(index_count, - reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); - ASSERT_EQ(index_count, uint64_t{1}); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 5, 1)); ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(1)); ASSERT_EQ(row, uint64_t{2}); ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); ASSERT_EQ(row, uint64_t{3}); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema())); - ASSERT_OK_AND_ASSIGN(index_count, - reader->GetReaderMetrics()->GetCounter("vortex.selection-index-count")); - ASSERT_EQ(index_count, uint64_t{0}); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 6, 0)); ASSERT_OK_AND_ASSIGN(first, ReadNextBatch(reader.get())); ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(0)); ASSERT_EQ(row, uint64_t{0}); @@ -1363,27 +1697,254 @@ TEST_F(VortexFileFormatTest, FilterPushdownReturnsMatchingRowsAndPhysicalRowIds) ASSERT_EQ(got_row_ids, (std::vector{4, 5, 6, 7, 8, 9})); } -TEST_F(VortexFileFormatTest, RejectsInvalidMembershipPushdownMode) { - for (const std::string mode : {"", "AUTO", "enabled"}) { - SCOPED_TRACE(fmt::format("mode={}", mode)); +TEST_F(VortexFileFormatTest, PredicateTypeErrorsPropagateFromLazyScanCreation) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const Literal int_value(static_cast(1)); + const Literal bigint_value(static_cast(1)); + const std::vector>> cases = { + {"comparison-field", PredicateBuilder::Equal(0, "id", FieldType::BIGINT, bigint_value)}, + {"comparison-literal", PredicateBuilder::Equal(0, "id", FieldType::INT, bigint_value)}, + {"in-field", PredicateBuilder::In(0, "id", FieldType::BIGINT, {bigint_value})}, + {"in-literal", PredicateBuilder::In(0, "id", FieldType::INT, {int_value, bigint_value})}, + {"not-in-field", PredicateBuilder::NotIn(0, "id", FieldType::BIGINT, {bigint_value})}, + {"not-in-literal", + PredicateBuilder::NotIn(0, "id", FieldType::INT, {int_value, bigint_value})}}; + const std::shared_ptr valid = + PredicateBuilder::Equal(0, "id", FieldType::INT, int_value); + for (const auto& [name, predicate] : cases) { + SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + // Schema configuration succeeds; predicate conversion runs when the first scan is opened. + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), predicate)); + const Result batch = reader->NextBatch(); + ASSERT_FALSE(batch.ok()); + ASSERT_TRUE(batch.status().IsInvalid()) << batch.status().ToString(); + const std::string message = batch.status().ToString(); + ASSERT_NE(message.find("field 'id'"), std::string::npos) << message; + ASSERT_NE(message.find("expected INT"), std::string::npos) << message; + ASSERT_NE(message.find("actual BIGINT"), std::string::npos) << message; + + // A failed conversion must not prevent a later valid scan on the same reader. + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), valid)); + ASSERT_OK_AND_ASSIGN(std::vector row_ids, CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(row_ids, (std::vector{1})); + } +} + +TEST_F(VortexFileFormatTest, NativePredicateErrorsPropagateFromLazyScanCreation) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const Literal invalid_utf8(FieldType::STRING, "\xff", 1); + const std::vector> invalid = { + PredicateBuilder::Equal(1, "name", FieldType::STRING, invalid_utf8), + PredicateBuilder::In(1, "name", FieldType::STRING, + {Literal(FieldType::STRING, "one", 3), invalid_utf8}), + PredicateBuilder::NotIn(1, "name", FieldType::STRING, + {Literal(FieldType::STRING, "one", 3), invalid_utf8})}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr unsupported, + PredicateBuilder::StartsWith(1, "name", FieldType::STRING, + Literal(FieldType::STRING, "o", 1))); + const std::shared_ptr valid = + PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(int32_t{1})); + for (const std::shared_ptr& predicate : invalid) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({unsupported, predicate})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({unsupported, predicate})); + for (const std::shared_ptr& candidate : {predicate, conjunction, disjunction}) { + SCOPED_TRACE(candidate->GetFunction().ToString()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), candidate)); + const Result batch = reader->NextBatch(); + ASSERT_FALSE(batch.ok()); + ASSERT_TRUE(batch.status().IsInvalid()) << batch.status().ToString(); + const std::string message = batch.status().ToString(); + ASSERT_NE(message.find("create Vortex predicate scalar"), std::string::npos) << message; + ASSERT_NE(message.find("invalid utf-8"), std::string::npos) << message; + + // A native conversion failure must release its handles and allow a valid new scan. + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), valid)); + ASSERT_OK_AND_ASSIGN(std::vector row_ids, + CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(row_ids, (std::vector{1})); + } + } +} + +TEST_F(VortexFileFormatTest, NativeTemporalLimitsFallBackAndPreserveResidualResults) { + const std::shared_ptr schema = + arrow::schema({arrow::field("date", arrow::date32()), + arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"))}); + const auto data = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(schema->fields()), "[[0,0],[1,1000],[null,null],[-1,-1000]]"); + ASSERT_TRUE(data.ok()) << data.status().ToString(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "temporal-fallback.vortex"); + ASSERT_OK(WriteFile(path, schema, data.ValueOrDie(), /*batch_size=*/2)); + const int64_t maximum = std::numeric_limits::max(); + struct TemporalCase { + int32_t index; + const char* name; + FieldType type; + Literal extreme; + Literal zero; + bool above_values; + }; + const std::vector cases = { + {0, "date", FieldType::DATE, Literal(FieldType::DATE, std::numeric_limits::max()), + Literal(FieldType::DATE, 0), true}, + {0, "date", FieldType::DATE, Literal(FieldType::DATE, std::numeric_limits::min()), + Literal(FieldType::DATE, 0), false}, + {1, "ts", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(maximum / 1000, + static_cast((maximum % 1000) * 1000))), + Literal(Timestamp::FromEpochMillis(0)), true}, + {1, "ts", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(std::numeric_limits::min() / 1000)), + Literal(Timestamp::FromEpochMillis(0)), false}, + // This constant is in Paimon's documented calendar range, but not Jiff Timestamp's. + {1, "ts", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(253402300799999, 999000)), + Literal(Timestamp::FromEpochMillis(0)), true}}; + const std::vector candidates = {0, 1, 2, 3}; + const std::vector non_null = {0, 1, 3}; + const std::vector empty; + for (const TemporalCase& item : cases) { + SCOPED_TRACE(item.name); + SCOPED_TRACE(item.above_values); + const std::vector, std::vector>> predicates = + {{PredicateBuilder::Equal(item.index, item.name, item.type, item.extreme), empty}, + {PredicateBuilder::NotEqual(item.index, item.name, item.type, item.extreme), non_null}, + {PredicateBuilder::LessThan(item.index, item.name, item.type, item.extreme), + item.above_values ? non_null : empty}, + {PredicateBuilder::LessOrEqual(item.index, item.name, item.type, item.extreme), + item.above_values ? non_null : empty}, + {PredicateBuilder::GreaterThan(item.index, item.name, item.type, item.extreme), + item.above_values ? empty : non_null}, + {PredicateBuilder::GreaterOrEqual(item.index, item.name, item.type, item.extreme), + item.above_values ? empty : non_null}, + {PredicateBuilder::In(item.index, item.name, item.type, {item.extreme, item.zero}), + {0}}, + {PredicateBuilder::NotIn(item.index, item.name, item.type, {item.zero, item.extreme}), + {1, 3}}}; + for (const auto& [predicate, expected] : predicates) { + SCOPED_TRACE(predicate->GetFunction().ToString()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), schema, predicate)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + std::move(reader), predicate, data.ValueOrDie(), candidates, expected)); + } + } + // The list budget changes whether native construction is attempted, never query success. + std::vector large_values(33, Literal(FieldType::DATE, 0)); + large_values[0] = Literal(FieldType::DATE, std::numeric_limits::max()); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "date", FieldType::DATE, large_values) + : PredicateBuilder::In(0, "date", FieldType::DATE, large_values); + for (const std::string limit : {"0", "32", "33"}) { + SCOPED_TRACE(limit); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreatePredicateReader(path, schema, predicate, limit)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + std::move(reader), predicate, data.ValueOrDie(), candidates, + negate ? std::vector{1, 3} : std::vector{0})); + } + } + const std::shared_ptr unsupported = PredicateBuilder::LessThan( + 0, "date", FieldType::DATE, Literal(FieldType::DATE, std::numeric_limits::max())); + const std::shared_ptr supported = + PredicateBuilder::Equal(0, "date", FieldType::DATE, Literal(FieldType::DATE, 0)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({unsupported, supported})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({supported, unsupported})); + for (const auto& [predicate, expected_candidates, expected_matches] : std::vector< + std::tuple, std::vector, std::vector>>{ + {conjunction, {0}, {0}}, {disjunction, candidates, non_null}}) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK(ConfigureReader(reader.get(), schema, predicate)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(std::move(reader), predicate, + data.ValueOrDie(), expected_candidates, + expected_matches)); + } +} + +TEST_F(VortexFileFormatTest, MembershipMaximumAcceptsDefaultsAndUnsignedValueCounts) { + const std::vector> options = { + {}, + {{kVortexReadInPushdownMaxValues, "0"}}, + {{kVortexReadInPushdownMaxValues, "1"}}, + {{kVortexReadInPushdownMaxValues, "32"}}, + {{kVortexReadInPushdownMaxValues, "128"}}, + {{kVortexReadInPushdownMaxValues, "4294967295"}}, + }; + for (const auto& option : options) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", option)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr builder, + format->CreateReaderBuilder(2)); + ASSERT_NE(builder, nullptr); + } +} + +TEST_F(VortexFileFormatTest, MembershipMaximumRejectsInvalidAndOverflowingValueCounts) { + for (const std::string value : {"", "auto", "force", "off", "-1", "-0", "1.5", "32 kb", " 32", + "32 ", "+1", "4294967296", "18446744073709551616"}) { + SCOPED_TRACE(value); ASSERT_OK_AND_ASSIGN( std::shared_ptr format, - FileFormatFactory::Get( - "vortex", {{"file.format", "vortex"}, {"vortex.read.membership-pushdown", mode}})); + FileFormatFactory::Get("vortex", {{kVortexReadInPushdownMaxValues, value}})); Result> result = format->CreateReaderBuilder(2); ASSERT_FALSE(result.ok()); ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); - ASSERT_NE(result.status().ToString().find("vortex.read.membership-pushdown"), + ASSERT_NE(result.status().ToString().find(kVortexReadInPushdownMaxValues), std::string::npos); } } -TEST_F(VortexFileFormatTest, RejectsInvalidIoCoalesceMode) { - for (const std::string mode : {"", "LOCAL", "auto", "none", "local "}) { - SCOPED_TRACE(fmt::format("mode={}", mode)); +TEST_F(VortexFileFormatTest, IoCoalesceOptionsAcceptDefaultsAndIndependentOverrides) { + const std::vector> option_sets = { + {}, + {{kVortexReadIoCoalesceMaxGap, "0"}}, + {{kVortexReadIoCoalesceMaxGap, "16777216"}}, + {{kVortexReadIoCoalesceMaxSize, "1048576"}}, + {{kVortexReadIoCoalesceMaxSize, "33554432"}}, + {{kVortexReadIoCoalesceMaxGap, "0"}, {kVortexReadIoCoalesceMaxSize, "1"}}, + {{kVortexReadIoCoalesceMaxGap, "4096"}, {kVortexReadIoCoalesceMaxSize, "4096"}}, + }; + for (size_t i = 0; i < option_sets.size(); ++i) { + SCOPED_TRACE(i); ASSERT_OK_AND_ASSIGN(std::shared_ptr format, - FileFormatFactory::Get("vortex", {{"file.format", "vortex"}, - {"vortex.read.io-coalesce", mode}})); + FileFormatFactory::Get("vortex", option_sets[i])); + ASSERT_OK_AND_ASSIGN(std::unique_ptr builder, + format->CreateReaderBuilder(2)); + ASSERT_NE(builder, nullptr); + } +} + +TEST_F(VortexFileFormatTest, IoCoalesceOptionsRejectInvalidByteCountsAndRelationships) { + for (const char* key : {kVortexReadIoCoalesceMaxGap, kVortexReadIoCoalesceMaxSize}) { + for (const std::string value : {"", "invalid", "-1", "1.5", "64 kb", " 64", "64 ", "+1", + "9223372036854775808", "18446744073709551616"}) { + SCOPED_TRACE(fmt::format("{}={}", key, value)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", {{key, value}})); + Result> result = format->CreateReaderBuilder(2); + ASSERT_FALSE(result.ok()); + ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); + ASSERT_NE(result.status().ToString().find(key), std::string::npos); + } + } + // A single override is checked against the other option's default, not clamped to it. + const std::vector> invalid_pairs = { + {{kVortexReadIoCoalesceMaxSize, "0"}}, + {{kVortexReadIoCoalesceMaxGap, "17"}, {kVortexReadIoCoalesceMaxSize, "16"}}, + {{kVortexReadIoCoalesceMaxGap, "16777217"}}, + {{kVortexReadIoCoalesceMaxSize, "1048575"}}, + }; + for (size_t i = 0; i < invalid_pairs.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", invalid_pairs[i])); Result> result = format->CreateReaderBuilder(2); ASSERT_FALSE(result.ok()); ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); @@ -1391,8 +1952,8 @@ TEST_F(VortexFileFormatTest, RejectsInvalidIoCoalesceMode) { } } -TEST_F(VortexFileFormatTest, IoCoalesceModesPreserveDataPredicatesAndRowIds) { - const std::string path = PathUtil::JoinPath(directory_->Str(), "io-coalesce-modes.vortex"); +TEST_F(VortexFileFormatTest, IoCoalesceOptionsPreserveDataPredicatesAndRowIds) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "io-coalesce-options.vortex"); const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), arrow::field("payload", arrow::utf8())}; const std::shared_ptr schema = arrow::schema(fields); @@ -1404,22 +1965,159 @@ TEST_F(VortexFileFormatTest, IoCoalesceModesPreserveDataPredicatesAndRowIds) { ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); const std::shared_ptr predicate = PredicateBuilder::Equal( /*field_index=*/0, "id", FieldType::INT, Literal(static_cast(30))); - - for (const std::string mode : {"object-storage", "local"}) { - SCOPED_TRACE(fmt::format("mode={}", mode)); - ASSERT_OK_AND_ASSIGN( - format_, FileFormatFactory::Get( - "vortex", {{"file.format", "vortex"}, {"vortex.read.io-coalesce", mode}})); + const auto projection = arrow::schema({fields[1]}); + const std::vector> option_sets = { + {}, + {{kVortexReadIoCoalesceMaxGap, "0"}}, + {{kVortexReadIoCoalesceMaxSize, "33554432"}}, + {{kVortexReadIoCoalesceMaxGap, "65536"}, {kVortexReadIoCoalesceMaxSize, "4194304"}}, + {{kVortexReadIoCoalesceMaxGap, "0"}, {kVortexReadIoCoalesceMaxSize, "1"}}, + {{kVortexReadIoCoalesceMaxGap, "4096"}, {kVortexReadIoCoalesceMaxSize, "4096"}}, + }; + for (size_t i = 0; i < option_sets.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_OK_AND_ASSIGN(format_, FileFormatFactory::Get("vortex", option_sets[i])); AssertReadWithBatchSizes(path, schema, data, {1, 2, 8}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, - CreatePredicateReader(path, schema, predicate, "auto", mode)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreatePredicateReader(path, projection, predicate, "", option_sets[i])); ASSERT_OK_AND_ASSIGN(std::vector row_ids, CollectNativeRowIds(std::move(reader))); ASSERT_EQ(row_ids, (std::vector{3})); } } -TEST_F(VortexFileFormatTest, MembershipPushdownModesStayLocalToEachReader) { - const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-modes.vortex"); +TEST_F(VortexFileFormatTest, IoCoalesceOptionsChangePhysicalReadRanges) { + // Both payload start offsets must fit within max-gap regardless of which read is polled first. + constexpr int32_t kRows = 1536; + constexpr int32_t kPayloadBytes = 512; + constexpr int64_t kMaxGap = 2LL << 20; + constexpr int64_t kSmallMaxSize = 2LL << 20; + constexpr int64_t kLargeMaxSize = 4LL << 20; + constexpr int64_t kLargeRead = 128LL << 10; + arrow::Int32Builder id_builder; + arrow::BinaryBuilder left_builder; + arrow::BinaryBuilder omitted_builder; + arrow::BinaryBuilder right_builder; + uint64_t state = 0x9e3779b97f4a7c15ULL; + for (int32_t row = 0; row < kRows; ++row) { + ASSERT_TRUE(id_builder.Append(row).ok()); + for (arrow::BinaryBuilder* builder : {&left_builder, &omitted_builder, &right_builder}) { + std::string bytes(kPayloadBytes, '\0'); + for (char& byte : bytes) { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + byte = static_cast(state & 0xff); + } + ASSERT_TRUE(builder->Append(bytes).ok()); + } + } + std::shared_ptr ids; + std::shared_ptr left; + std::shared_ptr omitted; + std::shared_ptr right; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(left_builder.Finish(&left).ok()); + ASSERT_TRUE(omitted_builder.Finish(&omitted).ok()); + ASSERT_TRUE(right_builder.Finish(&right).ok()); + const auto schema = arrow::schema({arrow::field("id", arrow::int32(), false), + arrow::field("left", arrow::binary(), false), + arrow::field("omitted", arrow::binary(), false), + arrow::field("right", arrow::binary(), false)}); + const auto data = + arrow::StructArray::Make({ids, left, omitted, right}, schema->fields()).ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "io-coalesce-ranges.vortex"); + // Distinct high-entropy columns leave a real unprojected payload between the two read columns. + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/kRows)); + const auto projection = arrow::schema({schema->field(1), schema->field(3)}); + const auto predicate = PredicateBuilder::GreaterOrEqual( + 0, "id", FieldType::INT, Literal(static_cast(kRows / 2))); + const auto& expected_left = checked_cast(*left); + const auto& expected_right = checked_cast(*right); + const std::vector> option_sets = { + {{kVortexReadIoCoalesceMaxGap, "0"}, + {kVortexReadIoCoalesceMaxSize, std::to_string(kLargeMaxSize)}}, + {{kVortexReadIoCoalesceMaxGap, std::to_string(kMaxGap)}, + {kVortexReadIoCoalesceMaxSize, std::to_string(kLargeMaxSize)}}, + {{kVortexReadIoCoalesceMaxGap, std::to_string(kMaxGap)}, + {kVortexReadIoCoalesceMaxSize, std::to_string(kSmallMaxSize)}}, + }; + std::vector>> physical_reads; + for (size_t i = 0; i < option_sets.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", option_sets[i])); + ASSERT_OK_AND_ASSIGN(std::unique_ptr builder, + format->CreateReaderBuilder(/*batch_size=*/64)); + builder->WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(path)); + auto tracking = std::make_shared(std::move(input)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, builder->Build(tracking)); + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate)); + // Opening the file reads its footer outside the segment coalescer under test. + tracking->ClearReadAtRanges(); + int64_t expected_row = kRows / 2; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); + const auto& batch = checked_cast(*array); + const auto& actual_left = checked_cast(*batch.field(0)); + const auto& actual_right = checked_cast(*batch.field(1)); + for (int64_t row = 0; row < array->length(); ++row) { + ASSERT_LT(expected_row, kRows); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(row))); + ASSERT_EQ(physical_row, static_cast(expected_row)); + ASSERT_EQ(actual_left.GetString(row), expected_left.GetString(expected_row)); + ASSERT_EQ(actual_right.GetString(row), expected_right.GetString(expected_row)); + ++expected_row; + } + } + ASSERT_EQ(expected_row, kRows); + reader.reset(); + physical_reads.push_back(tracking->GetReadAtRanges()); + } + + auto baseline = physical_reads[0]; + std::sort(baseline.begin(), baseline.end()); + const auto covers = [](const std::pair& merged, + const std::pair& first, + const std::pair& second) { + return merged.first <= first.first && + merged.first + merged.second >= second.first + second.second; + }; + bool observed_limited_merge = false; + for (size_t i = 0; i < baseline.size(); ++i) { + if (baseline[i].second <= kLargeRead) { + continue; + } + for (size_t j = i + 1; j < baseline.size(); ++j) { + const int64_t hole = baseline[j].first - (baseline[i].first + baseline[i].second); + const int64_t span = baseline[j].first + baseline[j].second - baseline[i].first; + if (baseline[j].second <= kLargeRead || hole <= kLargeRead || hole > kMaxGap || + span <= kSmallMaxSize + (64LL << 10)) { + continue; + } + const bool merged = std::any_of( + physical_reads[1].begin(), physical_reads[1].end(), + [&](const auto& range) { return covers(range, baseline[i], baseline[j]); }); + const bool limited = std::any_of( + physical_reads[2].begin(), physical_reads[2].end(), + [&](const auto& range) { return covers(range, baseline[i], baseline[j]); }); + observed_limited_merge |= merged && !limited; + } + } + ASSERT_TRUE(observed_limited_merge) + << "Expected a large read to bridge two baseline payload reads only with the larger span: " + << ::testing::PrintToString(physical_reads); +} + +TEST_F(VortexFileFormatTest, MembershipMaximumStaysLocalToEachReader) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-limits.vortex"); const arrow::FieldVector fields = {arrow::field("id", arrow::int32())}; const std::shared_ptr schema = arrow::schema(fields); const std::shared_ptr data = @@ -1438,58 +2136,107 @@ TEST_F(VortexFileFormatTest, MembershipPushdownModesStayLocalToEachReader) { PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(3)))}; const std::vector> filtered_rows = {{1, 3, 5}, {2, 4, 6, 7}, {4}}; const std::vector all_rows = {0, 1, 2, 3, 4, 5, 6, 7}; - // Empty mode means the option is omitted, exercising the default auto policy. Configure all - // readers before consuming any, so later force/off readers cannot alter an earlier reader. - const std::vector modes = {"force", "off", "", "auto", "force", "off"}; + // Configure all readers before consuming any. Omission uses the default of 32; creating a + // later reader with a larger or zero maximum must not alter previously configured readers. + const std::vector maxima = {"64", "0", "", "32", "63", "64"}; std::vector> readers; std::vector> expected; for (size_t predicate_idx = 0; predicate_idx < predicates.size(); ++predicate_idx) { - for (const std::string& mode : modes) { + for (const std::string& maximum : maxima) { ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, - CreatePredicateReader(path, schema, predicates[predicate_idx], mode)); + CreatePredicateReader(path, schema, predicates[predicate_idx], maximum)); readers.push_back(std::move(reader)); - expected.push_back(predicate_idx == 2 || mode == "force" ? filtered_rows[predicate_idx] + expected.push_back(predicate_idx == 2 || maximum == "64" ? filtered_rows[predicate_idx] : all_rows); } } for (size_t i = 0; i < readers.size(); ++i) { + const size_t predicate_idx = i / maxima.size(); SCOPED_TRACE( - fmt::format("predicate={}, mode={}", i / modes.size(), modes[i % modes.size()])); - ASSERT_OK_AND_ASSIGN(std::vector rows, - CollectNativeRowIds(std::move(readers[i]))); - ASSERT_EQ(rows, expected[i]); + fmt::format("predicate={}, maximum={}", predicate_idx, maxima[i % maxima.size()])); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(std::move(readers[i]), predicates[predicate_idx], data, + expected[i], filtered_rows[predicate_idx])); } } -TEST_F(VortexFileFormatTest, AutoPushesCompactLargeMembershipPredicates) { - const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-ranges.vortex"); - const arrow::FieldVector fields = {arrow::field("id", arrow::int32())}; +TEST_F(VortexFileFormatTest, ContinuousAndDuplicateListsUseRawLengthAndPreserveResidualResults) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-raw-length.vortex"); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::utf8())}; const std::shared_ptr schema = arrow::schema(fields); const std::shared_ptr data = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), - R"([[null],[0],[1],[2],[3],[4],[5],[1000]])") + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields), + R"([[null,"null-id"],[0,"zero"],[1,null],[2,"two"],[3,"three"],[4,null],[5,"five"],[1000,"thousand"]])") .ValueOrDie(); ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); - - std::vector contiguous; - for (int32_t i = 0; i < 64; ++i) { - contiguous.emplace_back(i); - } - const std::vector> predicates = { - PredicateBuilder::In(0, "id", FieldType::INT, contiguous), - PredicateBuilder::NotIn(0, "id", FieldType::INT, contiguous)}; - const std::vector> filtered_rows = {{1, 2, 3, 4, 5, 6}, {7}}; const std::vector all_rows = {0, 1, 2, 3, 4, 5, 6, 7}; - for (size_t predicate_idx = 0; predicate_idx < predicates.size(); ++predicate_idx) { - for (const std::string mode : {"", "auto", "force", "off"}) { - SCOPED_TRACE(fmt::format("predicate={}, mode={}", predicate_idx, mode)); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - CreatePredicateReader(path, schema, predicates[predicate_idx], mode)); - ASSERT_OK_AND_ASSIGN(std::vector rows, - CollectNativeRowIds(std::move(reader))); - ASSERT_EQ(rows, mode == "off" ? all_rows : filtered_rows[predicate_idx]); + for (bool duplicate : {false, true}) { + std::vector literals; + for (int32_t i = 0; i < 64; ++i) { + literals.emplace_back(duplicate ? int32_t{3} : i); + } + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "id", FieldType::INT, literals) + : PredicateBuilder::In(0, "id", FieldType::INT, literals); + const std::vector matches = + duplicate + ? (negate ? std::vector{1, 2, 3, 5, 6, 7} : std::vector{4}) + : (negate ? std::vector{7} : std::vector{1, 2, 3, 4, 5, 6}); + for (const std::string maximum : {"", "0", "32", "63", "64"}) { + SCOPED_TRACE( + fmt::format("duplicate={}, negate={}, maximum={}", duplicate, negate, maximum)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreatePredicateReader(path, schema, predicate, maximum)); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(std::move(reader), predicate, data, + maximum == "64" ? matches : all_rows, matches)); + } + } + } +} + +TEST_F(VortexFileFormatTest, MembershipMaximumBoundariesControlNativeRowsWithProjection) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "membership-boundaries.vortex"); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("payload", arrow::utf8())}; + const std::shared_ptr schema = arrow::schema(fields); + const auto data = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields), R"([[null,"null-id"],[0,"zero"],[100,"hundred"],[1,"one"]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/4)); + const auto projection = arrow::schema({fields[1]}); + const std::vector all_rows = {0, 1, 2, 3}; + for (const std::string maximum : {"", "1", "7", "32"}) { + const int32_t limit = maximum.empty() ? 32 : std::stoi(maximum); + for (int32_t size : {std::max(int32_t{1}, limit - 1), limit, limit + 1}) { + std::vector literals; + for (int32_t value = 0; value < size; ++value) { + literals.emplace_back(value); + } + for (bool negate : {false, true}) { + SCOPED_TRACE(fmt::format("maximum={}, size={}, negate={}", maximum, size, negate)); + const auto predicate = + negate ? PredicateBuilder::NotIn(0, "id", FieldType::INT, literals) + : PredicateBuilder::In(0, "id", FieldType::INT, literals); + const std::vector matches = + size == 1 ? (negate ? std::vector{2, 3} : std::vector{1}) + : (negate ? std::vector{2} : std::vector{1, 3}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr full_reader, + CreatePredicateReader(path, schema, predicate, maximum)); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(std::move(full_reader), predicate, data, + size <= limit ? matches : all_rows, matches)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr projected_reader, + CreatePredicateReader(path, projection, predicate, maximum)); + ASSERT_OK_AND_ASSIGN(std::vector rows, + CollectNativeRowIds(std::move(projected_reader))); + ASSERT_EQ(rows, size <= limit ? matches : all_rows); + } } } } @@ -1637,6 +2384,7 @@ TEST_F(VortexFileFormatTest, CombinedSelectionAndFilterPreservePhysicalRowsAfter PredicateBuilder::GreaterOrEqual(0, "id", FieldType::INT, Literal(static_cast(2))); ASSERT_OK( ConfigureReader(reader.get(), projection, predicate, RoaringBitmap32::From({0, 2, 3, 5}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 4, 2)); std::vector physical_rows; std::vector> batches; while (true) { @@ -1654,6 +2402,7 @@ TEST_F(VortexFileFormatTest, CombinedSelectionAndFilterPreservePhysicalRowsAfter batches.push_back(std::move(array)); } ASSERT_EQ(physical_rows, (std::vector{2, 3, 5})); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 4, 2)); auto actual = arrow::Concatenate(batches, arrow_pool_.get()).ValueOrDie(); auto expected = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(projection->fields()), R"([["two"],["three"],["five"]])") @@ -1661,6 +2410,110 @@ TEST_F(VortexFileFormatTest, CombinedSelectionAndFilterPreservePhysicalRowsAfter ASSERT_TRUE(actual->Equals(expected)); } +TEST_F(VortexFileFormatTest, BoundedStringBatchesPreserveProjectionRowIdsAndBufferOwnership) { + constexpr int32_t kRows = 40; + constexpr int32_t kBatchSize = 4; + const std::string repeated_text(4096, 't'); + std::string repeated_binary(4096, '\0'); + repeated_binary.front() = '\xff'; + repeated_binary.back() = '\x7f'; + arrow::Int32Builder id_builder; + arrow::StringBuilder text_builder; + arrow::BinaryBuilder binary_builder; + RoaringBitmap32 selection; + std::vector expected_rows; + for (int32_t row = 0; row < kRows; ++row) { + ASSERT_TRUE(id_builder.Append(row).ok()); + if (row % 4 == 2) { + ASSERT_TRUE(text_builder.AppendNull().ok()); + } else { + ASSERT_TRUE(text_builder.Append(row % 4 == 3 ? "" : repeated_text).ok()); + } + if (row % 5 == 2) { + ASSERT_TRUE(binary_builder.AppendNull().ok()); + } else { + ASSERT_TRUE(binary_builder.Append(row % 5 == 3 ? "" : repeated_binary).ok()); + } + if (row % 3 != 1) { + selection.Add(row); + if (row >= 5) { + expected_rows.push_back(row); + } + } + } + std::shared_ptr ids; + std::shared_ptr texts; + std::shared_ptr binaries; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(text_builder.Finish(&texts).ok()); + ASSERT_TRUE(binary_builder.Finish(&binaries).ok()); + const auto schema = arrow::schema({arrow::field("id", arrow::int32(), false), + arrow::field("text", arrow::utf8()), + arrow::field("payload", arrow::binary())}); + const auto data = + arrow::StructArray::Make({ids, texts, binaries}, schema->fields()).ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "bounded-strings.vortex"); + // A single write batch contains repeated values much larger than an inline string view. + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/kRows)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path, kBatchSize)); + const auto projection = arrow::schema({schema->field(2), schema->field(1)}); + const auto predicate = + PredicateBuilder::GreaterOrEqual(0, "id", FieldType::INT, Literal(static_cast(5))); + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, selection)); + std::vector retained; + size_t rows_seen = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + ASSERT_OK(paimon::test::ReadResultCollector::CheckBatchOffset(batch)); + ASSERT_GT(batch.first->length, 0); + ASSERT_LE(batch.first->length, kBatchSize); + for (int64_t i = 0; i < batch.first->length; ++i) { + ASSERT_LT(rows_seen, expected_rows.size()); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + ASSERT_EQ(physical_row, expected_rows[rows_seen++]); + } + retained.push_back(std::move(batch)); + } + ASSERT_EQ(rows_seen, expected_rows.size()); + ASSERT_GT(retained.size(), 1); + // Keep the raw Arrow exports across later batches, a new scan, and reader destruction. + ASSERT_OK(ConfigureReader(reader.get(), schema)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch restarted, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(restarted)); + reader.reset(); + auto restarted_array = + arrow::ImportArray(restarted.first.get(), restarted.second.get()).ValueOrDie(); + ASSERT_TRUE(restarted_array->ValidateFull().ok()); + + const auto& expected_texts = checked_cast(*texts); + const auto& expected_binaries = checked_cast(*binaries); + rows_seen = 0; + for (auto& batch : retained) { + auto imported = arrow::ImportArray(batch.first.get(), batch.second.get()).ValueOrDie(); + ASSERT_TRUE(imported->ValidateFull().ok()); + ASSERT_TRUE(imported->type()->Equals(arrow::struct_(projection->fields()))); + const auto& values = checked_cast(*imported); + const auto& actual_binaries = checked_cast(*values.field(0)); + const auto& actual_texts = checked_cast(*values.field(1)); + for (int64_t i = 0; i < imported->length(); ++i) { + const int64_t physical_row = static_cast(expected_rows[rows_seen++]); + ASSERT_EQ(actual_texts.IsNull(i), expected_texts.IsNull(physical_row)); + ASSERT_EQ(actual_binaries.IsNull(i), expected_binaries.IsNull(physical_row)); + if (!actual_texts.IsNull(i)) { + ASSERT_EQ(actual_texts.GetString(i), expected_texts.GetString(physical_row)); + } + if (!actual_binaries.IsNull(i)) { + ASSERT_EQ(actual_binaries.GetString(i), expected_binaries.GetString(physical_row)); + } + } + } + ASSERT_EQ(rows_seen, expected_rows.size()); +} + TEST_F(VortexFileFormatTest, ResetReadSchemaRestartsAndClearsSelectionAndFilter) { ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); @@ -1687,20 +2540,21 @@ TEST_F(VortexFileFormatTest, ResetReadSchemaRestartsAndClearsSelectionAndFilter) ASSERT_EQ(rows, (std::vector{0, 5})); } -TEST_F(VortexFileFormatTest, EmptyAndMissingProjectionsPreserveRowCountsAndPhysicalIds) { +TEST_F(VortexFileFormatTest, EmptyProjectionPreservesRowCountsAndPhysicalIds) { ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); - const std::vector> projections = { - arrow::schema({}), arrow::schema({arrow::field("missing", arrow::utf8())}), - arrow::schema({arrow::field("_row_idx", arrow::int64())}), - arrow::schema( - {arrow::field("_row_idx", arrow::int64()), arrow::field("_row_idx", arrow::int64())}), - arrow::schema( - {arrow::field("missing", arrow::struct_({arrow::field("x", arrow::int32())}))})}; - for (const auto& projection : projections) { - SCOPED_TRACE(projection->ToString()); + const auto projection = arrow::schema({}); + for (const bool filtered : {false, true}) { + SCOPED_TRACE(filtered); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); - ASSERT_OK(ConfigureReader(reader.get(), projection)); - uint64_t rows = 0; + const std::shared_ptr predicate = + filtered ? PredicateBuilder::GreaterThan(0, "id", FieldType::INT, + Literal(static_cast(2))) + : nullptr; + const std::optional selection = + filtered ? std::optional(RoaringBitmap32::From({0, 1, 3, 5})) + : std::nullopt; + ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, selection)); + std::vector rows; while (true) { ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); if (array == nullptr) { @@ -1708,76 +2562,108 @@ TEST_F(VortexFileFormatTest, EmptyAndMissingProjectionsPreserveRowCountsAndPhysi } ASSERT_TRUE(array->ValidateFull().ok()); ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); - const auto& values = checked_cast(*array); - for (const auto& child : values.fields()) { - ASSERT_EQ(child->null_count(), array->length()); - } - for (int64_t i = 0; i < array->length(); ++i) { - ASSERT_OK_AND_ASSIGN(uint64_t physical_row, - reader->GetPreviousBatchFileRowId(static_cast(i))); - ASSERT_EQ(physical_row, rows++); - } - } - ASSERT_EQ(rows, uint64_t{6}); - const auto predicate = PredicateBuilder::GreaterThan(0, "id", FieldType::INT, - Literal(static_cast(2))); - ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, - RoaringBitmap32::From({0, 1, 3, 5}))); - std::vector selected; - while (true) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); - if (array == nullptr) { - break; - } - const auto& values = checked_cast(*array); - for (const auto& child : values.fields()) { - ASSERT_EQ(child->null_count(), array->length()); - } for (int64_t i = 0; i < array->length(); ++i) { ASSERT_OK_AND_ASSIGN(uint64_t row, reader->GetPreviousBatchFileRowId(static_cast(i))); - // File readers may return candidates; the caller applies the complete predicate. - if (row > 2) { - selected.push_back(row); + // File readers may return candidates; retain the residual predicate. + if (!filtered || row > 2) { + rows.push_back(row); } } } - ASSERT_EQ(selected, (std::vector{3, 5})); + ASSERT_EQ(rows, filtered ? (std::vector{3, 5}) + : (std::vector{0, 1, 2, 3, 4, 5})); + ASSERT_FALSE(reader->GetPreviousBatchFileRowId(0).ok()); } } -TEST_F(VortexFileFormatTest, RepeatedProjectionFieldsKeepBothValues) { +TEST_F(VortexFileFormatTest, RepeatedProjectionFieldsAreRejected) { ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); const auto projection = arrow::schema( {arrow::field("id", arrow::int32(), false), arrow::field("id", arrow::int32(), false)}); - for (const bool filtered : {false, true}) { - SCOPED_TRACE(filtered); + const Status status = ConfigureReader(reader.get(), projection); + ASSERT_TRUE(status.IsInvalid()) << status.ToString(); + ASSERT_NE(status.message().find("duplicate field id"), std::string::npos); +} + +TEST_F(VortexFileFormatTest, NestedProjectionValidatesEveryRequestedField) { + const std::string path = PathUtil::JoinPath(directory_->Str(), "nested-schema-contract.vortex"); + const auto item_type = + arrow::struct_({arrow::field("a", arrow::int32()), arrow::field("b", arrow::utf8())}); + const auto schema = + arrow::schema({arrow::field("r", item_type), arrow::field("l", arrow::list(item_type)), + arrow::field("v", arrow::fixed_size_list(item_type, 2))}); + const std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(schema->fields()), + R"([[{"a":10,"b":"x"},[{"a":1,"b":"one"}],[{"a":2,"b":"two"},{"a":3,"b":null}]], + [{"a":20,"b":"y"},[{"a":4,"b":"four"}],[{"a":5,"b":"five"},{"a":6,"b":"six"}]]])") + .ValueOrDie(); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/1)); + const auto missing = arrow::struct_({arrow::field("missing", arrow::int32())}); + const auto duplicate = + arrow::struct_({arrow::field("a", arrow::int32()), arrow::field("a", arrow::int32())}); + const auto pruned = arrow::struct_({arrow::field("b", arrow::utf8())}); + const std::vector> invalid = { + arrow::schema({arrow::field("r", missing)}), arrow::schema({arrow::field("r", duplicate)}), + arrow::schema({arrow::field("l", arrow::list(missing))}), + arrow::schema({arrow::field("l", arrow::list(duplicate))}), + arrow::schema({arrow::field("v", arrow::fixed_size_list(missing, 2))}), + arrow::schema({arrow::field("v", arrow::fixed_size_list(duplicate, 2))}), + // An earlier valid sub-field projection must not hide later invalid fields. + arrow::schema({arrow::field("r", pruned), arrow::field("l", arrow::list(missing))})}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + for (const auto& projection : invalid) { + SCOPED_TRACE(projection->ToString()); + const Status status = ConfigureReader(reader.get(), projection); + ASSERT_TRUE(status.IsInvalid()) << status.ToString(); + } + + // The list wrapper name is not a physical STRUCT field and may differ from the file. + const auto projection = arrow::schema( + {arrow::field("l", arrow::list(arrow::field("element", pruned))), + arrow::field("v", arrow::fixed_size_list(arrow::field("element", pruned), 2)), + arrow::field("r", pruned)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadFile(path, projection, /*batch_size=*/1, /*expected_row_count=*/2)); + const std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(projection->fields()), + R"([[[{"b":"one"}],[{"b":"two"},{"b":null}],{"b":"x"}], + [[{"b":"four"}],[{"b":"five"},{"b":"six"}],{"b":"y"}]])") + .ValueOrDie(); + ASSERT_TRUE(actual->Equals(expected)) << actual->ToString(); +} + +TEST_F(VortexFileFormatTest, InvalidReadSchemaResetPreservesFilterSelectionAndPhysicalRows) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + const auto id = arrow::field("id", arrow::int32()); + const std::vector> invalid = { + arrow::schema({id, id}), arrow::schema({arrow::field("missing", arrow::utf8())}), + arrow::schema({arrow::field("_row_idx", arrow::uint64())})}; + for (const auto& projection : invalid) { + SCOPED_TRACE(projection->ToString()); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); - const std::shared_ptr predicate = - filtered ? PredicateBuilder::GreaterThan(0, "id", FieldType::INT, - Literal(static_cast(2))) - : nullptr; - ASSERT_OK(ConfigureReader(reader.get(), projection, predicate)); - uint64_t rows = 0; - while (true) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); - if (array == nullptr) { - break; - } - ASSERT_TRUE(array->ValidateFull().ok()); - ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); - const auto& values = checked_cast(*array); - ASSERT_TRUE(values.field(0)->Equals(values.field(1))); - rows += array->length(); - } - ASSERT_EQ(rows, filtered ? uint64_t{3} : uint64_t{6}); - if (filtered) { - ASSERT_OK(ConfigureReader(reader.get(), projection, predicate, - RoaringBitmap32::From({1, 3, 5}))); - ASSERT_OK_AND_ASSIGN(std::vector selected, - CollectNativeRowIds(std::move(reader))); - ASSERT_EQ(selected, (std::vector{3, 5})); - } + const auto predicate = PredicateBuilder::GreaterThan(0, "id", FieldType::INT, + Literal(static_cast(0))); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), predicate, + RoaringBitmap32::From({0, 1, 3, 4, 5}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 5, 1)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); + ASSERT_NE(first, nullptr); + ASSERT_EQ(first->length(), 2); + const Status status = + ConfigureReader(reader.get(), projection, nullptr, RoaringBitmap32::From({0})); + ASSERT_TRUE(status.IsInvalid()) << status.ToString(); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 5, 1)); + ASSERT_OK_AND_ASSIGN(uint64_t row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row, uint64_t{1}); + ASSERT_OK_AND_ASSIGN(row, reader->GetPreviousBatchFileRowId(1)); + ASSERT_EQ(row, uint64_t{3}); + ASSERT_OK_AND_ASSIGN(std::vector remaining, + CollectNativeRowIds(std::move(reader))); + ASSERT_EQ(remaining, (std::vector{4, 5})); } } @@ -1789,32 +2675,46 @@ TEST_F(VortexFileFormatTest, FilterDoesNotOverwriteAFileColumnNamedRowIdx) { .ValueOrDie(); const std::string path = PathUtil::JoinPath(directory_->Str(), "row-idx-column.vortex"); ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); const auto predicate = PredicateBuilder::Equal(0, "id", FieldType::INT, Literal(static_cast(3))); - ASSERT_OK(ConfigureReader(reader.get(), schema, predicate, RoaringBitmap32::From({1, 3}))); const auto filter = std::dynamic_pointer_cast(predicate); ASSERT_NE(filter, nullptr); - std::vector matched_rows; - while (true) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); - if (array == nullptr) { - break; - } - ASSERT_TRUE(array->ValidateFull().ok()); - const auto& values = checked_cast(*array); - const auto& original_column = checked_cast(*values.field(1)); - ASSERT_OK_AND_ASSIGN(std::vector mask, filter->Test(*array, arrow_pool_.get())); - for (int64_t i = 0; i < array->length(); ++i) { - ASSERT_OK_AND_ASSIGN(uint64_t physical_row, - reader->GetPreviousBatchFileRowId(static_cast(i))); - ASSERT_EQ(original_column.Value(i), 100 + static_cast(physical_row)); - if (mask[i] != 0) { - matched_rows.push_back(physical_row); + for (const auto& read_schema : {schema, arrow::schema({schema->field(0)})}) { + SCOPED_TRACE(read_schema->ToString()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + ASSERT_OK( + ConfigureReader(reader.get(), read_schema, predicate, RoaringBitmap32::From({1, 3}))); + std::vector scanned_rows; + std::vector matched_rows; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, ReadNextBatch(reader.get())); + if (array == nullptr) { + break; + } + ASSERT_TRUE(array->ValidateFull().ok()); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(read_schema->fields()))); + const auto& values = checked_cast(*array); + const auto& ids = checked_cast(*values.field(0)); + ASSERT_OK_AND_ASSIGN(std::vector mask, filter->Test(*array, arrow_pool_.get())); + for (int64_t i = 0; i < array->length(); ++i) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + scanned_rows.push_back(physical_row); + ASSERT_EQ(ids.Value(i), static_cast(physical_row)); + if (read_schema->num_fields() == 2) { + const auto& original_column = + checked_cast(*values.field(1)); + ASSERT_EQ(original_column.Value(i), 100 + static_cast(physical_row)); + } + if (mask[i] != 0) { + matched_rows.push_back(physical_row); + } } } + // The stored _row_idx prevents native filtering even when it is not projected. + ASSERT_EQ(scanned_rows, (std::vector{1, 3})); + ASSERT_EQ(matched_rows, (std::vector{3})); } - ASSERT_EQ(matched_rows, (std::vector{3})); } TEST_F(VortexFileFormatTest, InvalidSelectionResetDoesNotChangeTheActiveReader) { @@ -1822,10 +2722,12 @@ TEST_F(VortexFileFormatTest, InvalidSelectionResetDoesNotChangeTheActiveReader) ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({0, 2, 4}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 3, 3)); ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); ASSERT_EQ(first->length(), 2); ASSERT_NOK( ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({1, 6}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 3, 3)); ASSERT_OK_AND_ASSIGN(uint64_t physical_row, reader->GetPreviousBatchFileRowId(0)); ASSERT_EQ(physical_row, uint64_t{0}); ASSERT_OK_AND_ASSIGN(physical_row, reader->GetPreviousBatchFileRowId(1)); @@ -1839,10 +2741,12 @@ TEST_F(VortexFileFormatTest, InvalidArrowSchemaResetDoesNotChangeTheActiveReader ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), nullptr, RoaringBitmap32::From({0, 2, 4}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 3, 3)); ASSERT_OK_AND_ASSIGN(std::shared_ptr first, ReadNextBatch(reader.get())); ASSERT_EQ(first->length(), 2); ::ArrowSchema invalid_schema{}; ASSERT_NOK(reader->SetReadSchema(&invalid_schema, nullptr, RoaringBitmap32::From({1, 3}))); + ASSERT_NO_FATAL_FAILURE(AssertSelectionMetrics(reader.get(), 3, 3)); ASSERT_OK_AND_ASSIGN(uint64_t physical_row, reader->GetPreviousBatchFileRowId(0)); ASSERT_EQ(physical_row, uint64_t{0}); ASSERT_OK_AND_ASSIGN(std::vector remaining, CollectNativeRowIds(std::move(reader))); diff --git a/src/paimon/format/vortex/vortex_format_writer.cpp b/src/paimon/format/vortex/vortex_format_writer.cpp index ffa9be256..8d1e972f3 100644 --- a/src/paimon/format/vortex/vortex_format_writer.cpp +++ b/src/paimon/format/vortex/vortex_format_writer.cpp @@ -44,6 +44,42 @@ namespace { constexpr double COMPRESSION_RATIO = 0.25; +// Match Java's UTC LTZ storage without changing the epoch values or timestamp units. +std::shared_ptr NormalizeTimestampTimezones( + const std::shared_ptr& type) { + if (type->id() == arrow::Type::TIMESTAMP) { + const auto& timestamp = checked_cast(*type); + return timestamp.timezone().empty() || timestamp.timezone() == "UTC" + ? type + : arrow::timestamp(timestamp.unit(), "UTC"); + } + if (type->id() != arrow::Type::STRUCT && type->id() != arrow::Type::LIST && + type->id() != arrow::Type::FIXED_SIZE_LIST) { + return type; + } + arrow::FieldVector fields = type->fields(); + bool changed = false; + for (auto& field : fields) { + std::shared_ptr normalized = NormalizeTimestampTimezones(field->type()); + if (normalized != field->type()) { + field = field->WithType(normalized); + changed = true; + } + } + if (!changed) { + return type; + } + switch (type->id()) { + case arrow::Type::STRUCT: + return arrow::struct_(fields); + case arrow::Type::LIST: + return arrow::list(fields[0]); + default: + return arrow::fixed_size_list( + fields[0], checked_cast(*type).list_size()); + } +} + Status ValidateVisibleVectorElements(const std::shared_ptr& array, arrow::MemoryPool* pool) { if (!VectorUtils::ContainsVectorType(array->type())) { @@ -129,12 +165,18 @@ Result> VortexFormatWriter::Create( return Status::Invalid("Vortex writer requires non-null output, schema and arrow pool"); } PAIMON_RETURN_NOT_OK(SchemaValidation::ValidateVortexArrowSchema(schema)); + arrow::FieldVector fields = schema->fields(); + for (auto& field : fields) { + field = field->WithType(NormalizeTimestampTimezones(field->type())); + } + const auto write_schema = + std::make_shared(fields, schema->endianness(), schema->metadata()); VxSessionPtr session(vx_session_new(), vx_session_free); if (session == nullptr) { return Status::IOError("failed to create Vortex session"); } ::ArrowSchema ffi_schema = {}; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*write_schema, &ffi_schema)); vx_error* error = nullptr; // vx_dtype_from_arrow_schema consumes ffi_schema on both success and failure. VxDtypePtr dtype(vx_dtype_from_arrow_schema(&ffi_schema, &error), vx_dtype_free); @@ -150,7 +192,7 @@ Result> VortexFormatWriter::Create( output_context->GetCallbackStatus()); } return std::unique_ptr(new VortexFormatWriter( - output, schema, std::move(session), std::move(output_context), sink, arrow_pool)); + output, write_schema, std::move(session), std::move(output_context), sink, arrow_pool)); } VortexFormatWriter::~VortexFormatWriter() { diff --git a/src/paimon/format/vortex/vortex_io_callbacks.h b/src/paimon/format/vortex/vortex_io_callbacks.h index afa46fb49..fdeed239c 100644 --- a/src/paimon/format/vortex/vortex_io_callbacks.h +++ b/src/paimon/format/vortex/vortex_io_callbacks.h @@ -36,10 +36,6 @@ class OutputStream; namespace paimon::vortex { -/// I/O coalescing policy. Low-latency sources must explicitly opt into LOCAL; -/// the reader does not infer filesystem semantics from InputStream URIs. -enum class VortexIoCoalesceMode { OBJECT_STORAGE, LOCAL }; - /// Bridges Vortex's positional reads to a paimon `InputStream`. class VortexInputContext { public: @@ -56,7 +52,7 @@ class VortexInputContext { /// reader has dropped its own reference. /// /// @param context The context to forward callbacks to; a shared reference is retained. - /// @return The callback struct to hand to `vx_data_source_new_callback`. + /// @return The callback struct for `vx_data_source_new_callback_with_coalesce`. static vx_input_callbacks MakeCallbacks(const std::shared_ptr& context); /// `vx_read_at_fn`: fills `length` bytes at `offset` from the paimon stream. diff --git a/src/paimon/format/vortex/vortex_predicate_converter.cpp b/src/paimon/format/vortex/vortex_predicate_converter.cpp index b74bc2f88..afb54d255 100644 --- a/src/paimon/format/vortex/vortex_predicate_converter.cpp +++ b/src/paimon/format/vortex/vortex_predicate_converter.cpp @@ -18,19 +18,20 @@ #include "paimon/format/vortex/vortex_predicate_converter.h" -#include -#include -#include #include #include #include +#include #include -#include #include #include #include "arrow/type.h" +#include "arrow/util/decimal.h" #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/core/casting/decimal_to_decimal_cast_executor.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" @@ -47,16 +48,21 @@ VxExpressionPtr NullExpr() { return VxExpressionPtr(nullptr, vx_expression_free); } -void ClearError(vx_error* error) { - if (error != nullptr) { - vx_error_free(error); +Status CheckNativeResult(const std::string& operation, const void* handle, + vx_error* error = nullptr) { + PAIMON_RETURN_NOT_OK(VortexFfiError(operation, error)); + if (handle == nullptr) { + return Status::Invalid(operation, ": Vortex returned a null handle"); } + return Status::OK(); } // vx_expression_not returns an owned handle typed const; reclaim it as an owned non-const pointer. -VxExpressionPtr MakeNotExpr(const vx_expression* child) { - return VxExpressionPtr(const_cast(vx_expression_not(child)), - vx_expression_free); +Result MakeNotExpr(const vx_expression* child) { + VxExpressionPtr expression(const_cast(vx_expression_not(child)), + vx_expression_free); + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex NOT expression", expression.get())); + return expression; } // Reject precision loss and overflow before constructing a native timestamp scalar. @@ -65,8 +71,8 @@ bool GetTimestampValue(const Literal& literal, const std::shared_ptrid() != arrow::Type::TIMESTAMP) { return false; } - const auto& timestamp_type = checked_cast(*arrow_type); - if (!timestamp_type.timezone().empty() && timestamp_type.timezone() != "UTC") { + const auto timestamp_type = checked_pointer_cast(arrow_type); + if (!timestamp_type->timezone().empty() && timestamp_type->timezone() != "UTC") { return false; } const auto timestamp = literal.GetValue(); @@ -75,12 +81,11 @@ bool GetTimestampValue(const Literal& literal, const std::shared_ptrunit()) { case arrow::TimeUnit::SECOND: if (value % 1000 != 0 || nanos != 0) { return false; } - value /= 1000; unit = VX_TIME_UNIT_SECOND; break; case arrow::TimeUnit::MILLI: @@ -101,6 +106,8 @@ bool GetTimestampValue(const Literal& literal, const std::shared_ptr::min() / factor || value > std::numeric_limits::max() / factor) { @@ -110,13 +117,46 @@ bool GetTimestampValue(const Literal& literal, const std::shared_ptr std::numeric_limits::max() - fraction) { return false; } - *result = value + fraction; + *result = DateTimeUtils::TimestampToInteger( + timestamp, DateTimeUtils::GetTimeTypeFromArrowType(timestamp_type)); *result_unit = unit; return true; } -VxScalarPtr MakeLiteralScalar(const Literal& literal, - const std::shared_ptr& arrow_type) { +// Normalize constants only; shared CAST may round, so accept it only if the value is unchanged. +Result> NormalizeDecimalLiteral( + const Literal& literal, const std::shared_ptr& arrow_type) { + const auto decimal = literal.GetValue(); + const auto& target_type = checked_cast(*arrow_type); + // Keep literals within Paimon's decimal range and the shared scale table. + if (decimal.Precision() < Decimal::MIN_PRECISION || + decimal.Precision() > Decimal::MAX_PRECISION || decimal.Scale() < 0 || + decimal.Scale() > decimal.Precision() || target_type.precision() < Decimal::MIN_PRECISION || + target_type.precision() > Decimal::MAX_PRECISION || target_type.scale() < 0 || + target_type.scale() > target_type.precision()) { + return std::optional(); + } + if (!arrow::Decimal128(decimal.HighBits(), decimal.LowBits()) + .FitsInPrecision(decimal.Precision())) { + return std::optional(); + } + if (decimal.Precision() == target_type.precision() && decimal.Scale() == target_type.scale()) { + return std::make_optional(decimal); + } + const DecimalToDecimalCastExecutor cast_executor; + PAIMON_ASSIGN_OR_RAISE(Literal converted, cast_executor.Cast(literal, arrow_type)); + if (converted.IsNull()) { + return std::optional(); + } + const auto normalized = converted.GetValue(); + if (decimal.CompareTo(normalized) != 0) { + return std::optional(); + } + return std::make_optional(normalized); +} + +Result MakeLiteralScalar(const Literal& literal, + const std::shared_ptr& arrow_type) { VxScalarPtr scalar(nullptr, vx_scalar_free); if (literal.IsNull()) { return scalar; @@ -139,6 +179,10 @@ VxScalarPtr MakeLiteralScalar(const Literal& literal, // plain int scalar, so build a date-typed scalar. case FieldType::DATE: scalar.reset(vx_scalar_new_date(literal.GetValue(), false, &error)); + if (scalar == nullptr && error == nullptr) { + // A valid Paimon constant can exceed the native civil-date range. + return scalar; + } break; case FieldType::BIGINT: scalar.reset(vx_scalar_new_i64(literal.GetValue(), false)); @@ -153,7 +197,13 @@ VxScalarPtr MakeLiteralScalar(const Literal& literal, break; } case FieldType::DECIMAL: { - const auto decimal = literal.GetValue(); + PAIMON_ASSIGN_OR_RAISE(std::optional normalized, + NormalizeDecimalLiteral(literal, arrow_type)); + if (!normalized.has_value()) { + return scalar; + } + // Cast validated the full source value before selecting a narrower representation. + const Decimal& decimal = normalized.value(); const auto precision = static_cast(decimal.Precision()); const auto scale = static_cast(decimal.Scale()); if (decimal.IsCompact()) { @@ -180,61 +230,69 @@ VxScalarPtr MakeLiteralScalar(const Literal& literal, const auto& timestamp_type = checked_cast(*arrow_type); const char* zone = timestamp_type.timezone().empty() ? nullptr : "UTC"; scalar.reset(vx_scalar_new_timestamp(value, unit, zone, false, &error)); + if (scalar == nullptr && error == nullptr) { + // Preserve residual comparison when native temporal limits are narrower. + return scalar; + } break; } default: return scalar; } - ClearError(error); + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex predicate scalar", scalar.get(), error)); return scalar; } -VxExpressionPtr MakeLiteralExpr(const Literal& literal, - const std::shared_ptr& arrow_type) { - VxScalarPtr scalar = MakeLiteralScalar(literal, arrow_type); +Result MakeLiteralExpr(const Literal& literal, + const std::shared_ptr& arrow_type) { + PAIMON_ASSIGN_OR_RAISE(VxScalarPtr scalar, MakeLiteralScalar(literal, arrow_type)); if (scalar == nullptr) { return NullExpr(); } - vx_error* literal_error = nullptr; - vx_expression* expr = vx_expression_literal(scalar.get(), &literal_error); - ClearError(literal_error); - return VxExpressionPtr(expr, vx_expression_free); + vx_error* error = nullptr; + VxExpressionPtr expression(vx_expression_literal(scalar.get(), &error), vx_expression_free); + PAIMON_RETURN_NOT_OK( + CheckNativeResult("create Vortex literal expression", expression.get(), error)); + return expression; } -// Reject implicit coercions that could narrow values before the schema-evolution/residual path. -bool LiteralMatchesColumn(const Literal& literal, FieldType field_type, - const std::shared_ptr& arrow_type) { - if (literal.IsNull() || literal.GetType() != field_type || arrow_type == nullptr) { +// Type mismatches violate the predicate contract; valid but unsafe conversions fall back. +Result LiteralMatchesColumn(const Literal& literal, + const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type) { + if (literal.IsNull() || arrow_type == nullptr) { + return false; + } + const FieldType field_type = leaf->GetFieldType(); + if (literal.GetType() != field_type) { + return Status::Invalid("Literal type mismatch for field '", leaf->FieldName(), + "': expected ", FieldTypeUtils::FieldTypeToString(field_type), + ", actual ", FieldTypeUtils::FieldTypeToString(literal.GetType())); + } + const Result file_type = FieldTypeUtils::ConvertToFieldType(arrow_type->id()); + if (!file_type.ok()) { + // Arrow types outside the Paimon type mapping cannot be pushed. return false; } + if (file_type.value() != field_type) { + return Status::Invalid("Predicate type mismatch for field '", leaf->FieldName(), + "': expected ", FieldTypeUtils::FieldTypeToString(file_type.value()), + " from file type ", arrow_type->ToString(), ", actual ", + FieldTypeUtils::FieldTypeToString(field_type)); + } switch (field_type) { case FieldType::BOOLEAN: - return arrow_type->id() == arrow::Type::BOOL; case FieldType::TINYINT: - return arrow_type->id() == arrow::Type::INT8; case FieldType::SMALLINT: - return arrow_type->id() == arrow::Type::INT16; case FieldType::INT: - return arrow_type->id() == arrow::Type::INT32; case FieldType::BIGINT: - return arrow_type->id() == arrow::Type::INT64; case FieldType::STRING: - return arrow_type->id() == arrow::Type::STRING; case FieldType::DATE: - return arrow_type->id() == arrow::Type::DATE32; - case FieldType::DECIMAL: { - if (arrow_type->id() != arrow::Type::DECIMAL128) { - return false; - } - const auto& decimal_type = checked_cast(*arrow_type); - const auto decimal = literal.GetValue(); - return decimal.Precision() == decimal_type.precision() && - decimal.Scale() == decimal_type.scale(); - } + return true; + case FieldType::DECIMAL: + // Parameter normalization and lossless-value checks run once when building the scalar. + return true; case FieldType::TIMESTAMP: { - if (arrow_type->id() != arrow::Type::TIMESTAMP) { - return false; - } const auto& timestamp_type = checked_cast(*arrow_type); return timestamp_type.timezone().empty() || timestamp_type.timezone() == "UTC"; } @@ -244,58 +302,29 @@ bool LiteralMatchesColumn(const Literal& literal, FieldType field_type, } } -struct LiteralPointerHash { - size_t operator()(const Literal* literal) const { - return literal->HashCode(); - } -}; - -struct LiteralPointerEqual { - bool operator()(const Literal* lhs, const Literal* rhs) const { - return *lhs == *rhs; - } -}; - -// All literals must be validated before their values are inspected or used as hash keys. Decimal -// scale/precision is identical to the column, so Literal's same-scale hash equality is safe here. -bool ValidateAndDeduplicate(const std::shared_ptr& leaf, - const std::shared_ptr& arrow_type, - std::vector* literals) { - if (leaf->Literals().empty()) { - return false; - } - for (const Literal& literal : leaf->Literals()) { - if (!LiteralMatchesColumn(literal, leaf->GetFieldType(), arrow_type)) { - return false; - } - if (leaf->GetFieldType() == FieldType::TIMESTAMP) { - int64_t value = 0; - vx_time_unit unit = VX_TIME_UNIT_MILLISECOND; - if (!GetTimestampValue(literal, arrow_type, &value, &unit)) { - return false; - } - } - } - std::unordered_set unique; - unique.reserve(leaf->Literals().size()); - literals->reserve(leaf->Literals().size()); +// Check every basic type even when a value or the list budget requires fallback. +Result ValidateMembershipLiteralTypes(const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type) { + bool valid = true; for (const Literal& literal : leaf->Literals()) { - if (unique.insert(&literal).second) { - literals->push_back(&literal); + PAIMON_ASSIGN_OR_RAISE(bool matches, LiteralMatchesColumn(literal, leaf, arrow_type)); + if (!matches) { + valid = false; + continue; } } - return true; + return valid; } -VxExpressionPtr ConvertInList(const std::vector& literals, - const std::shared_ptr& arrow_type, - const vx_expression* column) { +Result ConvertInList(const std::vector& literals, + const std::shared_ptr& arrow_type, + const vx_expression* column) { std::vector scalars; std::vector raw; scalars.reserve(literals.size()); raw.reserve(literals.size()); - for (const Literal* literal : literals) { - VxScalarPtr scalar = MakeLiteralScalar(*literal, arrow_type); + for (const Literal& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(VxScalarPtr scalar, MakeLiteralScalar(literal, arrow_type)); if (scalar == nullptr) { return NullExpr(); } @@ -306,176 +335,52 @@ VxExpressionPtr ConvertInList(const std::vector& literals, VxScalarPtr list( vx_scalar_new_list(vx_scalar_dtype(raw[0]), raw.data(), raw.size(), false, &error), vx_scalar_free); - ClearError(error); - if (list == nullptr) { - return NullExpr(); - } + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex list scalar", list.get(), error)); error = nullptr; VxExpressionPtr literal_expr(vx_expression_literal(list.get(), &error), vx_expression_free); - ClearError(error); - if (literal_expr == nullptr) { - return NullExpr(); - } - return VxExpressionPtr(vx_expression_list_contains(literal_expr.get(), column), - vx_expression_free); -} - -bool IsSignedInteger(FieldType field_type) { - return field_type == FieldType::TINYINT || field_type == FieldType::SMALLINT || - field_type == FieldType::INT || field_type == FieldType::BIGINT; -} - -// Only called for non-null signed integers whose type has already been validated. -int64_t SignedIntegerValue(const Literal& literal) { - switch (literal.GetType()) { - case FieldType::TINYINT: - return literal.GetValue(); - case FieldType::SMALLINT: - return literal.GetValue(); - case FieldType::INT: - return literal.GetValue(); - case FieldType::BIGINT: - return literal.GetValue(); - default: - assert(false); - return 0; - } + PAIMON_RETURN_NOT_OK( + CheckNativeResult("create Vortex literal expression", literal_expr.get(), error)); + VxExpressionPtr expression(vx_expression_list_contains(literal_expr.get(), column), + vx_expression_free); + PAIMON_RETURN_NOT_OK( + CheckNativeResult("create Vortex list membership expression", expression.get())); + return expression; } -using IntegerRange = std::pair; - -std::vector PlanIntegerRanges(std::vector literals) { - std::sort(literals.begin(), literals.end(), [](const Literal* lhs, const Literal* rhs) { - return SignedIntegerValue(*lhs) < SignedIntegerValue(*rhs); - }); - std::vector ranges; - for (const Literal* literal : literals) { - const int64_t value = SignedIntegerValue(*literal); - const int64_t previous = ranges.empty() ? 0 : SignedIntegerValue(*ranges.back().second); - // Guard before addition, including when the stored type is BIGINT at INT64_MAX. - if (!ranges.empty() && previous != std::numeric_limits::max() && - value == previous + 1) { - ranges.back().second = literal; - } else { - ranges.emplace_back(literal, literal); - } +Result ConvertMembership(const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type, + const vx_expression* column, bool negate, + const VortexPredicateOptions& options) { + const std::vector& literals = leaf->Literals(); + PAIMON_ASSIGN_OR_RAISE(bool valid, ValidateMembershipLiteralTypes(leaf, arrow_type)); + // The query engine owns deduplication and integer-range rewrites. Apply value conversions + // and create native scalars only within the original-list limit, after type validation. + if (!valid || literals.empty() || literals.size() > options.max_in_values) { + return NullExpr(); } - return ranges; -} - -// Exclude input/literal nodes. Each range costs two comparisons and one connective; R ranges -// need R-1 additional connectives. -uint64_t IntegerRangeNodes(const std::vector& ranges) { - uint64_t nodes = ranges.size() - 1; - for (const auto& [lower, upper] : ranges) { - nodes += lower == upper ? 1 : 3; - } - return nodes; -} - -VxExpressionPtr ConvertIntegerRanges(const std::vector& ranges, - const std::shared_ptr& arrow_type, - const vx_expression* column, bool negate) { - std::vector expressions; - expressions.reserve(ranges.size()); - for (const auto& [lower, upper] : ranges) { - VxExpressionPtr lower_literal = MakeLiteralExpr(*lower, arrow_type); - if (lower_literal == nullptr) { - return NullExpr(); - } - if (lower == upper) { - VxExpressionPtr equal(vx_expression_binary(negate ? VX_OPERATOR_NOT_EQ : VX_OPERATOR_EQ, - column, lower_literal.get()), - vx_expression_free); - if (equal == nullptr) { + if (leaf->GetFieldType() == FieldType::TIMESTAMP) { + // Preserve whole-list fallback for unsafe timestamp values before native construction. + for (const Literal& literal : literals) { + int64_t value = 0; + vx_time_unit unit = VX_TIME_UNIT_MILLISECOND; + if (!GetTimestampValue(literal, arrow_type, &value, &unit)) { return NullExpr(); } - expressions.push_back(std::move(equal)); - continue; - } - VxExpressionPtr upper_literal = MakeLiteralExpr(*upper, arrow_type); - if (upper_literal == nullptr) { - return NullExpr(); - } - VxExpressionPtr lower_bound(vx_expression_binary(negate ? VX_OPERATOR_LT : VX_OPERATOR_GTE, - column, lower_literal.get()), - vx_expression_free); - VxExpressionPtr upper_bound(vx_expression_binary(negate ? VX_OPERATOR_GT : VX_OPERATOR_LTE, - column, upper_literal.get()), - vx_expression_free); - if (lower_bound == nullptr || upper_bound == nullptr) { - return NullExpr(); - } - const vx_expression* bounds[] = {lower_bound.get(), upper_bound.get()}; - VxExpressionPtr expression( - negate ? vx_expression_or(bounds, 2) : vx_expression_and(bounds, 2), - vx_expression_free); - if (expression == nullptr) { - return NullExpr(); } - expressions.push_back(std::move(expression)); - } - if (expressions.size() == 1) { - return std::move(expressions.front()); - } - std::vector raw; - raw.reserve(expressions.size()); - for (const VxExpressionPtr& expression : expressions) { - raw.push_back(expression.get()); - } - // IN is the union of closed ranges; NOT IN intersects their open complements. - return VxExpressionPtr(negate ? vx_expression_and(raw.data(), raw.size()) - : vx_expression_or(raw.data(), raw.size()), - vx_expression_free); -} - -VxExpressionPtr ConvertMembership(const std::shared_ptr& leaf, - const std::shared_ptr& arrow_type, - const vx_expression* column, bool negate, - const VortexPredicateOptions& options) { - if (options.membership_pushdown == MembershipPushdownMode::OFF) { - return NullExpr(); - } - std::vector literals; - if (!ValidateAndDeduplicate(leaf, arrow_type, &literals)) { - return NullExpr(); } - std::vector ranges; - if (IsSignedInteger(leaf->GetFieldType())) { - ranges = PlanIntegerRanges(literals); - } - // ListContains lowers N values to N equality nodes and N-1 OR nodes. NOT IN adds NOT plus - // the IS NOT NULL guard (IS NULL, NOT, AND); a complement-of-ranges needs only the guard. - const uint64_t list_nodes = 2 * static_cast(literals.size()) - 1 + (negate ? 4 : 0); - const uint64_t range_nodes = - ranges.empty() ? list_nodes : IntegerRangeNodes(ranges) + (negate ? 3 : 0); - const bool use_ranges = !ranges.empty() && range_nodes < list_nodes; - const uint64_t nodes = use_ranges ? range_nodes : list_nodes; - if (options.membership_pushdown == MembershipPushdownMode::AUTO && - nodes > options.max_membership_nodes) { - return NullExpr(); - } - // Allocate native scalars only after the plan passes the budget; ranges need only endpoints. - VxExpressionPtr expression = use_ranges - ? ConvertIntegerRanges(ranges, arrow_type, column, negate) - : ConvertInList(literals, arrow_type, column); + PAIMON_ASSIGN_OR_RAISE(VxExpressionPtr expression, ConvertInList(literals, arrow_type, column)); if (expression == nullptr || !negate) { return expression; } - if (!use_ranges) { - expression = MakeNotExpr(expression.get()); - } + PAIMON_ASSIGN_OR_RAISE(expression, MakeNotExpr(expression.get())); VxExpressionPtr is_null(vx_expression_is_null(column), vx_expression_free); - if (expression == nullptr || is_null == nullptr) { - return NullExpr(); - } - VxExpressionPtr is_not_null = MakeNotExpr(is_null.get()); - if (is_not_null == nullptr) { - return NullExpr(); - } + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex IS NULL expression", is_null.get())); + PAIMON_ASSIGN_OR_RAISE(VxExpressionPtr is_not_null, MakeNotExpr(is_null.get())); // Neither IN nor NOT IN may retain null rows. Validation also rejects null/empty lists. const vx_expression* children[] = {is_not_null.get(), expression.get()}; - return VxExpressionPtr(vx_expression_and(children, 2), vx_expression_free); + VxExpressionPtr combined(vx_expression_and(children, 2), vx_expression_free); + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex AND expression", combined.get())); + return combined; } bool ComparisonOperator(Function::Type type, vx_binary_operator* op) { @@ -503,9 +408,17 @@ bool ComparisonOperator(Function::Type type, vx_binary_operator* op) { } } -VxExpressionPtr ConvertLeaf(const std::shared_ptr& leaf, - const std::shared_ptr& file_schema, - const vx_expression* root, const VortexPredicateOptions& options) { +Result ConvertLeaf(const std::shared_ptr& leaf, + const std::shared_ptr& file_schema, + const vx_expression* root, + const VortexPredicateOptions& options) { + const Function::Type type = leaf->GetFunction().GetType(); + vx_binary_operator op = VX_OPERATOR_EQ; + const bool is_comparison = ComparisonOperator(type, &op); + if (is_comparison && leaf->Literals().size() != 1) { + return Status::Invalid("Comparison predicate for field '", leaf->FieldName(), + "' requires exactly one literal"); + } const int32_t field_index = file_schema != nullptr ? file_schema->GetFieldIndex(leaf->FieldName()) : -1; if (field_index < 0) { @@ -513,19 +426,14 @@ VxExpressionPtr ConvertLeaf(const std::shared_ptr& leaf, return NullExpr(); } const std::shared_ptr& arrow_type = file_schema->field(field_index)->type(); - const Function::Type type = leaf->GetFunction().GetType(); VxExpressionPtr column(vx_expression_get_item(leaf->FieldName().c_str(), root), vx_expression_free); - if (column == nullptr) { - return NullExpr(); - } - if (type == Function::Type::IS_NULL) { - return VxExpressionPtr(vx_expression_is_null(column.get()), vx_expression_free); - } - if (type == Function::Type::IS_NOT_NULL) { + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex field expression", column.get())); + if (type == Function::Type::IS_NULL || type == Function::Type::IS_NOT_NULL) { VxExpressionPtr is_null(vx_expression_is_null(column.get()), vx_expression_free); - if (is_null == nullptr) { - return NullExpr(); + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex IS NULL expression", is_null.get())); + if (type == Function::Type::IS_NULL) { + return is_null; } return MakeNotExpr(is_null.get()); } @@ -535,25 +443,30 @@ VxExpressionPtr ConvertLeaf(const std::shared_ptr& leaf, if (type == Function::Type::NOT_IN) { return ConvertMembership(leaf, arrow_type, column.get(), true, options); } - vx_binary_operator op = VX_OPERATOR_EQ; - if (!ComparisonOperator(type, &op)) { + if (!is_comparison) { return NullExpr(); } - if (leaf->Literals().empty() || - !LiteralMatchesColumn(leaf->Literals()[0], leaf->GetFieldType(), arrow_type)) { + PAIMON_ASSIGN_OR_RAISE(bool matches, + LiteralMatchesColumn(leaf->Literals()[0], leaf, arrow_type)); + if (!matches) { return NullExpr(); } - VxExpressionPtr literal_expr = MakeLiteralExpr(leaf->Literals()[0], arrow_type); + PAIMON_ASSIGN_OR_RAISE(VxExpressionPtr literal_expr, + MakeLiteralExpr(leaf->Literals()[0], arrow_type)); if (literal_expr == nullptr) { return NullExpr(); } - return VxExpressionPtr(vx_expression_binary(op, column.get(), literal_expr.get()), - vx_expression_free); + VxExpressionPtr expression(vx_expression_binary(op, column.get(), literal_expr.get()), + vx_expression_free); + PAIMON_RETURN_NOT_OK( + CheckNativeResult("create Vortex comparison expression", expression.get())); + return expression; } -VxExpressionPtr ConvertImpl(const std::shared_ptr& predicate, - const std::shared_ptr& file_schema, - const vx_expression* root, const VortexPredicateOptions& options) { +Result ConvertImpl(const std::shared_ptr& predicate, + const std::shared_ptr& file_schema, + const vx_expression* root, + const VortexPredicateOptions& options) { if (predicate == nullptr) { return NullExpr(); } @@ -561,21 +474,29 @@ VxExpressionPtr ConvertImpl(const std::shared_ptr& predicate, if (type == Function::Type::AND || type == Function::Type::OR) { const auto compound = std::dynamic_pointer_cast(predicate); if (compound == nullptr) { - return NullExpr(); + return Status::Invalid("AND/OR predicate must be a CompoundPredicate"); } + if (compound->Children().empty()) { + return Status::Invalid("AND/OR predicate must have at least one child"); + } + bool all_children_pushed = true; std::vector children; children.reserve(compound->Children().size()); for (const std::shared_ptr& child : compound->Children()) { - VxExpressionPtr converted = ConvertImpl(child, file_schema, root, options); + if (child == nullptr) { + return Status::Invalid("AND/OR predicate must not contain a null child"); + } + PAIMON_ASSIGN_OR_RAISE(VxExpressionPtr converted, + ConvertImpl(child, file_schema, root, options)); if (converted != nullptr) { children.push_back(std::move(converted)); - } else if (type == Function::Type::OR) { - // Dropping an OR branch could exclude matching rows. - return NullExpr(); + } else { + all_children_pushed = false; } - // AND: an unpushable child is dropped; the residual filter still applies it. + // Visit all branches so an unsupported OR child cannot hide a later error. } - if (children.empty()) { + // AND can retain supported children; dropping any OR branch could exclude matching rows. + if (children.empty() || (type == Function::Type::OR && !all_children_pushed)) { return NullExpr(); } std::vector raw; @@ -583,13 +504,22 @@ VxExpressionPtr ConvertImpl(const std::shared_ptr& predicate, for (const VxExpressionPtr& child : children) { raw.push_back(child.get()); } - vx_expression* combined = type == Function::Type::AND - ? vx_expression_and(raw.data(), raw.size()) - : vx_expression_or(raw.data(), raw.size()); - return VxExpressionPtr(combined, vx_expression_free); + VxExpressionPtr combined(type == Function::Type::AND + ? vx_expression_and(raw.data(), raw.size()) + : vx_expression_or(raw.data(), raw.size()), + vx_expression_free); + PAIMON_RETURN_NOT_OK( + CheckNativeResult("create Vortex compound expression", combined.get())); + return combined; } const auto leaf = std::dynamic_pointer_cast(predicate); if (leaf == nullptr) { + vx_binary_operator op = VX_OPERATOR_EQ; + if (ComparisonOperator(type, &op) || type == Function::Type::IS_NULL || + type == Function::Type::IS_NOT_NULL || type == Function::Type::IN || + type == Function::Type::NOT_IN) { + return Status::Invalid("Leaf predicate must be a LeafPredicate"); + } return NullExpr(); } return ConvertLeaf(leaf, file_schema, root, options); @@ -597,16 +527,14 @@ VxExpressionPtr ConvertImpl(const std::shared_ptr& predicate, } // namespace -VxExpressionPtr VortexPredicateConverter::Convert(const std::shared_ptr& predicate, - const std::shared_ptr& file_schema, - const VortexPredicateOptions& options) { +Result VortexPredicateConverter::Convert( + const std::shared_ptr& predicate, const std::shared_ptr& file_schema, + const VortexPredicateOptions& options) { if (predicate == nullptr) { return NullExpr(); } VxExpressionPtr root(vx_expression_root(), vx_expression_free); - if (root == nullptr) { - return NullExpr(); - } + PAIMON_RETURN_NOT_OK(CheckNativeResult("create Vortex root expression", root.get())); return ConvertImpl(predicate, file_schema, root.get(), options); } diff --git a/src/paimon/format/vortex/vortex_predicate_converter.h b/src/paimon/format/vortex/vortex_predicate_converter.h index 39ed7e9c8..023868345 100644 --- a/src/paimon/format/vortex/vortex_predicate_converter.h +++ b/src/paimon/format/vortex/vortex_predicate_converter.h @@ -18,10 +18,11 @@ #pragma once -#include #include +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_ffi_util.h" +#include "paimon/result.h" namespace arrow { class Schema; @@ -33,14 +34,6 @@ class Predicate; namespace paimon::vortex { -enum class MembershipPushdownMode { AUTO, FORCE, OFF }; - -struct VortexPredicateOptions { - MembershipPushdownMode membership_pushdown = MembershipPushdownMode::AUTO; - // Per-leaf comparison and boolean node budget after deduplication/range compression. - uint32_t max_membership_nodes = 64; -}; - /// Converts Paimon predicates into Vortex scan filters, leaving unsupported conditions to the /// caller's residual filter. class VortexPredicateConverter { @@ -50,11 +43,12 @@ class VortexPredicateConverter { /// @param predicate The predicate to convert. /// @param file_schema Stored column types, including timestamp units and zones. - /// @param options Membership pushdown policy and expression budget. - /// @return An owned expression, or nullptr if no condition can be pushed. - static VxExpressionPtr Convert(const std::shared_ptr& predicate, - const std::shared_ptr& file_schema, - const VortexPredicateOptions& options = {}); + /// @param options Maximum input list length for IN and NOT IN pushdown. + /// @return An owned expression, nullptr if no condition can be pushed, or an error if + /// the predicate is malformed, types do not match, or native expression construction fails. + static Result Convert(const std::shared_ptr& predicate, + const std::shared_ptr& file_schema, + const VortexPredicateOptions& options); }; } // namespace paimon::vortex diff --git a/src/paimon/format/vortex/vortex_predicate_converter_test.cpp b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp index ce63c7e54..b7c4218bc 100644 --- a/src/paimon/format/vortex/vortex_predicate_converter_test.cpp +++ b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp @@ -33,12 +33,18 @@ #include "arrow/ipc/json_simple.h" #include "fmt/format.h" #include "gtest/gtest.h" +#include "paimon/common/predicate/and.h" +#include "paimon/common/predicate/compound_predicate_impl.h" +#include "paimon/common/predicate/leaf_predicate_impl.h" +#include "paimon/common/predicate/or.h" +#include "paimon/common/predicate/predicate_filter.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_io_callbacks.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" @@ -67,9 +73,48 @@ std::shared_ptr TestSchema() { bool Pushed(const std::shared_ptr& predicate, const std::shared_ptr& schema, const VortexPredicateOptions& options = {}) { - return VortexPredicateConverter::Convert(predicate, schema, options) != nullptr; + const Result result = + VortexPredicateConverter::Convert(predicate, schema, options); + EXPECT_OK(result.status()); + return result.ok() && result.value() != nullptr; } +void AssertInvalidConversion(const std::shared_ptr& predicate, + const std::shared_ptr& schema, + const std::vector& message_parts, + const VortexPredicateOptions& options = {}) { + const Result result = + VortexPredicateConverter::Convert(predicate, schema, options); + ASSERT_FALSE(result.ok()); + ASSERT_TRUE(result.status().IsInvalid()) << result.status().ToString(); + for (const std::string& part : message_parts) { + ASSERT_NE(result.status().ToString().find(part), std::string::npos) + << result.status().ToString(); + } +} + +// Simulate a malformed external Predicate whose function disagrees with its node shape. +class FunctionOnlyPredicate : public Predicate { + public: + explicit FunctionOnlyPredicate(const Function& function) : function_(function) {} + + bool operator==(const Predicate& other) const override { + return this == &other; + } + const Function& GetFunction() const override { + return function_; + } + std::shared_ptr Negate() const override { + return nullptr; + } + std::string ToString() const override { + return "FunctionOnlyPredicate"; + } + + private: + const Function& function_; +}; + struct NativeInCase { const char* name; std::shared_ptr type; @@ -91,7 +136,8 @@ std::string NativeErrorMessage(const vx_error* error) { void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptr& predicate, const std::vector& expected_rows, bool project_rowid_only = false, const VortexPredicateOptions& options = {}, bool expect_native_filter = true, - const vx_expression* native_filter_override = nullptr) { + const vx_expression* native_filter_override = nullptr, + const std::vector* expected_residual_rows = nullptr) { SCOPED_TRACE(test_case.name); auto values_result = arrow::ipc::internal::json::ArrayFromJSON(test_case.type, test_case.values_json); @@ -171,13 +217,15 @@ void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptr input, file_system.Open(path)); ASSERT_OK_AND_ASSIGN(int64_t length, input->Length()); auto input_context = std::make_shared(input); - VxDataSourcePtr source( - vx_data_source_new_callback(session.get(), VortexInputContext::MakeCallbacks(input_context), - static_cast(length), &error), - vx_data_source_free); + VxDataSourcePtr source(vx_data_source_new_callback_with_coalesce( + session.get(), VortexInputContext::MakeCallbacks(input_context), + static_cast(length), kDefaultVortexReadIoCoalesceMaxGap, + kDefaultVortexReadIoCoalesceMaxSize, &error), + vx_data_source_free); ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); ASSERT_NE(source, nullptr); - VxExpressionPtr filter = VortexPredicateConverter::Convert(predicate, schema, options); + ASSERT_OK_AND_ASSIGN(VxExpressionPtr filter, + VortexPredicateConverter::Convert(predicate, schema, options)); ASSERT_EQ(filter != nullptr, expect_native_filter); vx_scan_options scan_options{}; scan_options.filter = native_filter_override == nullptr ? filter.get() : native_filter_override; @@ -196,6 +244,7 @@ void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptr row_ids; + std::vector residual_row_ids; while (true) { VxPartitionPtr partition(vx_scan_next_partition(scan.get(), &error), vx_partition_free); ASSERT_EQ(error, nullptr) << NativeErrorMessage(error); @@ -228,14 +277,55 @@ void AssertNativeRows(const NativeInCase& test_case, const std::shared_ptrtype_id(), arrow::Type::INT64); const auto& typed_ids = checked_cast(*ids); + std::vector mask; + if (expected_residual_rows != nullptr) { + ASSERT_FALSE(project_rowid_only); + const std::shared_ptr residual = + std::dynamic_pointer_cast(predicate); + ASSERT_NE(residual, nullptr); + const auto rows = + arrow::StructArray::Make(batch->columns(), batch->schema()->fields()); + ASSERT_TRUE(rows.ok()) << rows.status().ToString(); + ASSERT_OK_AND_ASSIGN( + mask, residual->Test(*rows.ValueOrDie(), arrow::default_memory_pool())); + ASSERT_EQ(mask.size(), static_cast(batch->num_rows())); + } for (int64_t row = 0; row < batch->num_rows(); ++row) { ASSERT_FALSE(typed_ids.IsNull(row)); - row_ids.push_back(typed_ids.Value(row)); + const int64_t row_id = typed_ids.Value(row); + row_ids.push_back(row_id); + if (expected_residual_rows != nullptr) { + ASSERT_GE(row_id, 0); + ASSERT_LT(row_id, values_result.ValueOrDie()->length()); + const auto actual = batch->column(1)->GetScalar(row); + const auto expected = values_result.ValueOrDie()->GetScalar(row_id); + ASSERT_TRUE(actual.ok()); + ASSERT_TRUE(expected.ok()); + ASSERT_TRUE(actual.ValueOrDie()->Equals(*expected.ValueOrDie())); + if (mask[row] != 0) { + residual_row_ids.push_back(row_id); + } + } } } } std::sort(row_ids.begin(), row_ids.end()); ASSERT_EQ(row_ids, expected_rows); + if (expected_residual_rows != nullptr) { + std::sort(residual_row_ids.begin(), residual_row_ids.end()); + ASSERT_EQ(residual_row_ids, *expected_residual_rows); + } +} + +void AssertNativeAndResidualRows(const NativeInCase& test_case, + const std::shared_ptr& predicate, + const std::vector& expected_candidates, + const std::vector& expected_matches, + const VortexPredicateOptions& options = {}, + bool expect_native_filter = true) { + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, predicate, expected_candidates, false, + options, expect_native_filter, nullptr, + &expected_matches)); } // Compare an unsafe direct native expression with the converter's conservative fallback. @@ -347,25 +437,338 @@ TEST(VortexPredicateConverterTest, ComparisonsRequireMatchingFileAndLiteralTypes {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, PredicateBuilder::GreaterOrEqual}) { - ASSERT_FALSE(Pushed(build(0, "i", FieldType::INT, Literal(int64_t{1})), TestSchema())); - ASSERT_FALSE(Pushed(build(0, "i", FieldType::BIGINT, Literal(int64_t{1})), TestSchema())); - ASSERT_FALSE(Pushed(build(0, "dt", FieldType::INT, Literal(int32_t{1})), TestSchema())); - ASSERT_FALSE(Pushed(build(0, "dec_compact", FieldType::DECIMAL, - Literal(Decimal::FromUnscaledLong(12345, 19, 2))), - TestSchema())); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(build(0, "i", FieldType::INT, Literal(int64_t{1})), + TestSchema(), {"i", "expected INT", "actual BIGINT"})); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(build(0, "i", FieldType::BIGINT, Literal(int64_t{1})), + TestSchema(), {"i", "expected INT", "actual BIGINT"})); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(build(0, "dt", FieldType::INT, Literal(int32_t{1})), + TestSchema(), {"dt", "expected DATE", "actual INT"})); + ASSERT_TRUE(Pushed(build(0, "dec_compact", FieldType::DECIMAL, + Literal(Decimal::FromUnscaledLong(12345, 19, 2))), + TestSchema())); ASSERT_FALSE(Pushed(build(0, "dec_compact", FieldType::DECIMAL, Literal(Decimal::FromUnscaledLong(12345, 10, 3))), TestSchema())); } } +TEST(VortexPredicateConverterTest, LosslessDecimalComparisonsNormalizeToFileType) { + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + const std::vector builders = { + PredicateBuilder::Equal, PredicateBuilder::NotEqual, + PredicateBuilder::LessThan, PredicateBuilder::LessOrEqual, + PredicateBuilder::GreaterThan, PredicateBuilder::GreaterOrEqual}; + const std::vector> positive_rows = {{3}, {0, 2}, {0, 2}, + {0, 2, 3}, {}, {3}}; + struct DecimalCase { + NativeInCase native; + std::vector> expected_rows; + }; + Decimal::int128_t maximum = 1; + for (int32_t digit = 0; digit < 37; ++digit) { + maximum *= 10; + } + --maximum; + const std::vector cases = { + {{"precision_only", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(19, 2, 120))}}, + positive_rows}, + {{"scale_down", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(20, 4, 12000))}}, + positive_rows}, + {{"scale_up", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(10, 1, 12))}}, + positive_rows}, + {{"negative_scale_down", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(20, 4, -12000))}}, + {{0}, {2, 3}, {}, {0}, {2, 3}, {0, 2, 3}}}, + {{"zero_scale_down", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(38, 38, 0))}}, + {{2}, {0, 3}, {0}, {0, 2}, {3}, {2, 3}}}, + {{"compact_to_wide", + arrow::decimal128(19, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(18, 1, 12))}}, + positive_rows}, + {{"38_digit_unscaled_value", + arrow::decimal128(37, 0), + FieldType::DECIMAL, + R"(["-9999999999999999999999999999999999999",null,"0","9999999999999999999999999999999999999"])", + {Literal(Decimal(38, 1, maximum * 10))}}, + positive_rows}, + }; + for (const DecimalCase& test_case : cases) { + for (size_t operation = 0; operation < builders.size(); ++operation) { + SCOPED_TRACE(operation); + const std::shared_ptr predicate = + builders[operation](1, "value", FieldType::DECIMAL, test_case.native.literals[0]); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + test_case.native, predicate, test_case.expected_rows[operation], + test_case.expected_rows[operation])); + } + } +} + +TEST(VortexPredicateConverterTest, LosslessDecimalMembershipNormalizesEachLiteral) { + struct DecimalCase { + NativeInCase native; + std::vector in_rows; + std::vector not_in_rows; + }; + Decimal::int128_t maximum = 1; + for (int32_t digit = 0; digit < 37; ++digit) { + maximum *= 10; + } + --maximum; + const std::vector cases = { + {{"mixed_parameters", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","4.56"])", + {Literal(Decimal(19, 1, -12)), Literal(Decimal(20, 4, 45600))}}, + {0, 3}, + {2}}, + {{"zero", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","4.56"])", + {Literal(Decimal(38, 38, 0))}}, + {2}, + {0, 3}}, + {{"compact_and_wide", + arrow::decimal128(19, 2), + FieldType::DECIMAL, + R"(["-1.20",null,"0.00","1.20"])", + {Literal(Decimal(18, 1, -12)), Literal(Decimal(38, 4, 12000))}}, + {0, 3}, + {2}}, + {{"38_digit_unscaled_values", + arrow::decimal128(37, 0), + FieldType::DECIMAL, + R"(["-9999999999999999999999999999999999999",null,"0","9999999999999999999999999999999999999"])", + {Literal(Decimal(38, 1, -maximum * 10)), Literal(Decimal(38, 1, maximum * 10))}}, + {0, 3}, + {2}}, + }; + for (const DecimalCase& test_case : cases) { + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "value", FieldType::DECIMAL, + test_case.native.literals) + : PredicateBuilder::In(1, "value", FieldType::DECIMAL, + test_case.native.literals); + const std::vector& expected = + negate ? test_case.not_in_rows : test_case.in_rows; + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case.native, predicate, expected, expected)); + } + } +} + +TEST(VortexPredicateConverterTest, UnrepresentableDecimalsPreserveNativeCandidatesAndResidualRows) { + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + const std::vector builders = { + PredicateBuilder::Equal, PredicateBuilder::NotEqual, + PredicateBuilder::LessThan, PredicateBuilder::LessOrEqual, + PredicateBuilder::GreaterThan, PredicateBuilder::GreaterOrEqual}; + struct DecimalCase { + NativeInCase native; + std::vector> expected_rows; + }; + Decimal::int128_t maximum = 1; + for (int32_t digit = 0; digit < 38; ++digit) { + maximum *= 10; + } + --maximum; + const std::vector cases = { + {{"lossy_positive", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + {Literal(Decimal(10, 3, 1234))}}, + {{}, {0, 2, 3}, {0, 2}, {0, 2}, {3}, {3}}}, + {{"lossy_negative", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + {Literal(Decimal(10, 3, -1234))}}, + {{}, {0, 2, 3}, {}, {}, {0, 2, 3}, {0, 2, 3}}}, + {{"precision_overflow", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + {Literal(Decimal(19, 2, 10000000000LL))}}, + {{}, {0, 2, 3}, {0, 2, 3}, {0, 2, 3}, {}, {}}}, + {{"wide_value_cannot_be_truncated_to_compact", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + // Its low 64 bits look like 1.23; the full source value must be checked first. + {Literal(Decimal(20, 2, (static_cast(1) << 64) + 123))}}, + {{}, {0, 2, 3}, {0, 2, 3}, {0, 2, 3}, {}, {}}}, + {{"128_bit_scale_overflow", + arrow::decimal128(38, 38), + FieldType::DECIMAL, + R"(["-0.12000000000000000000000000000000000000",null,"0.12000000000000000000000000000000000000","0.99000000000000000000000000000000000000"])", + {Literal(Decimal(38, 0, maximum))}}, + {{}, {0, 2, 3}, {0, 2, 3}, {0, 2, 3}, {}, {}}}, + }; + const std::vector all_rows = {0, 1, 2, 3}; + for (const DecimalCase& test_case : cases) { + for (size_t operation = 0; operation < builders.size(); ++operation) { + SCOPED_TRACE(operation); + const std::shared_ptr predicate = + builders[operation](1, "value", FieldType::DECIMAL, test_case.native.literals[0]); + // Rounding +/-1.234 to +/-1.23 would lose rows in strict inequalities. + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case.native, predicate, all_rows, + test_case.expected_rows[operation], {}, false)); + } + const auto& column_type = + checked_cast(*test_case.native.type); + const std::vector literals = { + Literal(Decimal(column_type.precision(), column_type.scale(), 0)), + test_case.native.literals[0]}; + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "value", FieldType::DECIMAL, literals) + : PredicateBuilder::In(1, "value", FieldType::DECIMAL, literals); + const std::vector expected = + negate ? std::vector{0, 2, 3} : std::vector{}; + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case.native, predicate, + all_rows, expected, {}, false)); + } + } +} + +TEST(VortexPredicateConverterTest, DecimalNormalizationRejectsInvalidLiteralMetadata) { + const std::shared_ptr schema = + arrow::schema({arrow::field("value", arrow::decimal128(10, 2))}); + const std::vector literals = { + Literal(Decimal(0, 0, 0)), Literal(Decimal(39, 2, 123)), + Literal(Decimal(10, -1, 123)), Literal(Decimal(10, 11, 123)), + Literal(Decimal(2, 0, 123)), Literal(Decimal(10, 2, 10000000000LL)), + }; + for (const Literal& literal : literals) { + ASSERT_FALSE( + Pushed(PredicateBuilder::Equal(0, "value", FieldType::DECIMAL, literal), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(0, "value", FieldType::DECIMAL, {literal}), schema)); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(0, "value", FieldType::DECIMAL, {literal}), schema)); + } +} + +TEST(VortexPredicateConverterTest, SameParameterDecimalOverflowPreservesNativeCandidates) { + const NativeInCase test_case = {"same-parameter-wide-overflow", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + {}}; + // The low 64 bits are 123, but the full value does not fit the advertised compact precision. + const Decimal::int128_t unscaled = (static_cast(1) << 64) + 123; + const Literal invalid(Decimal(10, 2, unscaled)); + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, + build(1, "value", FieldType::DECIMAL, invalid), + {0, 1, 2, 3}, false, {}, false)); + } + for (const std::shared_ptr& predicate : + {PredicateBuilder::In(1, "value", FieldType::DECIMAL, {invalid}), + PredicateBuilder::NotIn(1, "value", FieldType::DECIMAL, {invalid})}) { + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, false, {}, false)); + } +} + +TEST(VortexPredicateConverterTest, DecimalMembershipMaximumPrecedesNormalization) { + const NativeInCase test_case = {"decimal_raw_length", + arrow::decimal128(10, 2), + FieldType::DECIMAL, + R"(["-1.23",null,"1.23","1.24"])", + {}}; + const std::vector literals(33, Literal(Decimal(38, 4, 12300))); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "value", FieldType::DECIMAL, literals) + : PredicateBuilder::In(1, "value", FieldType::DECIMAL, literals); + const std::vector expected = + negate ? std::vector{0, 3} : std::vector{2}; + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case, predicate, {0, 1, 2, 3}, expected, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, predicate, expected, + expected, VortexPredicateOptions{33})); + } +} + +TEST(VortexPredicateConverterTest, MismatchedBasicTypesReturnInvalid) { + struct TypeCase { + std::shared_ptr arrow_type; + FieldType field_type; + const char* type_name; + }; + const std::vector cases = { + {arrow::boolean(), FieldType::BOOLEAN, "BOOLEAN"}, + {arrow::int8(), FieldType::TINYINT, "TINYINT"}, + {arrow::int16(), FieldType::SMALLINT, "SMALLINT"}, + {arrow::int32(), FieldType::INT, "INT"}, + {arrow::int64(), FieldType::BIGINT, "BIGINT"}, + {arrow::float32(), FieldType::FLOAT, "FLOAT"}, + {arrow::float64(), FieldType::DOUBLE, "DOUBLE"}, + {arrow::utf8(), FieldType::STRING, "STRING"}, + {arrow::binary(), FieldType::BINARY, "BINARY"}, + {arrow::date32(), FieldType::DATE, "DATE"}, + {arrow::decimal128(10, 2), FieldType::DECIMAL, "DECIMAL"}, + {arrow::timestamp(arrow::TimeUnit::MILLI), FieldType::TIMESTAMP, "TIMESTAMP"}, + }; + for (const TypeCase& test_case : cases) { + SCOPED_TRACE(test_case.type_name); + const std::shared_ptr schema = + arrow::schema({arrow::field("value", test_case.arrow_type)}); + const bool int_field = test_case.field_type == FieldType::INT; + const FieldType wrong_type = int_field ? FieldType::BIGINT : FieldType::INT; + const Literal wrong_literal = int_field ? Literal(int64_t{1}) : Literal(int32_t{1}); + const std::string expected = fmt::format("expected {}", test_case.type_name); + const std::string actual = int_field ? "actual BIGINT" : "actual INT"; + // Both the leaf-to-file contract and the literal-to-leaf contract are errors. + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(PredicateBuilder::Equal(0, "value", wrong_type, wrong_literal), + schema, {"value", expected, actual})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + PredicateBuilder::Equal(0, "value", test_case.field_type, wrong_literal), schema, + {"value", expected, actual})); + } +} + TEST(VortexPredicateConverterTest, UnboundPredicatesPreserveNativeCandidateRows) { const NativeInCase test_case = {"unbound", arrow::int32(), FieldType::INT, "[1,null,2,3]", {}}; const std::vector> predicates = { PredicateBuilder::Equal(2, "missing", FieldType::INT, Literal(int32_t{1})), PredicateBuilder::IsNull(2, "missing", FieldType::INT), PredicateBuilder::IsNotNull(2, "missing", FieldType::INT), - PredicateBuilder::Equal(1, "value", FieldType::BIGINT, Literal(int64_t{1})), }; for (const std::shared_ptr& predicate : predicates) { ASSERT_NO_FATAL_FAILURE( @@ -437,14 +840,14 @@ TEST(VortexPredicateConverterTest, RandomizedIntegerMembershipMatchesIndependent expected.push_back(row_ids[row]); } } - for (MembershipPushdownMode mode : - {MembershipPushdownMode::AUTO, MembershipPushdownMode::FORCE}) { - const VortexPredicateOptions options{mode, 64}; - const bool pushed = Pushed(predicate, schema, options); + for (uint32_t maximum : {uint32_t{32}, uint32_t{80}}) { + const VortexPredicateOptions options{maximum}; + const bool pushed = literals.size() <= maximum; + ASSERT_EQ(Pushed(predicate, schema, options), pushed); const std::vector candidates = pushed ? expected : std::vector{0, 1, 2, 3}; - ASSERT_NO_FATAL_FAILURE( - AssertNativeRows(test_case, predicate, candidates, false, options, pushed)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + test_case, predicate, candidates, expected, options, pushed)); } } } @@ -536,6 +939,61 @@ TEST(VortexPredicateConverterTest, LosslessTimestampComparisonsFilterNativeFiles } } +TEST(VortexPredicateConverterTest, TimestampComparisonsPreserveBoundaryDecisionsAndNativeValues) { + const std::vector cases = { + {"micros_calendar_limits", + arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(253402207200999, 999000)), + Literal(Timestamp::FromEpochMillis(-62167219200000)), + Literal(Timestamp::FromEpochMillis(-1, 999000)), + Literal(Timestamp::FromEpochMillis(9223372036854775, 808000)), + Literal(Timestamp::FromEpochMillis(-9223372036854776, 192000)), + Literal(Timestamp::FromEpochMillis(0, 2000))}}, + {"nanos_boundaries", + arrow::timestamp(arrow::TimeUnit::NANO), + FieldType::TIMESTAMP, + "[-1,null,7,2]", + {Literal(Timestamp::FromEpochMillis(9223372036854, 775807)), + Literal(Timestamp::FromEpochMillis(-9223372036854)), + Literal(Timestamp::FromEpochMillis(-1, 999999)), + Literal(Timestamp::FromEpochMillis(9223372036854, 775808)), + Literal(Timestamp::FromEpochMillis(-9223372036855, 224192)), + Literal(Timestamp::FromEpochMillis(0, 2))}}, + }; + for (const NativeInCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + const std::shared_ptr schema = + arrow::schema({arrow::field("value", test_case.type)}); + // Accepted expressions must also fit the native calendar range. + // The microsecond values cover year 0000 through Jiff's year 9999 limit. + ASSERT_TRUE( + Pushed(PredicateBuilder::Equal(0, "value", FieldType::TIMESTAMP, test_case.literals[0]), + schema)); + ASSERT_TRUE(Pushed( + PredicateBuilder::LessOrEqual(0, "value", FieldType::TIMESTAMP, test_case.literals[1]), + schema)); + ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + test_case, + PredicateBuilder::Equal(1, "value", FieldType::TIMESTAMP, test_case.literals[2]), {0})); + ASSERT_NO_FATAL_FAILURE( + AssertNativeRows(test_case, + PredicateBuilder::GreaterOrEqual(1, "value", FieldType::TIMESTAMP, + test_case.literals[5]), + {2, 3})); + // One unit beyond INT64_MAX must fall back rather than overflow. + ASSERT_FALSE( + Pushed(PredicateBuilder::Equal(0, "value", FieldType::TIMESTAMP, test_case.literals[3]), + schema)); + // Although the final value is INT64_MIN, the intermediate millisecond product + // overflows. Preserve the existing fallback before calling the common utility. + ASSERT_FALSE( + Pushed(PredicateBuilder::Equal(0, "value", FieldType::TIMESTAMP, test_case.literals[4]), + schema)); + } +} + TEST(VortexPredicateConverterTest, TruncatedTimestampComparisonFallsBackWithoutLosingRows) { const NativeInCase test_case = {"micros_fraction", arrow::timestamp(arrow::TimeUnit::MICRO), @@ -655,14 +1113,17 @@ TEST(VortexPredicateConverterTest, UnsupportedFunctionsAndLiteralsAreNotPushed) Pushed(PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(FieldType::INT)), schema)); } -TEST(VortexPredicateConverterTest, InSmallAndLargeListsArePushed) { +TEST(VortexPredicateConverterTest, InSmallAndLargeListsRespectConfiguredMaximum) { const std::shared_ptr schema = TestSchema(); for (const int32_t size : {1, 2, 20, 21, 1024}) { std::vector literals; for (int32_t value = 0; value < size; ++value) { literals.emplace_back(value); } - ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), schema)); + const std::shared_ptr predicate = + PredicateBuilder::In(0, "i", FieldType::INT, literals); + ASSERT_EQ(Pushed(predicate, schema), size <= 32); + ASSERT_TRUE(Pushed(predicate, schema, {1024})); } ASSERT_TRUE(Pushed( PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int32_t{1})}), @@ -808,13 +1269,20 @@ TEST(VortexPredicateConverterTest, NotInSmallAndLargeListsFilterNativeFiles) { } const std::vector expected_rows = size == 1 ? std::vector{2, 3} : std::vector{2}; - ASSERT_NO_FATAL_FAILURE(AssertNativeRows( - test_case, PredicateBuilder::NotIn(1, "value", FieldType::INT, literals), - expected_rows)); + const std::shared_ptr predicate = + PredicateBuilder::NotIn(1, "value", FieldType::INT, literals); + const bool pushed = size <= 32; + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + test_case, predicate, pushed ? expected_rows : std::vector{0, 1, 2, 3}, + expected_rows, {}, pushed)); + if (size <= 100) { + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( + test_case, predicate, expected_rows, expected_rows, {static_cast(size)})); + } } } -TEST(VortexPredicateConverterTest, IntegerRangesAtTypeLimitsFilterNativeFiles) { +TEST(VortexPredicateConverterTest, IntegerListsAtTypeLimitsFilterNativeFiles) { const std::vector cases = { IntegerRangeBoundaryCase("tinyint_range_limits", arrow::int8(), FieldType::TINYINT, "[-128,null,0,127]"), @@ -835,7 +1303,7 @@ TEST(VortexPredicateConverterTest, IntegerRangesAtTypeLimitsFilterNativeFiles) { } } -TEST(VortexPredicateConverterTest, IntegerRangesPreserveGapsAndSingletonsInNativeFiles) { +TEST(VortexPredicateConverterTest, IntegerListsPreserveGapsAndSingletonsInNativeFiles) { NativeInCase test_case = { "integer_ranges_with_gap", arrow::int32(), FieldType::INT, "[-7,null,0,6]", {}}; for (int32_t value = -8; value <= -1; ++value) { @@ -867,12 +1335,12 @@ TEST(VortexPredicateConverterTest, FullIntegerDomainAndDuplicateSingletonFilterN value <= std::numeric_limits::max(); ++value) { test_case.literals.emplace_back(static_cast(value)); } - ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( test_case, PredicateBuilder::In(1, "value", FieldType::TINYINT, test_case.literals), - {0, 2, 3})); - ASSERT_NO_FATAL_FAILURE(AssertNativeRows( + {0, 1, 2, 3}, {0, 2, 3}, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows( test_case, PredicateBuilder::NotIn(1, "value", FieldType::TINYINT, test_case.literals), - {})); + {0, 1, 2, 3}, {}, {}, false)); test_case.name = "duplicate_singleton"; test_case.literals = std::vector(8, Literal(int8_t{-128})); ASSERT_NO_FATAL_FAILURE(AssertNativeRows( @@ -882,7 +1350,7 @@ TEST(VortexPredicateConverterTest, FullIntegerDomainAndDuplicateSingletonFilterN {2, 3})); } -TEST(VortexPredicateConverterTest, IntegerRangesComposeWithNativePredicatesAndProjection) { +TEST(VortexPredicateConverterTest, IntegerListsComposeWithNativePredicatesAndProjection) { NativeInCase test_case = { "integer_range_compound", arrow::int32(), FieldType::INT, "[-4,null,8,4]", {}}; for (int32_t value = -4; value <= 4; ++value) { @@ -915,7 +1383,7 @@ TEST(VortexPredicateConverterTest, IntegerRangesComposeWithNativePredicatesAndPr ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, not_in, {2}, true)); } -TEST(VortexPredicateConverterTest, CompressibleIntegerListsStillValidateEveryLiteral) { +TEST(VortexPredicateConverterTest, IntegerListsValidateEveryLiteral) { const std::shared_ptr schema = TestSchema(); std::vector safe; for (int32_t value = 0; value < 16; ++value) { @@ -925,14 +1393,28 @@ TEST(VortexPredicateConverterTest, CompressibleIntegerListsStillValidateEveryLit for (size_t position : {size_t{0}, size_t{8}, safe.size()}) { std::vector literals = safe; literals.insert(literals.begin() + position, unsafe); - ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), schema)); - ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, literals), schema)); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + if (unsafe.IsNull()) { + ASSERT_FALSE(Pushed(predicate, schema)); + } else { + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, {"i", "expected INT", "actual BIGINT"})); + } + } } } - ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, safe), - arrow::schema({arrow::field("i", arrow::int64())}))); - ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, safe), - arrow::schema({arrow::field("i", arrow::int64())}))); + const std::shared_ptr bigint_schema = + arrow::schema({arrow::field("i", arrow::int64())}); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, safe) + : PredicateBuilder::In(0, "i", FieldType::INT, safe); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion(predicate, bigint_schema, + {"i", "expected BIGINT", "actual INT"})); + } } TEST(VortexPredicateConverterTest, NotInDuplicatesMissesAndAllExcludedFilterNativeFiles) { @@ -991,7 +1473,7 @@ TEST(VortexPredicateConverterTest, NotInCompoundPredicatesAndProjectionFilterNat arrow::field("value", arrow::int32())}))); } -TEST(VortexPredicateConverterTest, NotInUnsafeLiteralsAreNotPushed) { +TEST(VortexPredicateConverterTest, NotInUnsafeLiteralsFallBackAndTypeMismatchesFail) { const std::shared_ptr schema = TestSchema(); ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {}), schema)); ASSERT_FALSE( @@ -999,11 +1481,12 @@ TEST(VortexPredicateConverterTest, NotInUnsafeLiteralsAreNotPushed) { ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(FieldType::INT)}), schema)); - ASSERT_FALSE(Pushed( + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( PredicateBuilder::NotIn(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int64_t{2})}), - schema)); - ASSERT_FALSE( - Pushed(PredicateBuilder::NotIn(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema)); + schema, {"i", "expected INT", "actual BIGINT"})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + PredicateBuilder::NotIn(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema, + {"i", "expected INT", "actual BIGINT"})); ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(0, "absent", FieldType::INT, {Literal(int32_t{1})}), schema)); ASSERT_FALSE( @@ -1033,7 +1516,7 @@ TEST(VortexPredicateConverterTest, NotInUnsafeLiteralsAreNotPushed) { "ts", arrow::timestamp(arrow::TimeUnit::MICRO, "Asia/Shanghai"))}))); } -TEST(VortexPredicateConverterTest, InUnsupportedOrInconsistentLiteralsAreNotPushed) { +TEST(VortexPredicateConverterTest, InUnsafeLiteralsFallBackAndTypeMismatchesFail) { const std::shared_ptr schema = TestSchema(); ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, {}), schema)); ASSERT_FALSE( @@ -1041,11 +1524,12 @@ TEST(VortexPredicateConverterTest, InUnsupportedOrInconsistentLiteralsAreNotPush ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(FieldType::INT)}), schema)); - ASSERT_FALSE(Pushed( + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( PredicateBuilder::In(0, "i", FieldType::INT, {Literal(int32_t{1}), Literal(int64_t{2})}), - schema)); - ASSERT_FALSE( - Pushed(PredicateBuilder::In(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema)); + schema, {"i", "expected INT", "actual BIGINT"})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + PredicateBuilder::In(0, "i", FieldType::BIGINT, {Literal(int64_t{1})}), schema, + {"i", "expected INT", "actual BIGINT"})); ASSERT_FALSE( Pushed(PredicateBuilder::In(0, "absent", FieldType::INT, {Literal(int32_t{1})}), schema)); ASSERT_FALSE( @@ -1063,9 +1547,9 @@ TEST(VortexPredicateConverterTest, InUnsupportedOrInconsistentLiteralsAreNotPush {Literal(Decimal::FromUnscaledLong(123, 10, 2)), Literal(Decimal::FromUnscaledLong(123, 10, 3))}), decimal_schema)); - ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "dec", FieldType::DECIMAL, - {Literal(Decimal::FromUnscaledLong(123, 11, 2))}), - decimal_schema)); + ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "dec", FieldType::DECIMAL, + {Literal(Decimal::FromUnscaledLong(123, 11, 2))}), + decimal_schema)); } TEST(VortexPredicateConverterTest, InTimestampsRequireMatchingUnitsAndLosslessValues) { @@ -1153,86 +1637,486 @@ TEST(VortexPredicateConverterTest, AndDropsUnsupportedChildrenButOrRejects) { ASSERT_FALSE(Pushed(or_partial, schema)); } -TEST(VortexPredicateConverterTest, MembershipListBudgetCountsBooleanNodesAndNullGuards) { +TEST(VortexPredicateConverterTest, CompoundPredicatesPropagateTypeErrorsRegardlessOfBranchOrder) { const std::shared_ptr schema = TestSchema(); - std::vector literals; - for (int32_t value = 0; value < 33; ++value) { - const std::string text = fmt::format("value-{}", value); - literals.emplace_back(FieldType::STRING, text.data(), text.size()); + const std::shared_ptr supported = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int32_t{1})); + const std::shared_ptr invalid = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int64_t{1})); + const std::shared_ptr unsupported = + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(FieldType::INT)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({supported, invalid})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({supported, invalid})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_after_fallback, + PredicateBuilder::And({unsupported, invalid})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_before_fallback, + PredicateBuilder::Or({invalid, unsupported})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_after_fallback, + PredicateBuilder::Or({unsupported, invalid})); + for (const std::shared_ptr& predicate : + {and_predicate, or_predicate, and_after_fallback, or_before_fallback, or_after_fallback}) { + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(predicate, schema, {"i", "expected INT", "actual BIGINT"})); + } +} + +TEST(VortexPredicateConverterTest, MalformedPredicateStructuresReturnInvalid) { + const std::shared_ptr schema = TestSchema(); + const std::shared_ptr unsupported = + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(FieldType::INT)}); + const std::vector functions = {&And::Instance(), &Or::Instance()}; + for (const CompoundFunction* function : functions) { + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + std::make_shared(*function), schema, {"CompoundPredicate"})); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(std::make_shared( + *function, std::vector>{}), + schema, {"at least one child"})); + for (const std::vector>& children : + {std::vector>{nullptr, unsupported}, + std::vector>{unsupported, nullptr}}) { + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + std::make_shared(*function, children), schema, + {"null child"})); + } + } + const std::shared_ptr comparison = + PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int32_t{1})); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(std::make_shared(comparison->GetFunction()), + schema, {"LeafPredicate"})); +} + +TEST(VortexPredicateConverterTest, ComparisonArityReturnsInvalid) { + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + const auto valid = std::dynamic_pointer_cast( + build(0, "i", FieldType::INT, Literal(int32_t{1}))); + ASSERT_NE(valid, nullptr); + for (const std::vector& literals : + {std::vector{}, + std::vector{Literal(int32_t{1}), Literal(int32_t{2})}}) { + const std::shared_ptr malformed = std::make_shared( + valid->GetLeafFunction(), 0, "i", FieldType::INT, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion(malformed, TestSchema(), + {"i", "requires exactly one literal"})); + const std::shared_ptr unsupported = + PredicateBuilder::In(0, "i", FieldType::INT, {Literal(FieldType::INT)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({unsupported, malformed})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion(disjunction, TestSchema(), + {"i", "requires exactly one literal"})); + } + } +} + +TEST(VortexPredicateConverterTest, NestedCompoundsPreserveNativeCandidatesAndPropagateErrors) { + const NativeInCase test_case = { + "nested-fallback", arrow::utf8(), FieldType::STRING, R"(["aa",null,"ab","ba"])", {}}; + const Literal ab(FieldType::STRING, "ab", 2); + const std::shared_ptr supported = + PredicateBuilder::GreaterOrEqual(1, "value", FieldType::STRING, ab); + const std::shared_ptr equal = + PredicateBuilder::Equal(1, "value", FieldType::STRING, Literal(FieldType::STRING, "aa", 2)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr unsupported, + PredicateBuilder::StartsWith(1, "value", FieldType::STRING, + Literal(FieldType::STRING, "a", 1))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr partial_and, + PredicateBuilder::And({supported, unsupported})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr fallback_or, + PredicateBuilder::Or({unsupported, supported})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr nested_and, + PredicateBuilder::And({fallback_or, equal})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr nested_or, + PredicateBuilder::Or({partial_and, equal})); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, partial_and, {2, 3}, {2})); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case, fallback_or, {0, 1, 2, 3}, {0, 2, 3}, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, nested_and, {0}, {0})); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, nested_or, {0, 2, 3}, {0, 2})); + + const std::shared_ptr invalid = + PredicateBuilder::Equal(1, "value", FieldType::STRING, Literal(int64_t{1})); + using CompoundBuilder = decltype(&PredicateBuilder::And); + const std::shared_ptr schema = + arrow::schema({arrow::field("value", arrow::utf8())}); + for (CompoundBuilder outer : {PredicateBuilder::And, PredicateBuilder::Or}) { + for (CompoundBuilder inner : {PredicateBuilder::And, PredicateBuilder::Or}) { + for (bool invalid_first : {false, true}) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr nested, + inner(invalid_first + ? std::vector>{invalid, unsupported} + : std::vector>{unsupported, invalid})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate, + outer({unsupported, nested})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, {"value", "expected STRING", "actual BIGINT"})); + } + } + } +} + +TEST(VortexPredicateConverterTest, NativeScalarErrorsPropagateThroughAllPredicateBranches) { + const Literal invalid_utf8(FieldType::STRING, "\xff", 1); + std::vector> invalid; + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + invalid.push_back(build(1, "s", FieldType::STRING, invalid_utf8)); + } + invalid.push_back(PredicateBuilder::In(1, "s", FieldType::STRING, + {Literal(FieldType::STRING, "valid", 5), invalid_utf8})); + invalid.push_back(PredicateBuilder::NotIn( + 1, "s", FieldType::STRING, {Literal(FieldType::STRING, "valid", 5), invalid_utf8})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr unsupported, + PredicateBuilder::StartsWith(1, "s", FieldType::STRING, + Literal(FieldType::STRING, "v", 1))); + for (const std::shared_ptr& predicate : invalid) { + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, TestSchema(), {"create Vortex predicate scalar", "invalid utf-8"})); + using CompoundBuilder = decltype(&PredicateBuilder::And); + for (CompoundBuilder build : {PredicateBuilder::And, PredicateBuilder::Or}) { + for (bool invalid_first : {false, true}) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr compound, + build(invalid_first + ? std::vector>{predicate, unsupported} + : std::vector>{unsupported, predicate})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + compound, TestSchema(), {"create Vortex predicate scalar", "invalid utf-8"})); + } + } + } + // The list budget is still checked before inspecting a declined native scalar. + std::vector over_budget(33, invalid_utf8); + ASSERT_FALSE( + Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, over_budget), TestSchema())); + ASSERT_FALSE( + Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, over_budget), TestSchema())); + ASSERT_TRUE(Pushed( + PredicateBuilder::Equal(1, "s", FieldType::STRING, Literal(FieldType::STRING, "valid", 5)), + TestSchema())); +} + +TEST(VortexPredicateConverterTest, NativeTemporalScalarBoundsFallBackWithoutErrors) { + const int64_t maximum = std::numeric_limits::max(); + struct TemporalCase { + const char* name; + std::shared_ptr type; + FieldType field_type; + Literal literal; + }; + const std::vector cases = { + {"maximum-micros", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(maximum / 1000, + static_cast((maximum % 1000) * 1000)))}, + {"minimum-micros", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(std::numeric_limits::min() / 1000))}, + {"paimon-calendar-limit", arrow::timestamp(arrow::TimeUnit::MICRO), FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(253402300799999, 999000))}, + {"native-epoch-seconds", arrow::timestamp(arrow::TimeUnit::SECOND, "UTC"), + FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(253402207201000))}, + {"native-epoch-millis", arrow::timestamp(arrow::TimeUnit::MILLI), FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(253402207201000))}, + {"maximum-date", arrow::date32(), FieldType::DATE, + Literal(FieldType::DATE, std::numeric_limits::max())}, + {"minimum-date", arrow::date32(), FieldType::DATE, + Literal(FieldType::DATE, std::numeric_limits::min())}}; + using ComparisonBuilder = decltype(&PredicateBuilder::Equal); + for (const TemporalCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + const std::shared_ptr schema = + arrow::schema({arrow::field("value", test_case.type)}); + for (ComparisonBuilder build : + {PredicateBuilder::Equal, PredicateBuilder::NotEqual, PredicateBuilder::LessThan, + PredicateBuilder::LessOrEqual, PredicateBuilder::GreaterThan, + PredicateBuilder::GreaterOrEqual}) { + ASSERT_FALSE( + Pushed(build(0, "value", test_case.field_type, test_case.literal), schema)); + } + for (const std::shared_ptr& predicate : + {PredicateBuilder::In(0, "value", test_case.field_type, {test_case.literal}), + PredicateBuilder::NotIn(0, "value", test_case.field_type, {test_case.literal})}) { + ASSERT_FALSE(Pushed(predicate, schema)); + } + } + const std::shared_ptr nanos_schema = + arrow::schema({arrow::field("value", arrow::timestamp(arrow::TimeUnit::NANO, "UTC"))}); + // INT64_MAX nanoseconds is in year 2262 and remains a valid native timestamp. + ASSERT_TRUE(Pushed( + PredicateBuilder::Equal(0, "value", FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis( + maximum / 1000000, static_cast(maximum % 1000000)))), + nanos_schema)); +} + +TEST(VortexPredicateConverterTest, TemporalRangeFallbacksDoNotHideRealErrors) { + const std::shared_ptr schema = + arrow::schema({arrow::field("date", arrow::date32()), arrow::field("name", arrow::utf8())}); + const Literal extreme(FieldType::DATE, std::numeric_limits::max()); + const Literal wrong_type(int64_t{1}); + const std::shared_ptr unsupported = + PredicateBuilder::LessThan(0, "date", FieldType::DATE, extreme); + for (bool negate : {false, true}) { + for (bool extreme_first : {false, true}) { + const std::vector literals = extreme_first + ? std::vector{extreme, wrong_type} + : std::vector{wrong_type, extreme}; + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "date", FieldType::DATE, literals) + : PredicateBuilder::In(0, "date", FieldType::DATE, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, {"date", "expected DATE", "actual BIGINT"})); + } + } + const auto valid = + PredicateBuilder::Equal(0, "date", FieldType::DATE, Literal(FieldType::DATE, 0)); + const auto valid_leaf = std::dynamic_pointer_cast(valid); + ASSERT_NE(valid_leaf, nullptr); + const auto malformed = std::make_shared( + valid_leaf->GetLeafFunction(), 0, "date", FieldType::DATE, std::vector{}); + const std::vector, std::string>> errors = { + {PredicateBuilder::Equal(0, "date", FieldType::DATE, wrong_type), "Literal type mismatch"}, + {PredicateBuilder::Equal(1, "name", FieldType::STRING, + Literal(FieldType::STRING, "\xff", 1)), + "invalid utf-8"}, + {malformed, "requires exactly one literal"}}; + using CompoundBuilder = decltype(&PredicateBuilder::And); + for (const auto& [invalid, message] : errors) { + for (CompoundBuilder build : {PredicateBuilder::And, PredicateBuilder::Or}) { + for (bool invalid_first : {false, true}) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr compound, + build(invalid_first + ? std::vector>{invalid, unsupported} + : std::vector>{unsupported, invalid})); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion(compound, schema, {message})); + } + } } +} + +TEST(VortexPredicateConverterTest, MembershipFallbacksDoNotHideTypeErrors) { + const Literal null_value(FieldType::INT); + const Literal wrong_type(int64_t{1}); + const Literal valid(int32_t{1}); + for (bool negate : {false, true}) { + for (bool invalid_first : {false, true}) { + const std::vector literals = + invalid_first ? std::vector{wrong_type, null_value} + : std::vector{null_value, wrong_type}; + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, TestSchema(), {"i", "expected INT", "actual BIGINT"})); + } + const std::shared_ptr fallback = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, {null_value, valid}) + : PredicateBuilder::In(0, "i", FieldType::INT, {null_value, valid}); + ASSERT_FALSE(Pushed(fallback, TestSchema())); + } +} + +TEST(VortexPredicateConverterTest, MembershipTypeErrorsAreIndependentOfMaximum) { + struct TypeCase { + std::shared_ptr type; + FieldType field_type; + Literal valid; + const char* expected_type; + }; + const std::vector cases = { + {arrow::int32(), FieldType::INT, Literal(int32_t{1}), "INT"}, + {arrow::decimal128(10, 2), FieldType::DECIMAL, Literal(Decimal(38, 4, 12300)), "DECIMAL"}}; + for (const TypeCase& test_case : cases) { + const std::shared_ptr schema = + arrow::schema({arrow::field("value", test_case.type)}); + for (int32_t position : {0, 16, 32}) { + std::vector literals(33, test_case.valid); + literals[position] = Literal(int64_t{2}); + for (uint32_t maximum : {uint32_t{0}, uint32_t{32}, uint32_t{33}}) { + SCOPED_TRACE(fmt::format("type={}, maximum={}, position={}", + test_case.expected_type, maximum, position)); + for (bool negate : {false, true}) { + SCOPED_TRACE(negate); + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "value", test_case.field_type, literals) + : PredicateBuilder::In(0, "value", test_case.field_type, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, + {"value", fmt::format("expected {}", test_case.expected_type), + "actual BIGINT"}, + {maximum})); + } + } + } + } + + // File/predicate mismatches must also be independent of the membership maximum. + const std::shared_ptr bigint_schema = + arrow::schema({arrow::field("i", arrow::int64())}); + const std::vector literals(33, Literal(int32_t{1})); + for (uint32_t maximum : {uint32_t{0}, uint32_t{32}, uint32_t{33}}) { + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, bigint_schema, {"i", "expected BIGINT", "actual INT"}, {maximum})); + } + } +} + +TEST(VortexPredicateConverterTest, OversizedMembershipFallbacksDoNotHideTypeErrors) { + struct TypeCase { + std::shared_ptr type; + FieldType field_type; + Literal valid; + Literal first; + const char* expected_type; + }; + const std::vector cases = { + {arrow::int32(), FieldType::INT, Literal(int32_t{1}), Literal(FieldType::INT), "INT"}, + {arrow::float32(), FieldType::FLOAT, Literal(1.0f), Literal(1.0f), "FLOAT"}, + {arrow::float64(), FieldType::DOUBLE, Literal(1.0), Literal(1.0), "DOUBLE"}}; + for (const TypeCase& test_case : cases) { + const std::shared_ptr schema = + arrow::schema({arrow::field("value", test_case.type)}); + std::vector literals(33, test_case.valid); + literals.front() = test_case.first; + literals.back() = Literal(int64_t{2}); + for (uint32_t maximum : {uint32_t{0}, uint32_t{32}, uint32_t{33}}) { + SCOPED_TRACE(fmt::format("type={}, maximum={}", test_case.expected_type, maximum)); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "value", test_case.field_type, literals) + : PredicateBuilder::In(0, "value", test_case.field_type, literals); + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, + {"value", fmt::format("expected {}", test_case.expected_type), "actual BIGINT"}, + {maximum})); + } + } + } +} + +TEST(VortexPredicateConverterTest, MembershipTimestampValueChecksRemainWithinMaximum) { + const std::shared_ptr schema = + arrow::schema({arrow::field("value", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"))}); + const int64_t maximum = std::numeric_limits::max(); + const Literal native_out_of_range( + Timestamp::FromEpochMillis(maximum / 1000, static_cast((maximum % 1000) * 1000))); + const Literal precision_loss(Timestamp::FromEpochMillis(0, 1)); for (bool negate : {false, true}) { + for (bool unsafe_first : {false, true}) { + std::vector literals(33, native_out_of_range); + literals[unsafe_first ? 0 : 32] = precision_loss; + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "value", FieldType::TIMESTAMP, literals) + : PredicateBuilder::In(0, "value", FieldType::TIMESTAMP, literals); + for (uint32_t limit : {uint32_t{0}, uint32_t{32}, uint32_t{33}}) { + ASSERT_FALSE(Pushed(predicate, schema, {limit})); + } + } + const std::vector literals(33, native_out_of_range); const std::shared_ptr predicate = - negate ? PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals) - : PredicateBuilder::In(1, "s", FieldType::STRING, literals); - // ListContains lowers to 33 EQ + 32 OR nodes. NOT IN adds NOT plus the 3-node null guard. - const uint32_t nodes = 65 + (negate ? 4 : 0); + negate ? PredicateBuilder::NotIn(0, "value", FieldType::TIMESTAMP, literals) + : PredicateBuilder::In(0, "value", FieldType::TIMESTAMP, literals); ASSERT_FALSE(Pushed(predicate, schema)); - ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); - ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); - ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::FORCE, 0})); - ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::OFF, nodes})); - } - literals.erase(literals.begin() + 32, literals.end()); - ASSERT_TRUE(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, literals), schema)); - ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); - literals.erase(literals.begin() + 31, literals.end()); - ASSERT_FALSE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); - literals.erase(literals.begin() + 30, literals.end()); - ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); - literals.erase(literals.begin() + 21, literals.end()); - ASSERT_TRUE(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, literals), schema)); - ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals), schema)); -} - -TEST(VortexPredicateConverterTest, MembershipDeduplicatesBeforeApplyingBudget) { + ASSERT_FALSE(Pushed(predicate, schema, {33})); + } +} + +TEST(VortexPredicateConverterTest, MembershipMaximumCountsRawValuesForInAndNotIn) { const std::shared_ptr schema = TestSchema(); - std::vector strings(4096, Literal(FieldType::STRING, "same", 4)); - strings.emplace_back(FieldType::STRING, "other", 5); - const std::shared_ptr in = PredicateBuilder::In(1, "s", FieldType::STRING, strings); - const std::shared_ptr not_in = - PredicateBuilder::NotIn(1, "s", FieldType::STRING, strings); - ASSERT_FALSE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 2})); - ASSERT_TRUE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 3})); - ASSERT_FALSE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 6})); - ASSERT_TRUE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 7})); - - // Raw list length is not the decision: 4096 duplicates do not consume native nodes. - for (int32_t unique_count : {21, 32, 33}) { - std::vector repeated(4096, Literal(FieldType::STRING, "same", 4)); - for (int32_t value = 1; value < unique_count; ++value) { - const std::string text = fmt::format("value-{}", value); - repeated.emplace_back(FieldType::STRING, text.data(), text.size()); + for (uint32_t maximum : {uint32_t{1}, uint32_t{7}, uint32_t{32}, uint32_t{64}}) { + for (uint32_t size : {maximum - 1, maximum, maximum + 1}) { + SCOPED_TRACE(fmt::format("maximum={}, size={}", maximum, size)); + std::vector literals; + for (uint32_t value = 0; value < size; ++value) { + const std::string text = fmt::format("value-{}", value); + literals.emplace_back(FieldType::STRING, text.data(), text.size()); + } + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(1, "s", FieldType::STRING, literals) + : PredicateBuilder::In(1, "s", FieldType::STRING, literals); + ASSERT_EQ(Pushed(predicate, schema, {maximum}), size != 0 && size <= maximum); + ASSERT_EQ(Pushed(predicate, schema), size != 0 && size <= 32); + ASSERT_FALSE(Pushed(predicate, schema, {0})); + } } - ASSERT_EQ(Pushed(PredicateBuilder::In(1, "s", FieldType::STRING, repeated), schema), - unique_count <= 32); - ASSERT_EQ(Pushed(PredicateBuilder::NotIn(1, "s", FieldType::STRING, repeated), schema), - unique_count <= 30); } +} - const std::vector integers(4096, Literal(int32_t{7})); - ASSERT_TRUE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, integers), schema, - {MembershipPushdownMode::AUTO, 1})); - // The integer complement is one NEQ plus the 3-node null guard. - ASSERT_TRUE(Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, integers), schema, - {MembershipPushdownMode::AUTO, 4})); +TEST(VortexPredicateConverterTest, MembershipMaximumAppliesIndependentlyToEachLeaf) { + const std::shared_ptr schema = TestSchema(); + std::vector integers; + std::vector longs; + for (int32_t value = 0; value < 32; ++value) { + integers.emplace_back(value); + longs.emplace_back(static_cast(value)); + } + const auto in = PredicateBuilder::In(0, "i", FieldType::INT, integers); + const auto not_in = PredicateBuilder::NotIn(4, "l", FieldType::BIGINT, longs); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, + PredicateBuilder::And({in, not_in})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, + PredicateBuilder::Or({in, not_in})); + ASSERT_TRUE(Pushed(and_predicate, schema)); + ASSERT_TRUE(Pushed(or_predicate, schema)); + ASSERT_FALSE(Pushed(and_predicate, schema, {31})); + ASSERT_FALSE(Pushed(or_predicate, schema, {31})); } -TEST(VortexPredicateConverterTest, MembershipKeepsHugeContinuousSetsButBoundsManyRanges) { +TEST(VortexPredicateConverterTest, DuplicateValuesCountTowardMembershipMaximum) { + const std::shared_ptr schema = TestSchema(); + for (uint32_t size : {uint32_t{31}, uint32_t{32}, uint32_t{33}, uint32_t{4096}}) { + const std::vector strings(size, Literal(FieldType::STRING, "same", 4)); + const std::vector integers(size, Literal(int32_t{7})); + for (bool negate : {false, true}) { + const std::shared_ptr string_predicate = + negate ? PredicateBuilder::NotIn(1, "s", FieldType::STRING, strings) + : PredicateBuilder::In(1, "s", FieldType::STRING, strings); + const std::shared_ptr integer_predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, integers) + : PredicateBuilder::In(0, "i", FieldType::INT, integers); + ASSERT_EQ(Pushed(string_predicate, schema), size <= 32); + ASSERT_EQ(Pushed(integer_predicate, schema), size <= 32); + ASSERT_FALSE(Pushed(string_predicate, schema, {size - 1})); + ASSERT_FALSE(Pushed(integer_predicate, schema, {size - 1})); + if (size <= 33) { + ASSERT_TRUE(Pushed(string_predicate, schema, {size})); + ASSERT_TRUE(Pushed(integer_predicate, schema, {size})); + } + } + } +} + +TEST(VortexPredicateConverterTest, + HugeContinuousMembershipListsAreRejectedBeforeNativeConstruction) { const std::shared_ptr schema = TestSchema(); std::vector literals; for (int32_t value = 0; value < 65536; ++value) { literals.emplace_back(value); } - const std::shared_ptr in = PredicateBuilder::In(0, "i", FieldType::INT, literals); - const std::shared_ptr not_in = - PredicateBuilder::NotIn(0, "i", FieldType::INT, literals); - ASSERT_TRUE(Pushed(in, schema)); - ASSERT_TRUE(Pushed(not_in, schema)); - ASSERT_FALSE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 2})); - ASSERT_TRUE(Pushed(in, schema, {MembershipPushdownMode::AUTO, 3})); - ASSERT_FALSE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 5})); - ASSERT_TRUE(Pushed(not_in, schema, {MembershipPushdownMode::AUTO, 6})); - + // Do not construct a giant native expression even though the list is one continuous range. + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) + : PredicateBuilder::In(0, "i", FieldType::INT, literals); + ASSERT_FALSE(Pushed(predicate, schema)); + ASSERT_FALSE(Pushed(predicate, schema, {65535})); + ASSERT_FALSE(Pushed(predicate, schema, {0})); + } literals.clear(); for (int32_t start = 0; start < 200; start += 10) { for (int32_t value = start; value < start + 8; ++value) { @@ -1243,36 +2127,31 @@ TEST(VortexPredicateConverterTest, MembershipKeepsHugeContinuousSetsButBoundsMan const std::shared_ptr predicate = negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) : PredicateBuilder::In(0, "i", FieldType::INT, literals); - // 20 ranges, each 3 nodes, joined by 19 connectives; NOT IN adds the null guard. - const uint32_t nodes = 79 + (negate ? 3 : 0); ASSERT_FALSE(Pushed(predicate, schema)); - ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); - ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); - ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::FORCE, 0})); + ASSERT_FALSE(Pushed(predicate, schema, {159})); + ASSERT_TRUE(Pushed(predicate, schema, {160})); } } -TEST(VortexPredicateConverterTest, MembershipBudgetCountsMixedRangesAndSingletons) { +TEST(VortexPredicateConverterTest, MixedRangesAndSingletonsUseRawMembershipLength) { const std::shared_ptr schema = TestSchema(); std::vector literals = {Literal(int32_t{20}), Literal(int32_t{50})}; for (int32_t value = 1; value <= 10; ++value) { literals.emplace_back(value); literals.emplace_back(value + 29); } - // Two ranges + two singletons + three connectives = 11 nodes. for (bool negate : {false, true}) { const std::shared_ptr predicate = negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, literals) : PredicateBuilder::In(0, "i", FieldType::INT, literals); - const uint32_t nodes = 11 + (negate ? 3 : 0); - ASSERT_FALSE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes - 1})); - ASSERT_TRUE(Pushed(predicate, schema, {MembershipPushdownMode::AUTO, nodes})); + ASSERT_FALSE(Pushed(predicate, schema, {21})); + ASSERT_TRUE(Pushed(predicate, schema, {22})); } } -TEST(VortexPredicateConverterTest, ForceMembershipStillRejectsUnsafeLiterals) { +TEST(VortexPredicateConverterTest, LargerMembershipMaximumStillRejectsUnsafeLiterals) { const std::shared_ptr schema = TestSchema(); - const VortexPredicateOptions force{MembershipPushdownMode::FORCE, 0}; + const VortexPredicateOptions larger_maximum{2048}; std::vector literals; for (int32_t value = 0; value < 1024; ++value) { literals.emplace_back(value); @@ -1280,12 +2159,22 @@ TEST(VortexPredicateConverterTest, ForceMembershipStillRejectsUnsafeLiterals) { for (const Literal& unsafe : {Literal(FieldType::INT), Literal(int64_t{1024})}) { std::vector invalid = literals; invalid.push_back(unsafe); - ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, invalid), schema, force)); - ASSERT_FALSE( - Pushed(PredicateBuilder::NotIn(0, "i", FieldType::INT, invalid), schema, force)); + for (bool negate : {false, true}) { + const std::shared_ptr predicate = + negate ? PredicateBuilder::NotIn(0, "i", FieldType::INT, invalid) + : PredicateBuilder::In(0, "i", FieldType::INT, invalid); + if (unsafe.IsNull()) { + ASSERT_FALSE(Pushed(predicate, schema, larger_maximum)); + } else { + ASSERT_NO_FATAL_FAILURE(AssertInvalidConversion( + predicate, schema, {"i", "expected INT", "actual BIGINT"}, larger_maximum)); + } + } } - ASSERT_FALSE(Pushed(PredicateBuilder::In(0, "i", FieldType::INT, literals), - arrow::schema({arrow::field("i", arrow::int64())}), force)); + ASSERT_NO_FATAL_FAILURE( + AssertInvalidConversion(PredicateBuilder::In(0, "i", FieldType::INT, literals), + arrow::schema({arrow::field("i", arrow::int64())}), + {"i", "expected BIGINT", "actual INT"}, larger_maximum)); for (bool negate : {false, true}) { const std::vector timestamps = { Literal(Timestamp::FromEpochMillis(0)), @@ -1295,13 +2184,13 @@ TEST(VortexPredicateConverterTest, ForceMembershipStillRejectsUnsafeLiterals) { : PredicateBuilder::In(0, "ts", FieldType::TIMESTAMP, timestamps); ASSERT_FALSE(Pushed( predicate, arrow::schema({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::NANO))}), - force)); + larger_maximum)); } } -TEST(VortexPredicateConverterTest, MembershipOffPreservesOtherNativePredicates) { +TEST(VortexPredicateConverterTest, ZeroMembershipMaximumPreservesOtherNativePredicates) { const std::shared_ptr schema = TestSchema(); - const VortexPredicateOptions off{MembershipPushdownMode::OFF, 0}; + const VortexPredicateOptions off{0}; const std::shared_ptr equal = PredicateBuilder::Equal(0, "i", FieldType::INT, Literal(int32_t{1})); const std::shared_ptr in = @@ -1319,7 +2208,7 @@ TEST(VortexPredicateConverterTest, MembershipOffPreservesOtherNativePredicates) ASSERT_FALSE(Pushed(or_predicate, schema, off)); } -TEST(VortexPredicateConverterTest, SparseMembershipModesReturnNativeOrCandidateRows) { +TEST(VortexPredicateConverterTest, SparseMembershipMaximumReturnsNativeOrCandidateRows) { NativeInCase test_case = {"sparse_budget", arrow::int32(), FieldType::INT, "[0,null,63,2]", {}}; for (int32_t value = 0; value <= 64; value += 2) { test_case.literals.emplace_back(value); @@ -1331,15 +2220,15 @@ TEST(VortexPredicateConverterTest, SparseMembershipModesReturnNativeOrCandidateR const std::vector expected = negate ? std::vector{2} : std::vector{0, 3}; ASSERT_NO_FATAL_FAILURE( - AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, true, {}, false)); - ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, predicate, expected, true, - {MembershipPushdownMode::FORCE, 0})); - ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, predicate, {0, 1, 2, 3}, true, - {MembershipPushdownMode::OFF}, false)); + AssertNativeAndResidualRows(test_case, predicate, {0, 1, 2, 3}, expected, {}, false)); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case, predicate, expected, expected, {33})); + ASSERT_NO_FATAL_FAILURE( + AssertNativeAndResidualRows(test_case, predicate, {0, 1, 2, 3}, expected, {0}, false)); } } -TEST(VortexPredicateConverterTest, StringMembershipModesAndPartialAndOrUseSafeCandidates) { +TEST(VortexPredicateConverterTest, StringMembershipMaximumAndPartialAndOrUseSafeCandidates) { NativeInCase test_case = { "string_budget", arrow::utf8(), FieldType::STRING, R"(["s0",null,"absent","s32"])", {}}; for (int32_t value = 0; value < 33; ++value) { @@ -1348,9 +2237,9 @@ TEST(VortexPredicateConverterTest, StringMembershipModesAndPartialAndOrUseSafeCa } const std::shared_ptr in = PredicateBuilder::In(1, "value", FieldType::STRING, test_case.literals); - ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, in, {0, 1, 2, 3}, true, {}, false)); ASSERT_NO_FATAL_FAILURE( - AssertNativeRows(test_case, in, {0, 3}, true, {MembershipPushdownMode::FORCE, 0})); + AssertNativeAndResidualRows(test_case, in, {0, 1, 2, 3}, {0, 3}, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, in, {0, 3}, {0, 3}, {33})); const std::shared_ptr less_than_three = PredicateBuilder::LessThan(0, "rowid", FieldType::BIGINT, Literal(int64_t{3})); ASSERT_OK_AND_ASSIGN(std::shared_ptr and_predicate, @@ -1358,10 +2247,10 @@ TEST(VortexPredicateConverterTest, StringMembershipModesAndPartialAndOrUseSafeCa ASSERT_OK_AND_ASSIGN(std::shared_ptr or_predicate, PredicateBuilder::Or({in, less_than_three})); // AND keeps the other child, so the residual receives a safe superset, including null rows. - ASSERT_NO_FATAL_FAILURE(AssertNativeRows(test_case, and_predicate, {0, 1, 2}, true)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, and_predicate, {0, 1, 2}, {0})); // Dropping only the IN child of OR would lose row 3; reject the entire native OR instead. - ASSERT_NO_FATAL_FAILURE( - AssertNativeRows(test_case, or_predicate, {0, 1, 2, 3}, true, {}, false)); + ASSERT_NO_FATAL_FAILURE(AssertNativeAndResidualRows(test_case, or_predicate, {0, 1, 2, 3}, + {0, 1, 2, 3}, {}, false)); } TEST(VortexPredicateConverterTest, NullPredicateIsNotPushed) { diff --git a/src/paimon/format/vortex/vortex_reader_builder.h b/src/paimon/format/vortex/vortex_reader_builder.h index 74f69ec9d..f655cdcdd 100644 --- a/src/paimon/format/vortex/vortex_reader_builder.h +++ b/src/paimon/format/vortex/vortex_reader_builder.h @@ -23,6 +23,7 @@ #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/format/reader_builder.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_file_batch_reader.h" #include "paimon/format/vortex/vortex_predicate_converter.h" #include "paimon/memory/memory_pool.h" @@ -31,12 +32,12 @@ namespace paimon::vortex { class VortexReaderBuilder : public ReaderBuilder { public: - explicit VortexReaderBuilder( - int32_t batch_size, const VortexPredicateOptions& predicate_options = {}, - VortexIoCoalesceMode coalesce_mode = VortexIoCoalesceMode::OBJECT_STORAGE) + explicit VortexReaderBuilder(int32_t batch_size, + const VortexPredicateOptions& predicate_options, + const VortexIoCoalesceOptions& coalesce_options) : batch_size_(batch_size), predicate_options_(predicate_options), - coalesce_mode_(coalesce_mode), + coalesce_options_(coalesce_options), pool_(GetDefaultPool()), arrow_pool_(GetArrowPool(pool_)) {} @@ -55,14 +56,14 @@ class VortexReaderBuilder : public ReaderBuilder { if (pool_ == nullptr) { return Status::Invalid("Vortex reader memory pool is nullptr"); } - return VortexFileBatchReader::Create(input, batch_size_, arrow_pool_, predicate_options_, - coalesce_mode_); + return VortexFileBatchReader::Create(input, batch_size_, predicate_options_, + coalesce_options_, arrow_pool_); } private: int32_t batch_size_; VortexPredicateOptions predicate_options_; - VortexIoCoalesceMode coalesce_mode_; + VortexIoCoalesceOptions coalesce_options_; std::shared_ptr pool_; std::shared_ptr arrow_pool_; }; diff --git a/src/paimon/format/vortex/vortex_stats_extractor.cpp b/src/paimon/format/vortex/vortex_stats_extractor.cpp index c589e7c29..29ed43d90 100644 --- a/src/paimon/format/vortex/vortex_stats_extractor.cpp +++ b/src/paimon/format/vortex/vortex_stats_extractor.cpp @@ -28,6 +28,7 @@ #include "paimon/common/utils/math.h" #include "paimon/defs.h" #include "paimon/format/column_stats.h" +#include "paimon/format/vortex/vortex_defs.h" #include "paimon/format/vortex/vortex_ffi.h" #include "paimon/format/vortex/vortex_ffi_util.h" #include "paimon/format/vortex/vortex_io_callbacks.h" @@ -59,8 +60,9 @@ VortexStatsExtractor::ExtractWithFileInfo(const std::shared_ptr& fil } vx_error* error = nullptr; VxDataSourcePtr data_source( - vx_data_source_new_callback(session.get(), VortexInputContext::MakeCallbacks(input_context), - length, &error), + vx_data_source_new_callback_with_coalesce( + session.get(), VortexInputContext::MakeCallbacks(input_context), length, + kDefaultVortexReadIoCoalesceMaxGap, kDefaultVortexReadIoCoalesceMaxSize, &error), vx_data_source_free); if (data_source == nullptr) { return VortexCallbackError("open Vortex file for stats", error, diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 5de6c9fa7..33412e5f8 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -45,6 +45,7 @@ #include "paimon/catalog/identifier.h" #include "paimon/commit_context.h" #include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/reader/reader_utils.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" @@ -4968,6 +4969,169 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { ASSERT_OK(PrepareAndClose(writer.get())); } +TEST_F(RealtimeWriteInteTest, TestPredicateFilteringForMemoryAndSpilledCandidates) { + fields_ = {arrow::field("id", arrow::int64()), arrow::field("value", arrow::float64()), + arrow::field("payload", arrow::utf8())}; + schema_ = arrow::schema(fields_); + options_[Options::REALTIME_SPILL_ENABLED] = "true"; + options_[Options::REALTIME_STORE_STATS_MODE] = "full"; + options_[Options::READ_BATCH_SIZE] = "2"; + CreateTable(/*partition_keys=*/{}); + const std::string temp_directory = PathUtil::JoinPath(dir_->Str(), "late-filter-spill"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context, temp_directory)); + // Keep matching and non-matching rows in one stored batch: statistics may retain the whole + // batch, but exact filtering must evaluate individual rows after the read pipeline rebuilds it. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeUnpartitionedBatchFromJson( + R"([[0,0,-1,"minus"],[1,1,0,"zero"],[2,2,1,"one"],[3,3,null,"null"]])")); + ASSERT_OK(writer->Write(std::move(batch))); + + struct PredicateCase { + std::string field_name; + std::shared_ptr predicate; + std::shared_ptr type; + std::string expected; + std::string expected_with_payload; + }; + std::vector values = {Literal(int64_t{2})}; + for (int64_t value = 100; value < 132; ++value) { + values.emplace_back(value); + } + const std::vector cases = { + {"value", PredicateBuilder::GreaterThan(0, "value", FieldType::DOUBLE, Literal(0.0)), + arrow::float64(), R"([[0,1]])", R"([[0,1,"one"]])"}, + {"id", PredicateBuilder::In(0, "id", FieldType::BIGINT, values), arrow::int64(), + R"([[0,2]])", R"([[0,2,"one"]])"}, + {"id", PredicateBuilder::NotIn(0, "id", FieldType::BIGINT, values), arrow::int64(), + R"([[0,0],[0,1],[0,3]])", R"([[0,0,"minus"],[0,1,"zero"],[0,3,"null"]])"}}; + + double memory_before_spill = 0; + for (bool spilled : {false, true}) { + SCOPED_TRACE(spilled); + if (spilled) { + ASSERT_OK(writer->Seal()); + ASSERT_EQ(1, TestHelper::CountChannelFiles(dir_->GetFileSystem(), temp_directory)); + } + ASSERT_OK_AND_ASSIGN(RealtimeMetricValues metrics, + ReadRealtimeMetrics(writer->GetMetrics())); + ASSERT_EQ(4, metrics.total_rows); + if (spilled) { + // Full batch statistics stay resident after spilling the Arrow payload to IPC. + ASSERT_EQ(0, metrics.building_memory); + ASSERT_EQ(0, metrics.building_rows); + ASSERT_EQ(4, metrics.sealed_rows); + ASSERT_GT(metrics.sealed_memory, 0); + ASSERT_EQ(metrics.sealed_memory, metrics.total_memory); + ASSERT_LT(metrics.total_memory, memory_before_spill); + } else { + ASSERT_GT(metrics.building_memory, 0); + ASSERT_EQ(4, metrics.building_rows); + ASSERT_EQ(0, metrics.sealed_rows); + memory_before_spill = metrics.building_memory; + } + + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.predicate->ToString()); + for (bool with_payload : {false, true}) { + SCOPED_TRACE(with_payload); + std::vector projection = {test_case.field_name}; + arrow::FieldVector result_fields = { + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field(test_case.field_name, test_case.type)}; + if (with_payload) { + projection.push_back("payload"); + result_fields.push_back(fields_[2]); + } + const auto result_type = arrow::struct_(result_fields); + const std::string candidates = + test_case.field_name == "value" + ? (with_payload + ? R"([[0,-1,"minus"],[0,0,"zero"],[0,1,"one"],[0,null,"null"]])" + : R"([[0,-1],[0,0],[0,1],[0,null]])") + : (with_payload ? R"([[0,0,"minus"],[0,1,"zero"],[0,2,"one"],[0,3,"null"]])" + : R"([[0,0],[0,1],[0,2],[0,3]])"); + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (bool use_bitmap : {false, true}) { + SCOPED_TRACE(use_bitmap); + // Each read consumes its own real-time read-view ticket. + ASSERT_OK_AND_ASSIGN( + std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + auto split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + ASSERT_TRUE(split->DiskSplits().empty()); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames(projection) + .SetPredicate(test_case.predicate) + .EnablePredicateFilter(exact_filter) + .EnableLateMaterializing(late_materializing) + .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())); + std::shared_ptr result; + if (!use_bitmap) { + ASSERT_OK_AND_ASSIGN( + result, ReadResultCollector::CollectResult(std::move(reader))); + } else { + ScopeGuard close_reader([&reader]() { reader->Close(); }); + auto arrow_pool = std::shared_ptr( + arrow::default_memory_pool(), [](arrow::MemoryPool*) {}); + arrow::ArrayVector arrays; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap next, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(next)) { + break; + } + ASSERT_OK(ReadResultCollector::CheckBatchOffset(next.first)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch filtered, + ReaderUtils::ApplyBitmapToReadBatch( + std::move(next), arrow_pool)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr array, + ReadResultCollector::GetArray(std::move(filtered))); + arrays.push_back(std::move(array)); + } + auto chunked = arrow::ChunkedArray::Make(std::move(arrays)); + ASSERT_TRUE(chunked.ok()) << chunked.status().ToString(); + result = std::move(chunked).ValueOrDie(); + } + // Memory and IPC candidates do not pass through the file probe. + const std::string& expected_json = + exact_filter ? (with_payload ? test_case.expected_with_payload + : test_case.expected) + : candidates; + auto expected = arrow::ipc::internal::json::ArrayFromJSON(result_type, + expected_json) + .ValueOrDie(); + ASSERT_NE(nullptr, result); + ASSERT_TRUE( + std::make_shared(expected)->Equals(*result)) + << result->ToString(); + } + } + } + } + } + } + ASSERT_OK(PrepareAndClose(writer.get())); +} + TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { options_[Options::FILE_FORMAT] = "parquet"; options_[Options::WRITE_BATCH_SIZE] = "1"; diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index bed173fa4..6b3876639 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -38,8 +39,10 @@ #include "paimon/bucket/bucket_id_calculator.h" #include "paimon/common/factories/io_hook.h" #include "paimon/common/io/cache/lru_cache.h" +#include "paimon/common/reader/reader_utils.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" @@ -49,6 +52,8 @@ #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/deletion_file.h" #include "paimon/core/table/source/key_value_table_read.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/executor.h" #include "paimon/fs/file_system.h" @@ -3732,33 +3737,55 @@ namespace { Result> ReadVortexSdkRows( const std::string& table_path, const std::vector>& splits, const std::vector& projection, const std::shared_ptr& predicate, - bool exact_filter, bool late_materializing) { + bool exact_filter, bool late_materializing, bool use_bitmap = false) { ReadContextBuilder builder(table_path); builder.SetReadFieldNames(projection) .SetPredicate(predicate) .EnablePredicateFilter(exact_filter) .EnableLateMaterializing(late_materializing) - .EnablePrefetch(true) - .SetPrefetchBatchCount(2) - .SetPrefetchMaxParallelNum(2) .AddOption(Options::READ_BATCH_SIZE, "2"); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read, TableRead::Create(std::move(context))); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, read->CreateReader(splits)); - return ReadResultCollector::CollectResult(std::move(reader)); + if (!use_bitmap) { + return ReadResultCollector::CollectResult(std::move(reader)); + } + ScopeGuard close_reader([&reader]() { reader->Close(); }); + auto arrow_pool = + std::shared_ptr(arrow::default_memory_pool(), [](arrow::MemoryPool*) {}); + arrow::ArrayVector arrays; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_RETURN_NOT_OK(ReadResultCollector::CheckBatchOffset(batch.first)); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch filtered, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), arrow_pool)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr array, + ReadResultCollector::GetArray(std::move(filtered))); + arrays.push_back(std::move(array)); + } + if (arrays.empty()) { + return std::shared_ptr(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::ChunkedArray::Make(std::move(arrays))); + return result; } // Append splits need not have a stable file order. Compare full rows as a multiset. void AssertVortexSdkRows(const std::shared_ptr& actual, const std::shared_ptr& type, - const std::string& expected_json, bool exact_filter = true) { + const std::string& expected_json, bool expect_exact = true) { auto expected = arrow::ipc::internal::json::ArrayFromJSON(type, expected_json).ValueOrDie(); if (!actual) { ASSERT_EQ(expected->length(), 0); return; } ASSERT_TRUE(actual->type()->Equals(type)) << actual->type()->ToString(); - if (exact_filter) { + if (expect_exact) { ASSERT_EQ(actual->length(), expected->length()) << actual->ToString(); } else { ASSERT_GE(actual->length(), expected->length()); @@ -3818,7 +3845,7 @@ TEST_F(ScanAndReadInteTest, TestVortexSdkPredicatesWithResidualAndLateMaterializ PredicateBuilder::Equal(2, "tag", FieldType::STRING, Literal(FieldType::STRING, "keep", 4)); ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, PredicateBuilder::And({greater, tag})); - // Sparse IN(64) exceeds the auto expression budget and requires the SDK residual filter. + // IN(64) exceeds the default 32-value pushdown limit; exact results require SDK filtering. std::vector sparse = {Literal(int32_t{2}), Literal(int32_t{5})}; for (int32_t index = 0; index < 62; ++index) { sparse.emplace_back(1000 + index * 2); @@ -3839,9 +3866,265 @@ TEST_F(ScanAndReadInteTest, TestVortexSdkPredicatesWithResidualAndLateMaterializ std::shared_ptr actual, ReadVortexSdkRows(table_path, splits, {"payload", "id", "tag"}, predicate, exact_filter, late_materializing)); - // A non-exact reader may prune candidates, but it must retain every matching row. + // Every predicate leaves a payload column, so the late probe also filters exactly. + const bool expect_exact = exact_filter || late_materializing; + ASSERT_NO_FATAL_FAILURE( + AssertVortexSdkRows(actual, result_type, expected, expect_exact)); + } + } + } +} + +TEST_F(ScanAndReadInteTest, TestVortexSdkFallbackPredicatesWithoutPayload) { + for (bool floating_point : {true, false}) { + SCOPED_TRACE(floating_point); + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector fields = { + arrow::field("value", floating_point ? arrow::float64() : arrow::int32()), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "-1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), + R"([[-1,"minus"],[0,"zero"],[1,"one"],[null,"null"]])", {}, + 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 0, std::nullopt)); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + std::vector, std::string, std::string>> cases; + if (floating_point) { + // Floating-point comparisons cannot be pushed down safely for every NaN encoding. + cases.emplace_back( + PredicateBuilder::GreaterThan(0, "value", FieldType::DOUBLE, Literal(0.0)), + R"([[0,1]])", R"([[0,1,"one"]])"); + } else { + std::vector values = {Literal(int32_t{1})}; + for (int32_t value = 100; value < 132; ++value) { + values.emplace_back(value); + } + // Both lists exceed the default 32-value native pushdown limit. + cases.emplace_back(PredicateBuilder::In(0, "value", FieldType::INT, values), + R"([[0,1]])", R"([[0,1,"one"]])"); + cases.emplace_back(PredicateBuilder::NotIn(0, "value", FieldType::INT, values), + R"([[0,-1],[0,0]])", R"([[0,-1,"minus"],[0,0,"zero"]])"); + } + for (bool with_payload : {false, true}) { + SCOPED_TRACE(with_payload); + std::vector projection = {"value"}; + arrow::FieldVector result_fields = {arrow::field("_VALUE_KIND", arrow::int8()), + fields[0]}; + if (with_payload) { + projection.push_back("payload"); + result_fields.push_back(fields[1]); + } + const auto result_type = arrow::struct_(result_fields); + for (const auto& [predicate, without_payload, with_payload_result] : cases) { + SCOPED_TRACE(predicate->ToString()); + const std::string& expected = with_payload ? with_payload_result : without_payload; + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (bool use_bitmap : {false, true}) { + SCOPED_TRACE(use_bitmap); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, projection, predicate, + exact_filter, late_materializing, use_bitmap)); + // A probe-only projection bypasses the two-phase late reader. + const bool expect_exact = + exact_filter || (late_materializing && with_payload); + ASSERT_NO_FATAL_FAILURE( + AssertVortexSdkRows(actual, result_type, expected, expect_exact)); + } + } + } + } + } + } +} + +TEST_F(ScanAndReadInteTest, TestVortexSdkTemporalRangeFallbacksAreExact) { + struct TemporalCase { + const char* name; + std::shared_ptr type; + FieldType field_type; + Literal extreme; + Literal zero; + }; + const std::vector types = { + {"date", arrow::date32(), FieldType::DATE, + Literal(FieldType::DATE, std::numeric_limits::max()), + Literal(FieldType::DATE, 0)}, + {"timestamp", arrow::timestamp(arrow::TimeUnit::MICRO), FieldType::TIMESTAMP, + Literal(Timestamp::FromEpochMillis(253402300799999, 999000)), + Literal(Timestamp::FromEpochMillis(0))}, + {"timestamp-ltz", + arrow::timestamp(arrow::TimeUnit::MICRO, DateTimeUtils::GetLocalTimezoneName()), + FieldType::TIMESTAMP, Literal(Timestamp::FromEpochMillis(253402300799999, 999000)), + Literal(Timestamp::FromEpochMillis(0))}}; + for (const TemporalCase& item : types) { + SCOPED_TRACE(item.name); + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector fields = {arrow::field("value", item.type), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "-1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([[-1000,"minus"],[0,"zero"],[1000,"plus"],[null,"null"]])", {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 0, std::nullopt)); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + const auto less = PredicateBuilder::LessThan(0, "value", item.field_type, item.extreme); + const auto equal = PredicateBuilder::Equal(0, "value", item.field_type, item.zero); + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({less, equal})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({equal, less})); + const std::string all = R"([[0,-1000],[0,0],[0,1000]])"; + const std::string all_with_payload = R"([[0,-1000,"minus"],[0,0,"zero"],[0,1000,"plus"]])"; + const std::vector, std::string, std::string>> cases = + {{less, all, all_with_payload}, + {PredicateBuilder::GreaterThan(0, "value", item.field_type, item.extreme), "[]", "[]"}, + {PredicateBuilder::In(0, "value", item.field_type, {item.extreme, item.zero}), + R"([[0,0]])", R"([[0,0,"zero"]])"}, + {PredicateBuilder::NotIn(0, "value", item.field_type, {item.zero, item.extreme}), + R"([[0,-1000],[0,1000]])", R"([[0,-1000,"minus"],[0,1000,"plus"]])"}, + {conjunction, R"([[0,0]])", R"([[0,0,"zero"]])"}, + {disjunction, all, all_with_payload}}; + for (bool with_payload : {false, true}) { + SCOPED_TRACE(with_payload); + std::vector projection = {"value"}; + arrow::FieldVector result_fields = {arrow::field("_VALUE_KIND", arrow::int8()), + fields[0]}; + if (with_payload) { + projection.push_back("payload"); + result_fields.push_back(fields[1]); + } + const auto result_type = arrow::struct_(result_fields); + for (const auto& [predicate, without_payload, with_payload_result] : cases) { + SCOPED_TRACE(predicate->GetFunction().ToString()); + const std::string& expected = with_payload ? with_payload_result : without_payload; + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (bool use_bitmap : {false, true}) { + SCOPED_TRACE(use_bitmap); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, projection, predicate, + exact_filter, late_materializing, use_bitmap)); + // A probe-only projection bypasses the two-phase late reader. + const bool expect_exact = + exact_filter || (late_materializing && with_payload); + ASSERT_NO_FATAL_FAILURE( + AssertVortexSdkRows(actual, result_type, expected, expect_exact)); + } + } + } + } + } + } +} + +TEST_F(ScanAndReadInteTest, TestVortexSdkDecimalPredicatesWithProjectionAndResidualFiltering) { + auto dir = UniqueTestDirectory::Create("local"); + ASSERT_NE(dir, nullptr); + const arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("amount", arrow::decimal128(10, 2)), + arrow::field("payload", arrow::utf8())}; + const std::map options = {{Options::FILE_FORMAT, "vortex"}, + {Options::FILE_SYSTEM, "local"}, + {Options::BUCKET, "-1"}, + {Options::WRITE_ONLY, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, + TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {}, options, + /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), R"([ + [0,"-4.56","zero"],[1,"-1.23","one"],[2,"0.00","two"], + [3,"1.20","three"],[4,"1.23","four"],[5,"1.24","five"], + [6,"4.56","six"],[7,null,"seven"],[8,"99999999.99","eight"] + ])", + {}, 0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), 0, std::nullopt)); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), std::nullopt, false)); + const auto result_type = + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), fields[0], fields[1]}); + const auto decimal = [](int32_t precision, int64_t unscaled) { + return Literal(Decimal(precision, 2, unscaled)); + }; + const std::vector exact_values = {decimal(20, -123), decimal(17, 120), + decimal(19, 456)}; + // The SDK requires literals to have the column scale and fit its precision. Different + // precision metadata is valid; rescaling and overflow are tested at the format boundary. + std::vector large_values = {decimal(19, 123), decimal(20, 456)}; + for (int32_t index = 0; index < 31; ++index) { + large_values.push_back(decimal(19, 10000 + index * 2)); + } + const std::vector, std::string>> cases = { + {PredicateBuilder::Equal(1, "amount", FieldType::DECIMAL, decimal(19, 123)), + R"([[0,4,"1.23"]])"}, + {PredicateBuilder::NotEqual(1, "amount", FieldType::DECIMAL, decimal(20, 123)), + R"([[0,0,"-4.56"],[0,1,"-1.23"],[0,2,"0.00"],[0,3,"1.20"], + [0,5,"1.24"],[0,6,"4.56"],[0,8,"99999999.99"]])"}, + {PredicateBuilder::LessThan(1, "amount", FieldType::DECIMAL, decimal(20, 123)), + R"([[0,0,"-4.56"],[0,1,"-1.23"],[0,2,"0.00"],[0,3,"1.20"]])"}, + {PredicateBuilder::LessOrEqual(1, "amount", FieldType::DECIMAL, decimal(20, 123)), + R"([[0,0,"-4.56"],[0,1,"-1.23"],[0,2,"0.00"],[0,3,"1.20"],[0,4,"1.23"]])"}, + {PredicateBuilder::GreaterThan(1, "amount", FieldType::DECIMAL, decimal(20, 123)), + R"([[0,5,"1.24"],[0,6,"4.56"],[0,8,"99999999.99"]])"}, + {PredicateBuilder::GreaterOrEqual(1, "amount", FieldType::DECIMAL, decimal(19, 123)), + R"([[0,4,"1.23"],[0,5,"1.24"],[0,6,"4.56"],[0,8,"99999999.99"]])"}, + {PredicateBuilder::Equal(1, "amount", FieldType::DECIMAL, decimal(9, -123)), + R"([[0,1,"-1.23"]])"}, + {PredicateBuilder::Equal(1, "amount", FieldType::DECIMAL, decimal(38, 9999999999)), + R"([[0,8,"99999999.99"]])"}, + {PredicateBuilder::In(1, "amount", FieldType::DECIMAL, exact_values), + R"([[0,1,"-1.23"],[0,3,"1.20"],[0,6,"4.56"]])"}, + {PredicateBuilder::NotIn(1, "amount", FieldType::DECIMAL, exact_values), + R"([[0,0,"-4.56"],[0,2,"0.00"],[0,4,"1.23"],[0,5,"1.24"],[0,8,"99999999.99"]])"}, + // IN(33) and NOT IN(33) exceed the default limit and require residual filtering. + {PredicateBuilder::In(1, "amount", FieldType::DECIMAL, large_values), + R"([[0,4,"1.23"],[0,6,"4.56"]])"}, + {PredicateBuilder::NotIn(1, "amount", FieldType::DECIMAL, large_values), + R"([[0,0,"-4.56"],[0,1,"-1.23"],[0,2,"0.00"],[0,3,"1.20"], + [0,5,"1.24"],[0,8,"99999999.99"]])"}}; + for (bool late_materializing : {false, true}) { + SCOPED_TRACE(late_materializing); + for (const auto& [predicate, expected] : cases) { + SCOPED_TRACE(predicate->ToString()); + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, {"id", "amount"}, predicate, exact_filter, + late_materializing)); + // The id column is the late-materialization payload for predicates on amount. + const bool expect_exact = exact_filter || late_materializing; ASSERT_NO_FATAL_FAILURE( - AssertVortexSdkRows(actual, result_type, expected, exact_filter)); + AssertVortexSdkRows(actual, result_type, expected, expect_exact)); } } } @@ -3860,17 +4143,19 @@ TEST_F(ScanAndReadInteTest, TestVortexSdkPrimaryKeyMergeBeforeValueFiltering) { ASSERT_OK_AND_ASSIGN(std::unique_ptr helper, TestHelper::Create(dir->Str(), arrow::schema(fields), {}, {"pk"}, options, /*is_streaming_mode=*/true)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr original, - TestHelper::MakeRecordBatch( - arrow::struct_(fields), - R"([["a",10,"old-a"],["b",20,"old-b"],["c",30,"old-c"]])", {}, 0, {})); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr original, + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([["a",10,"old-a"],["b",20,"old-b"],["c",30,"old-c"],["e",1,"old-e"]])", {}, 0, {})); ASSERT_OK(helper->WriteAndCommit(std::move(original), 0, std::nullopt)); ASSERT_OK_AND_ASSIGN( std::unique_ptr changes, - TestHelper::MakeRecordBatch(arrow::struct_(fields), - R"([["a",1,"new-a"],["b",20,"old-b"],["d",40,"new-d"]])", {}, 0, - {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, - RecordBatch::RowKind::INSERT})); + TestHelper::MakeRecordBatch( + arrow::struct_(fields), + R"([["a",1,"new-a"],["b",20,"old-b"],["d",40,"new-d"],["e",20,"new-e"]])", {}, 0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT})); ASSERT_OK(helper->WriteAndCommit(std::move(changes), 1, std::nullopt)); const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); ASSERT_OK_AND_ASSIGN(std::vector> splits, @@ -3884,18 +4169,52 @@ TEST_F(ScanAndReadInteTest, TestVortexSdkPrimaryKeyMergeBeforeValueFiltering) { Literal(FieldType::STRING, "c", 1)}); ASSERT_OK_AND_ASSIGN(std::shared_ptr both, PredicateBuilder::And({key, value})); const std::vector, std::string>> cases = { - {nullptr, R"([[0,"a",1,"new-a"],[0,"c",30,"old-c"],[0,"d",40,"new-d"]])"}, - {value, R"([[0,"c",30,"old-c"],[0,"d",40,"new-d"]])"}, + {nullptr, R"([[0,"a",1,"new-a"],[0,"c",30,"old-c"],[0,"d",40,"new-d"], + [0,"e",20,"new-e"]])"}, + {value, R"([[0,"c",30,"old-c"],[0,"d",40,"new-d"],[0,"e",20,"new-e"]])"}, {key, R"([[0,"a",1,"new-a"],[0,"c",30,"old-c"]])"}, {both, R"([[0,"c",30,"old-c"]])"}}; + const std::shared_ptr merged_rows = + arrow::ipc::internal::json::ArrayFromJSON(result_type, cases.front().second).ValueOrDie(); for (bool late_materializing : {false, true}) { SCOPED_TRACE(late_materializing); - for (const auto& [predicate, expected] : cases) { - ASSERT_OK_AND_ASSIGN( - std::shared_ptr actual, - ReadVortexSdkRows(table_path, splits, {"pk", "value", "payload"}, predicate, - /*exact_filter=*/true, late_materializing)); - ASSERT_NO_FATAL_FAILURE(AssertVortexSdkRows(actual, result_type, expected)); + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + for (const auto& [predicate, expected] : cases) { + for (bool use_bitmap : {false, true}) { + SCOPED_TRACE(use_bitmap); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, {"pk", "value", "payload"}, predicate, + exact_filter, late_materializing, use_bitmap)); + // Candidate filtering must not resurrect old versions or deleted keys. + if (actual != nullptr) { + ASSERT_TRUE(actual->type()->Equals(result_type)); + std::vector used(static_cast(merged_rows->length()), false); + for (const auto& chunk : actual->chunks()) { + ASSERT_TRUE(chunk->ValidateFull().ok()); + for (int64_t row = 0; row < chunk->length(); ++row) { + bool found = false; + for (int64_t merged_row = 0; merged_row < merged_rows->length(); + ++merged_row) { + const auto index = static_cast(merged_row); + if (!used[index] && chunk->Slice(row, 1)->Equals( + merged_rows->Slice(merged_row, 1))) { + used[index] = true; + found = true; + break; + } + } + ASSERT_TRUE(found) << "Unexpected or duplicate merged row: " + << chunk->Slice(row, 1)->ToString(); + } + } + } + // Value predicates may be withheld until after primary-key merging. + ASSERT_NO_FATAL_FAILURE(AssertVortexSdkRows( + actual, result_type, expected, predicate == nullptr || exact_filter)); + } + } } } } @@ -3944,12 +4263,22 @@ TEST_F(ScanAndReadInteTest, TestVortexSdkSchemaEvolutionWithMissingRowIndexColum [0,2,null,null,null]])"}}; for (bool late_materializing : {false, true}) { SCOPED_TRACE(late_materializing); - for (const auto& [predicate, expected] : cases) { - ASSERT_OK_AND_ASSIGN( - std::shared_ptr actual, - ReadVortexSdkRows(table_path, splits, {"id", "payload", "added", "_row_idx"}, - predicate, /*exact_filter=*/true, late_materializing)); - ASSERT_NO_FATAL_FAILURE(AssertVortexSdkRows(actual, result_type, expected)); + for (bool exact_filter : {false, true}) { + SCOPED_TRACE(exact_filter); + for (const auto& [predicate, expected] : cases) { + SCOPED_TRACE(predicate->ToString()); + for (bool use_bitmap : {false, true}) { + SCOPED_TRACE(use_bitmap); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadVortexSdkRows(table_path, splits, + {"id", "payload", "added", "_row_idx"}, predicate, + exact_filter, late_materializing, use_bitmap)); + // Missing file fields can leave the late reader without the full predicate. + ASSERT_NO_FATAL_FAILURE( + AssertVortexSdkRows(actual, result_type, expected, exact_filter)); + } + } } } } From 6bc9835faa6bd9ec8183856654804301469d18b7 Mon Sep 17 00:00:00 2001 From: zhouyc-ali <301200591+zhouyc-ali@users.noreply.github.com> Date: Thu, 5 Nov 2026 06:23:30 +0000 Subject: [PATCH 4/4] fix(vortex): address reader lifecycle and scan review Release selection and row-index state on close, build dense selections from compressed complements, and push empty projections into native scans. Preserve Arrow schema error details, clarify pushdown eligibility and document fallback and LTZ metadata rules. Add focused regressions and fix cast declarations reported by clang-tidy. --- docs/source/api/file_format.rst | 30 ++ .../vortex/vortex_file_batch_reader.cpp | 77 +++-- .../format/vortex/vortex_file_batch_reader.h | 1 + .../format/vortex/vortex_file_format_test.cpp | 263 +++++++++++++++++- .../vortex/vortex_predicate_converter.cpp | 35 ++- 5 files changed, 365 insertions(+), 41 deletions(-) diff --git a/docs/source/api/file_format.rst b/docs/source/api/file_format.rst index 9ae2ca7d4..20674bf76 100644 --- a/docs/source/api/file_format.rst +++ b/docs/source/api/file_format.rst @@ -81,6 +81,25 @@ safety checks. ``0`` disables only IN/NOT IN pushdown. A list exceeding the limi is not pushed down and is never truncated. Other eligible predicates may still be pushed down; an OR expression requires every branch to be convertible. +Increasing the membership limit does not bypass these fallback conditions: + +- ``FLOAT``/``DOUBLE`` comparisons and IN/NOT IN are not pushed down, including + predicates with finite constants, because native NaN ordering differs from + Paimon's comparisons. +- Timestamp constants must convert to the file's timestamp unit without precision + loss or integer overflow and remain within the supported native temporal range. + This checks the predicate constant; stored values are not truncated. +- Timestamp value predicates require an empty stored timezone or the exact + timezone label ``UTC``. Other stored labels fall back, independently of the + timezone requested by the reader. +- Decimal constants must normalize to the file's precision and scale without + rounding or overflow. Empty or NULL-containing IN/NOT IN lists also fall back. + +``IS NULL`` and ``IS NOT NULL`` remain eligible for floating-point and timestamp +fields. If the file contains a top-level user column named ``_row_idx``, however, +the reader disables all native predicate filtering, even when that column is not +projected. Projection and bitmap selection remain available. + The Vortex adapter does not deduplicate lists or compress consecutive integer values into ranges. The query engine is responsible for these optimizations. To allow up to 128 entries per IN/NOT IN predicate: @@ -111,3 +130,14 @@ For example, configure a 64 KiB gap and a 4 MiB merging threshold: read_context_builder.AddOption("vortex.read.io-coalesce.max-gap", "65536"); read_context_builder.AddOption("vortex.read.io-coalesce.max-size", "4194304"); + +Vortex Timestamp Storage +======================== + +For newly written ``TIMESTAMP_LTZ`` fields, the writer normalizes non-empty +timezone metadata to ``UTC``, including timestamps nested in supported struct, +list and fixed-size-list types. Epoch values and timestamp units are unchanged; +the represented instants are not shifted. Timestamp fields without a timezone +retain empty timezone metadata. This changes the observable stored file schema; +existing files are not rewritten. Value predicates on existing timestamp fields +with other stored timezone labels follow the fallback rules above. diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.cpp b/src/paimon/format/vortex/vortex_file_batch_reader.cpp index a8b284398..d2f7ee223 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.cpp +++ b/src/paimon/format/vortex/vortex_file_batch_reader.cpp @@ -161,29 +161,45 @@ Result VortexFileBatchReader::SelectionSt } state.selected_row_count_ = static_cast(cardinality); state.excludes_ = state.selected_row_count_ > total_rows / 2; + + // Seek to the first invalid row without enumerating a dense selection. Bitmap values + // above INT32_MAX appear as negative int32_t values and remain invalid for this reader. + const uint64_t max_bitmap_row = std::numeric_limits::max(); + const auto first_invalid_row = static_cast(std::min(total_rows, max_bitmap_row)); + const auto end = bitmap->End(); + auto invalid_row = bitmap->EqualOrLarger(first_invalid_row); + if (total_rows > max_bitmap_row && invalid_row != end && + *invalid_row == std::numeric_limits::max()) { + ++invalid_row; + } + if (invalid_row != end) { + return Status::Invalid(fmt::format("Vortex selected row {} is out of range {}", + static_cast(*invalid_row), total_rows)); + } + if (state.selected_row_count_ == total_rows) { + return state; + } + state.indices_.reserve(static_cast( state.excludes_ ? total_rows - state.selected_row_count_ : state.selected_row_count_)); - uint64_t next_row = 0; - const auto end = bitmap->End(); - for (auto iter = bitmap->Begin(); iter != end; ++iter) { - const auto row_id = static_cast(*iter); - if (row_id >= total_rows) { - return Status::Invalid( - fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows)); + if (state.excludes_) { + RoaringBitmap32 excluded; + excluded.AddRange(0, static_cast(std::min(total_rows, max_bitmap_row))); + if (total_rows > max_bitmap_row) { + excluded.Add(std::numeric_limits::max()); } - if (state.excludes_) { - // Enumerate gaps directly instead of first expanding all selected rows. - for (; next_row < row_id; ++next_row) { - state.indices_.push_back(next_row); - } - next_row = row_id + 1; - } else { - state.indices_.push_back(row_id); + excluded -= *bitmap; + const auto excluded_end = excluded.End(); + for (auto iter = excluded.Begin(); iter != excluded_end; ++iter) { + state.indices_.push_back(static_cast(*iter)); } - } - if (state.excludes_) { - for (; next_row < total_rows; ++next_row) { - state.indices_.push_back(next_row); + // Rows beyond the bitmap's nonnegative int32_t domain are all excluded. + for (uint64_t row = max_bitmap_row + 1; row < total_rows; ++row) { + state.indices_.push_back(row); + } + } else { + for (auto iter = bitmap->Begin(); iter != end; ++iter) { + state.indices_.push_back(static_cast(*iter)); } } } @@ -364,16 +380,18 @@ Status VortexFileBatchReader::CreateScan() { VxExpressionPtr root(nullptr, vx_expression_free); VxExpressionPtr projection(nullptr, vx_expression_free); std::vector name_ptrs; - if (read_schema_ != nullptr && read_schema_->num_fields() > 0) { + if (read_schema_ != nullptr) { name_ptrs.reserve(read_schema_->num_fields()); // The schema owns these names; Vortex copies them when creating the projection. for (const std::shared_ptr& field : read_schema_->fields()) { name_ptrs.push_back(field->name().c_str()); } root = VxExpressionPtr(vx_expression_root(), vx_expression_free); - projection = - VxExpressionPtr(vx_expression_select(name_ptrs.data(), name_ptrs.size(), root.get()), - vx_expression_free); + // The FFI requires non-null names even for an explicit zero-column projection. + const char* empty_name = nullptr; + const char* const* names = name_ptrs.empty() ? &empty_name : name_ptrs.data(); + projection = VxExpressionPtr(vx_expression_select(names, name_ptrs.size(), root.get()), + vx_expression_free); options.projection = projection.get(); } selection_.ApplyToScan(&options); @@ -429,19 +447,26 @@ Result VortexFileBatchReader::OpenNextPartitionStream() { } current_stream_ = stream; stream_active_ = true; + PAIMON_RETURN_NOT_OK(ImportCurrentStreamSchema()); + return true; +} + +Status VortexFileBatchReader::ImportCurrentStreamSchema() { // Import batches with the scan's actual output schema (projected columns, standard types). ::ArrowSchema ffi_schema = {}; if (current_stream_.get_schema(¤t_stream_, &ffi_schema) != 0) { + const char* message = current_stream_.get_last_error(¤t_stream_); + std::string message_copy = message == nullptr ? "unknown" : message; Status callback_status = input_context_->GetCallbackStatus(); ReleaseStream(); if (!callback_status.ok()) { return callback_status.WithMessage("read Vortex stream schema", ": ", callback_status.message()); } - return Status::IOError("Vortex Arrow stream get_schema failed"); + return Status::IOError("Vortex Arrow stream get_schema failed: ", message_copy); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(current_stream_schema_, arrow::ImportSchema(&ffi_schema)); - return true; + return Status::OK(); } Result> VortexFileBatchReader::ReadNextArray() { @@ -600,6 +625,8 @@ void VortexFileBatchReader::CloseInternal() { ReleaseStream(); current_batch_.reset(); scan_.reset(); + selection_ = SelectionState{}; + row_position_.Reset(); data_source_.reset(); session_.reset(); input_context_.reset(); diff --git a/src/paimon/format/vortex/vortex_file_batch_reader.h b/src/paimon/format/vortex/vortex_file_batch_reader.h index 8a8b122cb..8d4dcbbb2 100644 --- a/src/paimon/format/vortex/vortex_file_batch_reader.h +++ b/src/paimon/format/vortex/vortex_file_batch_reader.h @@ -139,6 +139,7 @@ class VortexFileBatchReader : public FileBatchReader { /// Open the next partition and fill `current_stream_` with bounded Arrow batches. /// Returns false when there are no more partitions. Result OpenNextPartitionStream(); + Status ImportCurrentStreamSchema(); /// Reshape a scanned batch to `read_schema_`, selecting columns by name and recursing into /// nested fields. Returns `array` unchanged when SetReadSchema has not been called. Result> AlignToReadSchema( diff --git a/src/paimon/format/vortex/vortex_file_format_test.cpp b/src/paimon/format/vortex/vortex_file_format_test.cpp index b3f6345ee..3466aa74e 100644 --- a/src/paimon/format/vortex/vortex_file_format_test.cpp +++ b/src/paimon/format/vortex/vortex_file_format_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -59,6 +60,7 @@ #include "paimon/format/format_writer.h" #include "paimon/format/reader_builder.h" #include "paimon/format/vortex/vortex_defs.h" +#include "paimon/format/vortex/vortex_file_batch_reader.h" #include "paimon/format/vortex/vortex_io_callbacks.h" #include "paimon/format/writer_builder.h" #include "paimon/fs/file_system.h" @@ -1507,6 +1509,77 @@ TEST_F(VortexFileFormatTest, SelectionPushdownReturnsSelectedRowsAndPhysicalRowI ASSERT_EQ(got_row_ids, (std::vector{1, 3, 5, 7})); } +TEST(VortexSelectionStateTest, LargeDenseSelectionsPreserveBitmapAndPhysicalRanks) { + constexpr int32_t kRows = 100000000; + for (const int32_t deleted_row : {-1, 0, kRows / 2, kRows - 1}) { + SCOPED_TRACE(deleted_row); + RoaringBitmap32 bitmap; + bitmap.AddRange(0, kRows); + if (deleted_row >= 0) { + bitmap.RemoveRange(deleted_row, deleted_row + 1); + } + const RoaringBitmap32 original = bitmap; + ASSERT_OK_AND_ASSIGN(VortexFileBatchReader::SelectionState state, + VortexFileBatchReader::SelectionState::Create(bitmap, kRows)); + ASSERT_TRUE(bitmap == original); + const uint64_t expected_count = kRows - (deleted_row >= 0 ? 1 : 0); + ASSERT_EQ(state.SelectedRowCount(), expected_count); + ASSERT_EQ(state.IndexCount(), deleted_row >= 0 ? 1 : 0); + for (const uint64_t rank : {uint64_t{0}, expected_count / 2, expected_count - 1}) { + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, state.GetPhysicalRowId(rank)); + ASSERT_EQ(physical_row, + rank + (deleted_row >= 0 && rank >= static_cast(deleted_row))); + } + } +} + +TEST(VortexSelectionStateTest, DenseSelectionSupportsRowsBeyondSignedBitmapRange) { + constexpr int32_t kMaxRow = std::numeric_limits::max(); + for (const bool selects_max_row : {false, true}) { + SCOPED_TRACE(selects_max_row); + RoaringBitmap32 bitmap; + bitmap.AddRange(0, kMaxRow); + if (selects_max_row) { + bitmap.RemoveRange(kMaxRow - 1, kMaxRow); + bitmap.Add(kMaxRow); + ASSERT_TRUE(VortexFileBatchReader::SelectionState::Create(bitmap, kMaxRow) + .status() + .IsInvalid()); + } + const RoaringBitmap32 original = bitmap; + const uint64_t total_rows = static_cast(kMaxRow) + 2; + ASSERT_OK_AND_ASSIGN(VortexFileBatchReader::SelectionState state, + VortexFileBatchReader::SelectionState::Create(bitmap, total_rows)); + ASSERT_TRUE(bitmap == original); + ASSERT_EQ(state.SelectedRowCount(), static_cast(kMaxRow)); + const auto excluded = static_cast(selects_max_row ? kMaxRow - 1 : kMaxRow); + ASSERT_EQ(state.indices_, (std::vector{excluded, total_rows - 1})); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, state.GetPhysicalRowId(kMaxRow - 1)); + ASSERT_EQ(physical_row, static_cast(selects_max_row ? kMaxRow : kMaxRow - 1)); + ASSERT_FALSE(state.GetPhysicalRowId(kMaxRow).ok()); + } +} + +TEST(VortexSelectionStateTest, SelectionBoundsAreValidatedBeforeFastPaths) { + for (const auto& rows : std::vector>{{0, 1, 2, 4}, {-1, 0, 1, 2}}) { + const RoaringBitmap32 bitmap = RoaringBitmap32::From(rows); + ASSERT_TRUE(VortexFileBatchReader::SelectionState::Create(bitmap, 4).status().IsInvalid()); + } + constexpr int32_t kMaxRow = std::numeric_limits::max(); + const RoaringBitmap32 bitmap = RoaringBitmap32::From({0, kMaxRow}); + ASSERT_TRUE( + VortexFileBatchReader::SelectionState::Create(bitmap, kMaxRow).status().IsInvalid()); + ASSERT_OK_AND_ASSIGN( + VortexFileBatchReader::SelectionState state, + VortexFileBatchReader::SelectionState::Create(bitmap, static_cast(kMaxRow) + 1)); + ASSERT_OK_AND_ASSIGN(uint64_t physical_row, state.GetPhysicalRowId(1)); + ASSERT_EQ(physical_row, static_cast(kMaxRow)); + ASSERT_TRUE(VortexFileBatchReader::SelectionState::Create(RoaringBitmap32::From({0, -1}), + static_cast(kMaxRow) + 2) + .status() + .IsInvalid()); +} + TEST_F(VortexFileFormatTest, DenseSelectionKeepsOnlyExcludedIndicesAndPreservesNullValues) { constexpr int32_t kRows = 4096; constexpr int32_t kDeletedRow = 2048; @@ -2500,7 +2573,7 @@ TEST_F(VortexFileFormatTest, BoundedStringBatchesPreserveProjectionRowIdsAndBuff const auto& actual_binaries = checked_cast(*values.field(0)); const auto& actual_texts = checked_cast(*values.field(1)); for (int64_t i = 0; i < imported->length(); ++i) { - const int64_t physical_row = static_cast(expected_rows[rows_seen++]); + const auto physical_row = static_cast(expected_rows[rows_seen++]); ASSERT_EQ(actual_texts.IsNull(i), expected_texts.IsNull(physical_row)); ASSERT_EQ(actual_binaries.IsNull(i), expected_binaries.IsNull(physical_row)); if (!actual_texts.IsNull(i)) { @@ -2562,6 +2635,13 @@ TEST_F(VortexFileFormatTest, EmptyProjectionPreservesRowCountsAndPhysicalIds) { } ASSERT_TRUE(array->ValidateFull().ok()); ASSERT_TRUE(array->type()->Equals(arrow::struct_(projection->fields()))); + const auto* native_reader = checked_cast(reader.get()); + // Inspect the native export before C++ alignment can discard payload columns. + ASSERT_EQ(native_reader->current_stream_schema_->num_fields(), filtered ? 1 : 0); + ASSERT_EQ(native_reader->current_batch_->type()->num_fields(), filtered ? 1 : 0); + if (filtered) { + ASSERT_EQ(native_reader->current_stream_schema_->field(0)->name(), "_row_idx"); + } for (int64_t i = 0; i < array->length(); ++i) { ASSERT_OK_AND_ASSIGN(uint64_t row, reader->GetPreviousBatchFileRowId(static_cast(i))); @@ -2577,6 +2657,79 @@ TEST_F(VortexFileFormatTest, EmptyProjectionPreservesRowCountsAndPhysicalIds) { } } +TEST_F(VortexFileFormatTest, EmptyProjectionAvoidsPayloadReads) { + constexpr int32_t kRows = 64; + arrow::Int32Builder id_builder; + arrow::StringBuilder payload_builder; + uint32_t seed = 13; + std::string payload(16384, '\0'); + for (int32_t row = 0; row < kRows; ++row) { + ASSERT_TRUE(id_builder.Append(row).ok()); + for (char& byte : payload) { + seed = seed * 1664525U + 1013904223U; + byte = static_cast('!' + ((seed >> 24) % 90)); + } + ASSERT_TRUE(payload_builder.Append(payload).ok()); + } + std::shared_ptr ids; + std::shared_ptr payloads; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(payload_builder.Finish(&payloads).ok()); + const auto schema = arrow::schema( + {arrow::field("id", arrow::int32(), false), arrow::field("payload", arrow::utf8(), false)}); + const auto data = arrow::StructArray::Make({ids, payloads}, schema->fields()).ValueOrDie(); + const std::string path = PathUtil::JoinPath(directory_->Str(), "empty-projection-io.vortex"); + ASSERT_OK(WriteFile(path, schema, data, kRows)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr format, + FileFormatFactory::Get("vortex", {{"file.format", "vortex"}, + {kVortexReadIoCoalesceMaxGap, "0"}, + {kVortexReadIoCoalesceMaxSize, "1"}})); + for (const bool filtered : {false, true}) { + SCOPED_TRACE(filtered); + std::vector bytes_read; + for (const bool empty_projection : {false, true}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(path)); + ASSERT_OK_AND_ASSIGN(int64_t file_bytes, input->Length()); + // Keep payload outside the native footer read-ahead cache for the I/O comparison. + ASSERT_GT(file_bytes, int64_t{256 * 1024}); + auto tracked_input = std::make_shared(input); + ASSERT_OK_AND_ASSIGN(std::unique_ptr builder, + format->CreateReaderBuilder(/*batch_size=*/16)); + builder->WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + builder->Build(tracked_input)); + const std::shared_ptr predicate = + filtered ? PredicateBuilder::GreaterThan(0, "id", FieldType::INT, + Literal(static_cast(31))) + : nullptr; + ASSERT_OK(ConfigureReader(reader.get(), empty_projection ? arrow::schema({}) : schema, + predicate)); + tracked_input->ClearReadAtRanges(); + uint64_t rows = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, + ReadNextBatch(reader.get())); + if (batch == nullptr) { + break; + } + ASSERT_TRUE(batch->ValidateFull().ok()); + ASSERT_EQ(batch->type()->num_fields(), empty_projection ? 0 : 2); + rows += static_cast(batch->length()); + } + ASSERT_EQ(rows, filtered ? kRows / 2 : kRows); + reader->Close(); + uint64_t bytes = 0; + for (const auto& range : tracked_input->GetReadAtRanges()) { + bytes += static_cast(range.second); + } + bytes_read.push_back(bytes); + } + ASSERT_GT(bytes_read[0], uint64_t{0}); + // Metadata and predicate columns may still be read; the large payload must be omitted. + ASSERT_LT(bytes_read[1], bytes_read[0] / 4); + } +} + TEST_F(VortexFileFormatTest, RepeatedProjectionFieldsAreRejected) { ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); @@ -2753,6 +2906,114 @@ TEST_F(VortexFileFormatTest, InvalidArrowSchemaResetDoesNotChangeTheActiveReader ASSERT_EQ(remaining, (std::vector{4})); } +TEST_F(VortexFileFormatTest, CloseReleasesSelectionAndRowPositionWhileReaderRemainsAlive) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto predicate = + PredicateBuilder::GreaterThan(0, "id", FieldType::INT, Literal(static_cast(0))); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), predicate, + RoaringBitmap32::From({1, 3, 5}))); + auto* native_reader = checked_cast(reader.get()); + ASSERT_GT(native_reader->selection_.indices_.capacity(), size_t{0}); + std::weak_ptr row_ids; + { + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch, ReadNextBatch(reader.get())); + ASSERT_NE(batch, nullptr); + ASSERT_NE(native_reader->row_position_.previous_batch_row_idx_, nullptr); + row_ids = native_reader->row_position_.previous_batch_row_idx_; + } + ASSERT_FALSE(row_ids.expired()); + reader->Close(); + ASSERT_NE(reader, nullptr); + // The index vector uses the standard allocator, so Arrow pool counters cannot check it. + ASSERT_EQ(native_reader->selection_.indices_.capacity(), size_t{0}); + ASSERT_TRUE(row_ids.expired()); + reader->Close(); + ASSERT_FALSE(reader->NextBatch().ok()); +} + +TEST_F(VortexFileFormatTest, ReturnedFilteredBatchOutlivesClosedReader) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + const auto predicate = + PredicateBuilder::GreaterThan(0, "id", FieldType::INT, Literal(static_cast(2))); + ASSERT_OK(ConfigureReader(reader.get(), ReaderTestSchema(), predicate, + RoaringBitmap32::From({0, 3, 4, 5}))); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch retained, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(retained)); + ASSERT_OK_AND_ASSIGN(uint64_t row, reader->GetPreviousBatchFileRowId(0)); + ASSERT_EQ(row, uint64_t{3}); + reader->Close(); + reader.reset(); + ASSERT_OK(paimon::test::ReadResultCollector::CheckBatchOffset(retained)); + auto batch = arrow::ImportArray(retained.first.get(), retained.second.get()).ValueOrDie(); + ASSERT_TRUE(batch->ValidateFull().ok()); + auto expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(ReaderTestSchema()->fields()), + R"([[3,"three"],[4,null],[5,"five"]])") + .ValueOrDie(); + ASSERT_TRUE(batch->Equals(expected->Slice(0, batch->length()))); +} + +TEST_F(VortexFileFormatTest, SchemaStreamErrorsPreserveDetailsAndReleaseStream) { + ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); + for (const bool has_native_message : {false, true}) { + for (const bool has_callback_error : {false, true}) { + SCOPED_TRACE(has_native_message); + SCOPED_TRACE(has_callback_error); + struct SchemaFailure { + std::string message = "native schema export sentinel"; + bool has_message = false; + bool error_read_before_release = false; + int32_t releases = 0; + } failure; + failure.has_message = has_native_message; + const std::string native_message = failure.message; + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); + auto* native_reader = checked_cast(reader.get()); + ::ArrowArrayStream stream{}; + stream.private_data = &failure; + stream.get_schema = [](::ArrowArrayStream*, ::ArrowSchema*) -> int32_t { return EIO; }; + stream.get_last_error = [](::ArrowArrayStream* current) -> const char* { + auto* state = static_cast(current->private_data); + state->error_read_before_release = state->releases == 0; + return state->has_message ? state->message.c_str() : nullptr; + }; + stream.release = [](::ArrowArrayStream* current) { + auto* state = static_cast(current->private_data); + ++state->releases; + // Invalidate the stream-owned message without introducing a dangling pointer. + std::fill(state->message.begin(), state->message.end(), '#'); + current->release = nullptr; + }; + native_reader->current_stream_ = stream; + native_reader->stream_active_ = true; + if (has_callback_error) { + native_reader->input_context_->SetCallbackStatus( + Status::IOError("filesystem schema callback sentinel")); + } + const Status status = native_reader->ImportCurrentStreamSchema(); + ASSERT_TRUE(status.IsIOError()) << status.ToString(); + ASSERT_EQ(failure.releases, 1); + ASSERT_FALSE(native_reader->stream_active_); + ASSERT_EQ(native_reader->current_stream_schema_, nullptr); + if (has_callback_error) { + ASSERT_NE(status.message().find("filesystem schema callback sentinel"), + std::string::npos); + ASSERT_EQ(status.message().find(native_message), std::string::npos); + } else if (has_native_message) { + ASSERT_TRUE(failure.error_read_before_release); + ASSERT_NE(status.message().find(native_message), std::string::npos); + } else { + ASSERT_NE(status.message().find("get_schema failed"), std::string::npos); + ASSERT_NE(status.message().find("unknown"), std::string::npos); + } + reader->Close(); + ASSERT_EQ(failure.releases, 1); + } + } +} + TEST_F(VortexFileFormatTest, ReturnedBatchOutlivesClosedReaderAndClosedReaderRejectsReset) { ASSERT_OK_AND_ASSIGN(std::string path, WriteReaderTestFile()); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, CreateReader(path)); diff --git a/src/paimon/format/vortex/vortex_predicate_converter.cpp b/src/paimon/format/vortex/vortex_predicate_converter.cpp index afb54d255..286771063 100644 --- a/src/paimon/format/vortex/vortex_predicate_converter.cpp +++ b/src/paimon/format/vortex/vortex_predicate_converter.cpp @@ -257,9 +257,10 @@ Result MakeLiteralExpr(const Literal& literal, } // Type mismatches violate the predicate contract; valid but unsafe conversions fall back. -Result LiteralMatchesColumn(const Literal& literal, - const std::shared_ptr& leaf, - const std::shared_ptr& arrow_type) { +// True passes basic type eligibility only; scalar construction still checks value safety. +Result CheckLiteralTypePushdownEligibility( + const Literal& literal, const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type) { if (literal.IsNull() || arrow_type == nullptr) { return false; } @@ -303,17 +304,19 @@ Result LiteralMatchesColumn(const Literal& literal, } // Check every basic type even when a value or the list budget requires fallback. -Result ValidateMembershipLiteralTypes(const std::shared_ptr& leaf, - const std::shared_ptr& arrow_type) { - bool valid = true; +Result CheckMembershipTypePushdownEligibility( + const std::shared_ptr& leaf, + const std::shared_ptr& arrow_type) { + bool all_literal_types_eligible = true; for (const Literal& literal : leaf->Literals()) { - PAIMON_ASSIGN_OR_RAISE(bool matches, LiteralMatchesColumn(literal, leaf, arrow_type)); - if (!matches) { - valid = false; + PAIMON_ASSIGN_OR_RAISE(bool type_eligible, + CheckLiteralTypePushdownEligibility(literal, leaf, arrow_type)); + if (!type_eligible) { + all_literal_types_eligible = false; continue; } } - return valid; + return all_literal_types_eligible; } Result ConvertInList(const std::vector& literals, @@ -352,10 +355,12 @@ Result ConvertMembership(const std::shared_ptr& const vx_expression* column, bool negate, const VortexPredicateOptions& options) { const std::vector& literals = leaf->Literals(); - PAIMON_ASSIGN_OR_RAISE(bool valid, ValidateMembershipLiteralTypes(leaf, arrow_type)); + PAIMON_ASSIGN_OR_RAISE(bool all_literal_types_eligible, + CheckMembershipTypePushdownEligibility(leaf, arrow_type)); // The query engine owns deduplication and integer-range rewrites. Apply value conversions // and create native scalars only within the original-list limit, after type validation. - if (!valid || literals.empty() || literals.size() > options.max_in_values) { + if (!all_literal_types_eligible || literals.empty() || + literals.size() > options.max_in_values) { return NullExpr(); } if (leaf->GetFieldType() == FieldType::TIMESTAMP) { @@ -446,9 +451,9 @@ Result ConvertLeaf(const std::shared_ptr& leaf, if (!is_comparison) { return NullExpr(); } - PAIMON_ASSIGN_OR_RAISE(bool matches, - LiteralMatchesColumn(leaf->Literals()[0], leaf, arrow_type)); - if (!matches) { + PAIMON_ASSIGN_OR_RAISE(bool type_eligible, CheckLiteralTypePushdownEligibility( + leaf->Literals()[0], leaf, arrow_type)); + if (!type_eligible) { return NullExpr(); } PAIMON_ASSIGN_OR_RAISE(VxExpressionPtr literal_expr,