SteNicholas commented on code in PR #376: URL: https://github.com/apache/paimon-cpp/pull/376#discussion_r4255423476
########## crates/vortex_callback_io/callback_io.rs: ########## @@ -0,0 +1,781 @@ +// 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. + +//! Callback-based I/O for the Vortex C API. +//! +//! The stock C API can only read a whole file already materialized in memory +//! (`vx_data_source_new_buffer`) or a path that Vortex itself resolves +//! (`vx_data_source_new`), and it can only write to a local filesystem path +//! (`vx_array_sink_open_file`, which uses `async_fs::File::create`). paimon owns its own pluggable +//! filesystem, so none of these fit: reading forces the entire file into memory, and writing forces +//! a local temporary file that then has to be copied back. +//! +//! This module adds entry points that take host callbacks instead, so Vortex issues positional +//! reads and sequential writes straight against paimon's `InputStream` / `OutputStream`. Vortex's +//! core I/O traits are the plug points; the shape of the bridge mirrors `vortex-jni`'s +//! `JavaReadable` and `JavaWrite`. +//! +//! The write side needs its own handle type rather than reusing `vx_array_sink`, whose fields are +//! private to `crate::sink`. +//! +//! The sink entry points below (`vx_callback_sink_push` / `_close` / `_abort`) are adapted from +//! vortex-ffi's `src/sink.rs` (SPDX-License-Identifier: Apache-2.0, Copyright the Vortex +//! contributors), reworked to drive the write pipeline through host callbacks instead of a +//! Vortex-resolved filesystem path. +//! +//! This file is maintained in the paimon-cpp tree (`crates/vortex_callback_io/`) and copied into +//! `vortex-ffi/src/` at build time; see `cmake_modules/vortex.diff`. +//! +//! The C declarations that mirror the types and entry points below live in paimon-cpp's +//! `src/paimon/format/vortex/vortex_ffi.h` (cbindgen does not see this module, so they are written +//! by hand). The two sides are only checked by human review, and a mismatch in field order or a +//! function signature is silent memory corruption, not a compile error. Any change to a +//! `#[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::c_void; +use std::io; +use std::sync::Arc; + +use futures::FutureExt; +use futures::SinkExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::channel::mpsc::Sender; +use futures::future::BoxFuture; +use vortex::array::ArrayRef; +use vortex::array::buffer::BufferHandle; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::buffer::Alignment; +use vortex::buffer::ByteBufferMut; +use vortex::dtype::DType; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::file::WriteSummary; +use vortex::io::CoalesceConfig; +use vortex::io::IoBuf; +use vortex::io::VortexReadAt; +use vortex::io::VortexWrite; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::Handle; +use vortex::io::runtime::Task; +use vortex::io::session::RuntimeSessionExt; + +use crate::RUNTIME; +use crate::array::vx_array; +use crate::data_source::vx_data_source; +use crate::dtype::vx_dtype; +use crate::error::try_or; +use crate::error::try_or_default; +use crate::error::vx_error; +use crate::session::vx_session; + +/// 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. +/// The host is expected to keep its own error detail on the context, since only a status code +/// crosses the boundary. +pub type vx_read_at_fn = + unsafe extern "C" fn(ctx: *mut c_void, offset: u64, dst: *mut u8, length: usize) -> i32; + +/// Release the host context. Called exactly once, when the owning handle is freed. +pub type vx_release_fn = unsafe extern "C" fn(ctx: *mut c_void); + +/// 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. +/// Unlike reads, writes are sequential and never concurrent for one sink. +pub type vx_write_fn = unsafe extern "C" fn(ctx: *mut c_void, src: *const u8, length: usize) -> i32; + +/// Flush whatever the host has buffered. Returns 0 on success and non-zero on failure. +pub type vx_flush_fn = unsafe extern "C" fn(ctx: *mut c_void) -> i32; + +/// Host callbacks backing a positional reader. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_input_callbacks { + /// Opaque host context, passed back to every callback. + pub ctx: *mut c_void, + /// Positional read. Required. + /// + /// Must be thread-safe: Vortex issues concurrent positional reads, so this is called from + /// several blocking threads at once for the same context. + pub read_at_fn: Option<vx_read_at_fn>, + /// Context destructor. Optional; when set, it is the last callback invoked. + pub release_fn: Option<vx_release_fn>, +} + +/// Host callbacks backing a sequential writer. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_output_callbacks { + /// Opaque host context, passed back to every callback. + pub ctx: *mut c_void, + /// Sequential write. Required. + pub write_fn: Option<vx_write_fn>, + /// Flush. Optional; when null, flush requests are ignored. + pub flush_fn: Option<vx_flush_fn>, + /// Context destructor. Optional; when set, it is the last callback invoked. + pub release_fn: Option<vx_release_fn>, +} + +/// The host context, owned by the reader. +/// +/// Held behind an `Arc` so in-flight reads keep the context alive, and so `release_fn` runs exactly +/// once when the last reference goes away. It carries only the context, not the read callback, so +/// that ownership can be taken over before anything else is validated. +struct CallbackTarget { + ctx: *mut c_void, + release_fn: Option<vx_release_fn>, +} + +// SAFETY: the host contract for `vx_input_callbacks` requires `read_at_fn` to be thread-safe for +// concurrent calls on the same `ctx`, which is what makes it sound to move the context across +// threads and share it between concurrent reads. +unsafe impl Send for CallbackTarget {} +// SAFETY: see the `Send` impl above. +unsafe impl Sync for CallbackTarget {} + +impl Drop for CallbackTarget { + fn drop(&mut self) { + if let Some(release) = self.release_fn { + // SAFETY: `ctx` was supplied by the host together with `release_fn`, and an `Arc` + // guarantees this runs once, after every in-flight read has finished. + unsafe { release(self.ctx) }; + } + } +} + +/// A [`VortexReadAt`] that forwards positional reads to host callbacks. +/// +/// Reads run on the runtime's blocking pool, since the host callback is synchronous and must not +/// occupy an async executor thread. The file size is supplied at construction time, so `size()` +/// never crosses the FFI boundary. +struct CallbackReadAt { + target: Arc<CallbackTarget>, + read_at_fn: vx_read_at_fn, + len: u64, + handle: Handle, + concurrency: usize, +} + +/// Number of concurrent host read callbacks Vortex may have in flight for one file. +/// +/// The host filesystem may be remote, where concurrency hides latency; it is also the bound on +/// blocking threads occupied by one reader. +const DEFAULT_CONCURRENCY: usize = 16; + +impl VortexReadAt for CallbackReadAt { + fn coalesce_config(&self) -> Option<CoalesceConfig> { + // 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()) + } + + fn concurrency(&self) -> usize { + self.concurrency + } + + fn size(&self) -> BoxFuture<'static, VortexResult<u64>> { + let len = self.len; + async move { Ok(len) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult<BufferHandle>> { + let target = Arc::clone(&self.target); + let read_at_fn = self.read_at_fn; + let len = self.len; + let handle = self.handle.clone(); + + async move { + handle + .spawn_blocking(move || { + let end = offset + .checked_add(length as u64) + .ok_or_else(|| vortex_err!("read {offset}+{length} overflows u64"))?; + if end > len { + vortex_bail!("read {offset}..{end} out of bounds for file of length {len}"); + } + + let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment); + if length > 0 { + // SAFETY: the spare capacity covers `length` bytes of live allocation, and + // the host contract requires the callback to fill exactly that many bytes + // before returning success. + let status = unsafe { + read_at_fn( + target.ctx, + offset, + buffer.spare_capacity_mut().as_mut_ptr().cast::<u8>(), + length, + ) + }; + if status != 0 { + vortex_bail!( + "host read callback failed with status {status} for {offset}..{end}" + ); + } + } + // SAFETY: the callback reported success, which per its contract means all + // `length` bytes were written. + unsafe { buffer.set_len(length) }; + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .await + } + .boxed() + } +} + +unsafe fn data_source_new_callback( + session: *const vx_session, + callbacks: vx_input_callbacks, + size: u64, +) -> 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. + let target = Arc::new(CallbackTarget { + ctx: callbacks.ctx, + release_fn: callbacks.release_fn, + }); + + vortex_ensure!(!session.is_null()); + let read_at_fn = callbacks + .read_at_fn + .ok_or_else(|| vortex_err!("vx_input_callbacks.read_at_fn is required"))?; + + let session = vx_session::as_ref(session).clone(); + let reader: Arc<dyn VortexReadAt> = Arc::new(CallbackReadAt { + target, + read_at_fn, + len: size, + handle: session.handle(), + concurrency: DEFAULT_CONCURRENCY, + }); + + let file = RUNTIME.block_on(async { session.open_options().open(reader).await })?; + Ok(vx_data_source::new(file.data_source()?)) +} + +/// Create a data source that reads through host callbacks. +/// +/// `size` is the total length of the file in bytes; the host knows it up front, so it is passed +/// here instead of being fetched through another callback. +/// +/// The callbacks (and the context they carry) are owned by the returned data source and released +/// when it is freed, including when this call fails. +/// +/// On error, returns NULL and sets "err". +/// +/// # Safety +/// +/// `session` must be a valid `vx_session`. `callbacks.read_at_fn` must be non-null, thread-safe, +/// and valid for as long as the data source lives. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_new_callback( + session: *const vx_session, + callbacks: vx_input_callbacks, + size: u64, + err: *mut *mut vx_error, +) -> *const vx_data_source { + try_or(err, std::ptr::null(), || unsafe { + data_source_new_callback(session, callbacks, size) + }) +} + +/// A [`VortexWrite`] that forwards sequential writes to host callbacks. +/// +/// Writes run inline on the runtime thread driving the write task, the same way `vortex-jni` +/// forwards to Java sinks: the host callback is synchronous and the layout writer hands over one +/// buffer at a time, so there is nothing to overlap. +struct CallbackWrite { + target: Arc<CallbackTarget>, + write_fn: vx_write_fn, + flush_fn: Option<vx_flush_fn>, +} + +impl CallbackWrite { + fn write_slice(&self, bytes: &[u8]) -> io::Result<()> { + if bytes.is_empty() { + return Ok(()); + } + // SAFETY: `bytes` points at a live slice for the duration of the call, and the host + // contract forbids retaining it afterwards. + let status = unsafe { (self.write_fn)(self.target.ctx, bytes.as_ptr(), bytes.len()) }; + if status != 0 { + return Err(io::Error::other(format!( + "host write callback failed with status {status}" + ))); + } + Ok(()) + } + + fn flush_host(&self) -> io::Result<()> { + let Some(flush) = self.flush_fn else { + return Ok(()); + }; + // SAFETY: the context is kept alive by `target`. + let status = unsafe { flush(self.target.ctx) }; + if status != 0 { + return Err(io::Error::other(format!( + "host flush callback failed with status {status}" + ))); + } + Ok(()) + } +} + +impl VortexWrite for CallbackWrite { + async fn write_all<B: IoBuf>(&mut self, buffer: B) -> io::Result<B> { + self.write_slice(buffer.as_slice())?; + Ok(buffer) + } + + async fn flush(&mut self) -> io::Result<()> { + self.flush_host() + } + + async fn shutdown(&mut self) -> io::Result<()> { + // The host owns its stream and closes it itself, so this only flushes. + self.flush_host() + } +} + +/// A sink that writes a Vortex file through host callbacks. +/// +/// Mirrors `vx_array_sink`: pushed arrays go through a channel feeding a write task on the session's +/// runtime, and errors surface when the sink is closed. This is a separate type because +/// `vx_array_sink`'s fields are private to its own module. +pub struct vx_callback_sink { + sink: Sender<VortexResult<ArrayRef>>, + writer: Task<VortexResult<WriteSummary>>, + dtype: DType, +} + +unsafe fn callback_sink_open( + session: *const vx_session, + callbacks: vx_output_callbacks, + dtype: *const vx_dtype, +) -> VortexResult<*mut vx_callback_sink> { + // Take ownership of the host context before anything else, so an early error below still + // releases it exactly once when this `Arc` is dropped. + let target = Arc::new(CallbackTarget { + ctx: callbacks.ctx, + release_fn: callbacks.release_fn, + }); + + vortex_ensure!(!session.is_null()); + vortex_ensure!(!dtype.is_null()); + let write_fn = callbacks + .write_fn + .ok_or_else(|| vortex_err!("vx_output_callbacks.write_fn is required"))?; + + let session = vx_session::as_ref(session).clone(); + let dtype = vx_dtype::as_ref(dtype).clone(); + let write = CallbackWrite { + target, + write_fn, + flush_fn: callbacks.flush_fn, + }; + + // Small on purpose so the producing thread feels back-pressure, as in vortex-jni's writer. + let (sink, rx) = mpsc::channel(4); + let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); + + let strategy = WriteStrategyBuilder::default().build(); + let writer_session = session.clone(); + let writer = session.handle().spawn(async move { + writer_session + .write_options() + .with_strategy(strategy) + .write(write, array_stream) + .await + }); + + Ok(Box::into_raw(Box::new(vx_callback_sink { + sink, + writer, + dtype, + }))) +} + +/// Open a sink that writes a Vortex file through `callbacks`. +/// +/// The callbacks (and the context they carry) are owned by the returned sink and released when it is +/// closed or aborted, including when this call fails. Write errors are reported when the sink is +/// closed, since the bytes are produced by a background task. +/// +/// On error, returns NULL and sets "err". +/// +/// # Safety +/// +/// `session` and `dtype` must be valid. `callbacks.write_fn` must be non-null and valid for as long +/// as the sink lives. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_callback_sink_open( + session: *const vx_session, + callbacks: vx_output_callbacks, + dtype: *const vx_dtype, + err: *mut *mut vx_error, +) -> *mut vx_callback_sink { + try_or(err, std::ptr::null_mut(), || unsafe { + callback_sink_open(session, callbacks, dtype) + }) +} + +/// Push an array into a callback sink. Does not take ownership of `array`. +/// +/// Errors if the array's DType does not match the sink's. +/// +/// # Safety +/// +/// `sink` must come from `vx_callback_sink_open` and must not have been closed or aborted. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_callback_sink_push( + sink: *mut vx_callback_sink, + array: *const vx_array, + error_out: *mut *mut vx_error, +) { + try_or_default(error_out, || { + vortex_ensure!(!sink.is_null()); + vortex_ensure!(!array.is_null()); + + let array = vx_array::as_ref(array); + let sink = unsafe { &mut *sink }; + + vortex_ensure!( + *array.dtype() == sink.dtype, + "array dtype {} does not match sink dtype {}", + array.dtype(), + sink.dtype + ); + RUNTIME + .block_on(sink.sink.send(Ok(array.clone()))) + .map_err(|e| vortex_err!("Send error: {e}")) + }) +} + +/// Close a callback sink, flushing everything pushed so far and writing the file footer. +/// +/// Consumes `sink` even when it reports an error; do not use it afterwards. +/// +/// # Safety +/// +/// `sink` must come from `vx_callback_sink_open` and must not have been closed or aborted. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_callback_sink_close( + sink: *mut vx_callback_sink, + error_out: *mut *mut vx_error, +) { + try_or_default(error_out, || { + vortex_ensure!(!sink.is_null()); + let vx_callback_sink { + sink, + writer, + dtype: _, + } = *unsafe { Box::from_raw(sink) }; + // Dropping the sender signals end of input to the write task. + drop(sink); + + RUNTIME.block_on(async { + let _summary = writer.await?; + VortexResult::Ok(()) + })?; + + Ok(()) + }) +} + +/// Abort a callback sink. No footer is written, so whatever the host received is not a valid Vortex +/// file. Do not use `sink` afterwards. +/// +/// # Safety +/// +/// `sink` must come from `vx_callback_sink_open` and must not have been closed or aborted. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_callback_sink_abort(sink: *mut vx_callback_sink) { + if sink.is_null() { + return; + } + drop(unsafe { Box::from_raw(sink) }); Review Comment: **[P1] Drain write callbacks before abort returns.** Dropping the sink only requests cancellation; it does not wait for an executing synchronous write callback. With concurrent Vortex operations driving the shared runtime, that callback can continue after abort returns. `SingleFileWriter::Abort` also closes the `OutputStream` before destroying the format writer, allowing `Write`/`fwrite` to race with `Close`/`fclose`. Please ensure cancellation and callback completion happen before the stream is closed; retaining the context alone does not keep the underlying stream open. An isolated test of the underlying task primitive confirmed that its callback can outlive task destruction. ########## src/paimon/format/vortex/vortex_stats_extractor.cpp: ########## @@ -0,0 +1,133 @@ +/* + * 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_stats_extractor.h" + +#include <optional> +#include <utility> + +#include "arrow/api.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/math.h" +#include "paimon/defs.h" +#include "paimon/format/column_stats.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/fs/file_system.h" + +namespace paimon::vortex { + +Result<std::pair<ColumnStatsVector, FormatStatsExtractor::FileInfo>> +VortexStatsExtractor::ExtractWithFileInfo(const std::shared_ptr<FileSystem>& file_system, + const std::string& path, + const std::shared_ptr<MemoryPool>& pool) { + if (file_system == nullptr) { + return Status::Invalid("Vortex stats extractor requires a file system"); + } + if (schema_ == nullptr) { + return Status::Invalid("Vortex stats extractor has no schema"); + } + (void)pool; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input, file_system->Open(path)); + // Only the row count is needed; Vortex keeps it in the file footer. + PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Vortex input length")); + auto length = static_cast<uint64_t>(signed_length); + auto input_context = std::make_shared<VortexInputContext>(input); + + VxSessionPtr session(vx_session_new(), vx_session_free); + if (session == nullptr) { + return Status::IOError("failed to create Vortex session"); + } + vx_error* error = nullptr; + VxDataSourcePtr data_source( + 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 file for stats", error, + input_context->GetCallbackStatus()); + } + vx_estimate row_count{}; + vx_data_source_get_row_count(data_source.get(), &row_count); + if (row_count.type == VX_ESTIMATE_UNKNOWN) { + return Status::Invalid("Vortex file did not report a row count"); + } + const auto rows = static_cast<int64_t>(row_count.estimate); + // Vortex keeps its statistics internal, so every field reports an unknown entry; consumers + // index this vector by field position and require one entry per write-schema field. + ColumnStatsVector stats; + stats.reserve(schema_->num_fields()); + for (const std::shared_ptr<arrow::Field>& field : schema_->fields()) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ColumnStats> column_stats, + CreateUnknownStats(field->type())); + stats.push_back(std::move(column_stats)); + } + return std::make_pair(std::move(stats), FileInfo(rows)); +} + +Result<std::unique_ptr<ColumnStats>> VortexStatsExtractor::CreateUnknownStats( + const std::shared_ptr<arrow::DataType>& type) const { + switch (type->id()) { + case arrow::Type::BOOL: + return ColumnStats::CreateBooleanColumnStats(std::nullopt, std::nullopt, std::nullopt); + case arrow::Type::INT8: + return ColumnStats::CreateTinyIntColumnStats(std::nullopt, std::nullopt, std::nullopt); + case arrow::Type::INT16: + return ColumnStats::CreateSmallIntColumnStats(std::nullopt, std::nullopt, std::nullopt); + case arrow::Type::INT32: + return ColumnStats::CreateIntColumnStats(std::nullopt, std::nullopt, std::nullopt); + case arrow::Type::INT64: + return ColumnStats::CreateBigIntColumnStats(std::nullopt, std::nullopt, std::nullopt); Review Comment: **[P2] Handle TIME columns in the statistics extractor.** Schema validation accepts `TIME32`, but this switch falls through to `Invalid` for that type. Append-table finalization always extracts statistics, so writing a valid table with a top-level `TIME(3)` column fails when producing its commit metadata. The extracted production helper reproduced `cannot create unknown Vortex statistics for type time32[ms]`. Please add the corresponding unknown-statistics entry and a TIME write/finalize regression test. ########## src/paimon/format/vortex/vortex_file_batch_reader.cpp: ########## @@ -0,0 +1,457 @@ +/* + * 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_file_batch_reader.h" + +#include <algorithm> +#include <limits> +#include <string> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/math.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/format/vortex/vortex_ffi_util.h" +#include "paimon/fs/file_system.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<arrow::DataType> NormalizeViewType(const std::shared_ptr<arrow::DataType>& 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<arrow::Field>& 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<const arrow::FixedSizeListType&>(*type).list_size()); + default: + return type; + } +} + +std::shared_ptr<arrow::Schema> NormalizeViewSchema(const std::shared_ptr<arrow::Schema>& schema) { + arrow::FieldVector fields; + fields.reserve(schema->num_fields()); + for (const std::shared_ptr<arrow::Field>& 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<std::shared_ptr<arrow::Array>> NormalizeViewArray(const std::shared_ptr<arrow::Array>& 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<const arrow::StringViewArray&>(*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<arrow::Array> out; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); + return out; + } + case arrow::Type::BINARY_VIEW: { + const auto& view = checked_cast<const arrow::BinaryViewArray&>(*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<arrow::Array> out; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); + return out; + } + case arrow::Type::STRUCT: { + const auto& struct_array = checked_cast<const arrow::StructArray&>(*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<arrow::Array>& child = struct_array.field(i); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> 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::StructArray>( + 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<const arrow::ListArray&>(*array); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> values, + NormalizeViewArray(list.values(), pool)); + if (values.get() == list.values().get()) { + return array; + } + return std::make_shared<arrow::ListArray>( + 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<const arrow::FixedSizeListArray&>(*array); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> values, + NormalizeViewArray(list.values(), pool)); + if (values.get() == list.values().get()) { + return array; + } + const auto& fsl_type = checked_cast<const arrow::FixedSizeListType&>(*array->type()); + return std::make_shared<arrow::FixedSizeListArray>( + 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; + } +} + +} // namespace + +VortexFileBatchReader::VortexFileBatchReader( + const std::shared_ptr<InputStream>& input, int32_t batch_size, + std::shared_ptr<VortexInputContext> input_context, VxSessionPtr session, + VxDataSourcePtr data_source, VxScanPtr scan, const std::shared_ptr<arrow::Schema>& file_schema, + const std::shared_ptr<arrow::DataType>& struct_type, uint64_t total_rows, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) + : 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), + total_rows_(total_rows), + arrow_pool_(arrow_pool), + metrics_(std::make_shared<MetricsImpl>()) {} + +Result<std::unique_ptr<VortexFileBatchReader>> VortexFileBatchReader::Create( + const std::shared_ptr<InputStream>& input, int32_t batch_size, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (input == nullptr || arrow_pool == nullptr || batch_size <= 0) { + return Status::Invalid( + "Vortex reader requires non-null input and Arrow memory pool, and positive batch size"); + } + PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Vortex input length")); + auto length = static_cast<uint64_t>(signed_length); + auto input_context = std::make_shared<VortexInputContext>(input); + + VxSessionPtr session(vx_session_new(), vx_session_free); + if (session == nullptr) { + return Status::IOError("failed to create Vortex session"); + } + vx_error* error = nullptr; + VxDataSourcePtr data_source( + 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, + input_context->GetCallbackStatus()); + } + + // The data source dtype is available without consuming a scan. + // NOTE(vortex 0.75): vx_data_source_dtype returns a BORROWED pointer (no refcount bump), so it + // must NOT be freed here: freeing decrements the DType Arc and segfaults later in the scan. + // Upstream >=0.77 returns an owned clone that MUST be freed with vx_dtype_free. + const vx_dtype* dtype = vx_data_source_dtype(data_source.get()); + if (dtype == nullptr) { + return Status::IOError("failed to read Vortex data source dtype"); + } + ::ArrowSchema ffi_schema = {}; + error = nullptr; + if (vx_dtype_to_arrow_schema(dtype, &ffi_schema, &error) != 0) { + return VortexFfiError("convert Vortex dtype to Arrow schema", error); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> 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<arrow::DataType> struct_type = arrow::struct_(file_schema->fields()); + + vx_estimate row_count{}; + vx_data_source_get_row_count(data_source.get(), &row_count); + if (row_count.type == VX_ESTIMATE_UNKNOWN) { + return Status::Invalid("Vortex data source did not report a row count"); + } + const uint64_t total_rows = row_count.estimate; + + // The scan is created lazily on first read (and re-created by SetReadSchema), since a Vortex + // scan may be consumed only once. + return std::unique_ptr<VortexFileBatchReader>(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)); +} + +VortexFileBatchReader::~VortexFileBatchReader() { + CloseInternal(); +} + +Result<bool> VortexFileBatchReader::OpenNextPartitionStream() { + 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 + options.ordered = true; + VxScanPtr scan( + vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), + vx_scan_free); + if (scan == nullptr) { + return VortexCallbackError("create Vortex scan", error, + input_context_->GetCallbackStatus()); + } + scan_ = std::move(scan); + } + vx_error* error = nullptr; + VxPartitionPtr partition(vx_scan_next_partition(scan_.get(), &error), vx_partition_free); + if (error != nullptr) { + return VortexCallbackError("advance Vortex scan partition", error, + input_context_->GetCallbackStatus()); + } + if (partition == nullptr) { + return false; // Scan exhausted. + } + // vx_partition_scan_arrow consumes the partition (even on error), so release ownership here. + ::ArrowArrayStream stream{}; + error = nullptr; + if (vx_partition_scan_arrow(session_.get(), partition.release(), &stream, &error) != 0) { + return VortexCallbackError("scan Vortex partition to Arrow", error, + input_context_->GetCallbackStatus()); + } + current_stream_ = stream; + stream_active_ = true; + return true; +} + +Result<std::shared_ptr<arrow::Array>> VortexFileBatchReader::ReadNextArray() { + while (true) { + if (!stream_active_) { + PAIMON_ASSIGN_OR_RAISE(bool opened, OpenNextPartitionStream()); + if (!opened) { + return std::shared_ptr<arrow::Array>(); // End of scan. + } + } + ::ArrowArray ffi_array = {}; + const int32_t rc = current_stream_.get_next(¤t_stream_, &ffi_array); + if (rc != 0) { + // get_last_error returns a pointer owned by the stream; the Arrow C stream contract + // invalidates it once the stream is released, so copy it before ReleaseStream(). + const char* message = current_stream_.get_last_error(¤t_stream_); + std::string message_copy = message == nullptr ? "unknown" : message; + // A read failure surfaces here as an opaque stream error, so prefer the status the IO + // callback stashed. + Status callback_status = input_context_->GetCallbackStatus(); + ReleaseStream(); + if (!callback_status.ok()) { + return callback_status.WithMessage("read Vortex batch", ": ", + callback_status.message()); + } + return Status::IOError("Vortex Arrow stream error: ", message_copy); + } + if (ffi_array.release == nullptr) { + ReleaseStream(); // Partition exhausted; try the next one. + continue; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(&ffi_array, struct_type_)); + if (array->length() == 0) { + continue; // Skip empty batches. + } + return array; + } +} + +Result<std::shared_ptr<arrow::Array>> VortexFileBatchReader::AlignToReadSchema( + const std::shared_ptr<arrow::Array>& array) const { + if (read_schema_ == nullptr) { + return array; + } + // TODO(zhouyc-ali): push top-level column pruning into the Vortex scan + return NestedProjectionUtils::AlignArrayToReadType( + array, arrow::struct_(read_schema_->fields()), arrow_pool_.get()); Review Comment: **[P2] Normalize timestamp timezones before schema alignment.** Vortex preserves the stored timestamp timezone, while Paimon reconstructs `TIMESTAMP_LTZ` using the reader's local timezone. `AlignArrayToReadType` rejects that difference: UTC data requested as `Asia/Shanghai` fails with `unsupported leaf type change`. This prevents reading the Java fixture's `TIMESTAMP_LTZ` column outside UTC; its compatibility test masks the issue with `TimezoneGuard("UTC")`. Please retype timezone-bearing timestamps without changing their epoch values before alignment, including nested columns, and test a non-UTC reader. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
