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..801e18e3a --- /dev/null +++ b/cmake_modules/vortex-fmm-cache-shards.diff @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +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..8c24d3e58 --- /dev/null +++ b/cmake_modules/vortex-point-read.diff @@ -0,0 +1,374 @@ +# 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 +@@ -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..fa7aa6324 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,68 @@ VortexFileBatchReader::~VortexFileBatchReader() { } Result VortexFileBatchReader::OpenNextPartitionStream() { + if (has_selection_ && selected_row_count_ == 0) { + 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(); + } + } + // These FFI modes take individual indices. A dense bitmap sends only excluded rows; + // an all-selected bitmap needs no indices and uses the default INCLUDE_ALL mode. + if (has_selection_ && !selection_indices_.empty()) { + options.selection.idx = selection_indices_.data(); + options.selection.idx_len = selection_indices_.size(); + options.selection.include = + selection_excludes_ ? VX_SELECTION_EXCLUDE_RANGE : VX_SELECTION_INCLUDE_RANGE; + } + // Filtering drops rows; include_row_idx preserves their original physical positions. + // The scan clones the filter, so it can be freed after scan creation. + VxExpressionPtr filter(nullptr, vx_expression_free); + // The hidden index must not collide with a file field or a requested missing field. + const bool has_row_idx_column = + file_schema_->GetFieldIndex(kRowIdxColumn) >= 0 || + (read_schema_ != nullptr && !read_schema_->GetAllFieldIndices(kRowIdxColumn).empty()); + if (!has_row_idx_column) { + filter = VortexPredicateConverter::Convert(predicate_filter_, file_schema_, + predicate_options_); + } + if (filter != nullptr) { + options.filter = filter.get(); + options.include_row_idx = true; + has_filter_ = true; + } VxScanPtr scan( vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), vx_scan_free); @@ -347,6 +290,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 +334,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 +349,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 +374,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 +409,59 @@ 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. + std::vector selection_indices; + bool selection_excludes = false; + uint64_t selected_row_count = total_rows_; + if (selection_bitmap) { + const int32_t cardinality = selection_bitmap->Cardinality(); + if (cardinality < 0 || static_cast(cardinality) > total_rows_) { + return Status::Invalid("Vortex selection cardinality exceeds the file row count"); + } + selected_row_count = static_cast(cardinality); + selection_excludes = selected_row_count > total_rows_ / 2; + selection_indices.reserve(static_cast( + selection_excludes ? total_rows_ - selected_row_count : selected_row_count)); + uint64_t next_row = 0; + const auto end = selection_bitmap->End(); + for (auto iter = selection_bitmap->Begin(); iter != end; ++iter) { + const auto row_id = static_cast(*iter); + if (row_id >= total_rows_) { + return Status::Invalid( + fmt::format("Vortex selected row {} is out of range {}", row_id, total_rows_)); + } + if (selection_excludes) { + // Enumerate gaps directly instead of first expanding all selected rows. + for (; next_row < row_id; ++next_row) { + selection_indices.push_back(next_row); + } + next_row = row_id + 1; + } else { + selection_indices.push_back(row_id); + } + } + if (selection_excludes) { + for (; next_row < total_rows_; ++next_row) { + selection_indices.push_back(next_row); + } + } + } + // 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_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; @@ -466,6 +469,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 +484,41 @@ 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; + // Ordered scans emit selected rows in physical order, regardless of which side is stored. + if (has_selection_) { + if (emitted_offset >= selected_row_count_) { + return Status::Invalid(fmt::format("emitted row offset {} exceeds selection size {}", + emitted_offset, selected_row_count_)); + } + if (!selection_excludes_) { + return selection_indices_[emitted_offset]; + } + // For excluded row i, excluded[i] - i is the number of selected rows before it. + // Count exclusions at or before this selected rank without expanding a row mapping. + size_t begin = 0; + size_t end = selection_indices_.size(); + while (begin < end) { + const size_t middle = begin + (end - begin) / 2; + if (selection_indices_[middle] - middle <= emitted_offset) { + begin = middle + 1; + } else { + end = middle; + } + } + return emitted_offset + begin; + } + return emitted_offset; } Result VortexFileBatchReader::GetNumberOfRows() const { @@ -498,6 +537,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..5175b10ad 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,21 @@ class VortexFileBatchReader : public FileBatchReader { uint64_t previous_first_row_ = std::numeric_limits::max(); uint64_t previous_batch_row_count_ = 0; + // Keep the smaller side of the bitmap. Dense selections store excluded physical row ids, + // avoiding an O(file rows) index list both here and in the native scan's copy. + bool has_selection_ = false; + bool selection_excludes_ = false; + uint64_t selected_row_count_ = 0; + std::vector selection_indices_; + + // A pushed filter returns physical row ids in a hidden _row_idx column. Retain the last + // batch's column because emitted offsets no longer identify physical rows after filtering. + bool has_filter_ = false; + std::shared_ptr previous_batch_row_idx_; + 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..ae8791de8 100644 --- a/src/paimon/format/vortex/vortex_file_format_test.cpp +++ b/src/paimon/format/vortex/vortex_file_format_test.cpp @@ -19,13 +19,17 @@ // Round-trip unit tests for the Vortex file format: write, then read back via vortex-ffi. +#include #include #include #include #include #include +#include +#include #include #include +#include #include #include #include @@ -35,13 +39,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 +62,15 @@ #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" #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 +170,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 +1109,831 @@ 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, 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; + 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..ce63c7e54 --- /dev/null +++ b/src/paimon/format/vortex/vortex_predicate_converter_test.cpp @@ -0,0 +1,1372 @@ +/* + * 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); + LocalFileSystem file_system; + { + ASSERT_OK_AND_ASSIGN(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); + 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