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/2] 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/2] 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);