diff --git a/vortex-duckdb/cpp/copy_function.cpp b/vortex-duckdb/cpp/copy_function.cpp index 6f639f42f3a..840d159f266 100644 --- a/vortex-duckdb/cpp/copy_function.cpp +++ b/vortex-duckdb/cpp/copy_function.cpp @@ -11,6 +11,7 @@ #include "duckdb/main/capi/capi_internal.hpp" #include "duckdb/main/client_context.hpp" #include "duckdb/main/connection.hpp" +#include "duckdb/parser/keyword_helper.hpp" #include "duckdb/parser/parsed_data/create_copy_function_info.hpp" unique_ptr copy_to_bind(ClientContext &, @@ -39,7 +40,7 @@ unique_ptr copy_to_bind(ClientContext &, throw BinderException(IntoErrString(error_out)); } auto cdata = unique_ptr(reinterpret_cast(ffi_bind_data)); - return make_uniq(std::move(cdata)); + return make_uniq(std::move(cdata), names); } unique_ptr @@ -77,15 +78,78 @@ void copy_to_sink(ExecutionContext &, } } -void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstate) { - const VortexCopyGlobalState &global = gstate.Cast(); +// CopyToFileInfo::file_stats is owned by the operator's sink state, which outlives gstate. +void copy_to_get_written_statistics(ClientContext &, + FunctionData &, + GlobalFunctionData &gstate, + CopyFunctionFileStatistics &statistics) { + gstate.Cast().written_stats = &statistics; +} +void copy_to_finalize(ClientContext &, FunctionData &bind_data, GlobalFunctionData &gstate) { + auto &global = gstate.Cast(); void *const ffi_global = global.ffi_global->DataPtr(); duckdb_vx_error error_out = nullptr; duckdb_copy_function_copy_to_finalize(ffi_global, &error_out); if (error_out) { throw ExecutorException(IntoErrString(error_out)); } + + if (!global.written_stats) { + return; + } + auto &names = bind_data.Cast().column_names; + duckdb_vx_written_file_statistics file_stats; + if (!duckdb_copy_function_get_written_file_statistics(ffi_global, &file_stats)) { + // Statistics were requested (written_stats is set) but the finished write produced none; + // that is an internal inconsistency, not a silently empty result. + throw InternalException("vortex COPY: written statistics were requested but not produced"); + } + if (file_stats.num_columns != names.size()) { + throw InternalException("vortex COPY: %llu statistics columns for %llu written columns", + file_stats.num_columns, + names.size()); + } + D_ASSERT(global.written_stats != nullptr); + global.written_stats->row_count = file_stats.row_count; + global.written_stats->file_size_bytes = file_stats.file_size_bytes; + global.written_stats->footer_size_bytes = Value::UBIGINT(file_stats.footer_size_bytes); + // Keyed by top-level column name only. The vortex footer reports one statistics set per + // top-level field, so nested struct/list leaf columns get no statistics here (unlike parquet, + // which recurses to leaf paths). Flat tables are fully covered. + for (idx_t i = 0; i < file_stats.num_columns; i++) { + duckdb_vx_written_column_statistics col_stats {}; + duckdb_vx_error col_error = nullptr; + if (!duckdb_copy_function_get_written_column_statistics(ffi_global, i, &col_stats, &col_error)) { + if (col_error) { + throw ExecutorException(IntoErrString(col_error)); + } + throw InternalException("vortex COPY: no statistics for column %llu after finalize", i); + } + case_insensitive_map_t column; + column["num_values"] = Value::UBIGINT(col_stats.num_values); + if (col_stats.has_column_size) { + column["column_size_bytes"] = Value::UBIGINT(col_stats.column_size_bytes); + } + if (col_stats.has_null_count) { + column["null_count"] = Value::UBIGINT(col_stats.null_count); + } + if (col_stats.min) { + column["min"] = Value(reinterpret_cast(col_stats.min)->ToString()); + duckdb_destroy_value(&col_stats.min); + } + if (col_stats.max) { + column["max"] = Value(reinterpret_cast(col_stats.max)->ToString()); + duckdb_destroy_value(&col_stats.max); + } + if (col_stats.has_nan_stat) { + column["has_nan"] = Value::BOOLEAN(col_stats.contains_nan); + } + // DuckLake keys column statistics by a quoted, dot-separated path (see + // DuckLakeUtil::ParseQuotedList); match the parquet writer, which quotes each name. + global.written_stats->column_statistics.emplace(KeywordHelper::WriteQuoted(names[i], '"'), + std::move(column)); + } } unique_ptr copy_to_prepare_batch(ClientContext &, @@ -143,6 +207,7 @@ extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) fn.copy_to_finalize = copy_to_finalize; fn.prepare_batch = copy_to_prepare_batch; fn.flush_batch = copy_to_flush_batch; + fn.copy_to_get_written_statistics = copy_to_get_written_statistics; fn.extension = "vortex"; fn.execution_mode = [](bool preserve_insertion_order, bool supports_batch_index) { diff --git a/vortex-duckdb/cpp/include/copy_function.hpp b/vortex-duckdb/cpp/include/copy_function.hpp index 02307280dc2..ec8cccd7048 100644 --- a/vortex-duckdb/cpp/include/copy_function.hpp +++ b/vortex-duckdb/cpp/include/copy_function.hpp @@ -8,15 +8,20 @@ using namespace duckdb; struct VortexCopyBindData final : TableFunctionData { - VortexCopyBindData(unique_ptr ffi_bind) : ffi_bind(std::move(ffi_bind)) { + VortexCopyBindData(unique_ptr ffi_bind, vector column_names) + : ffi_bind(std::move(ffi_bind)), column_names(std::move(column_names)) { } unique_ptr ffi_bind; + // Column names in write order, used to key WRITTEN_FILE_STATISTICS. + vector column_names; }; struct VortexCopyGlobalState final : GlobalFunctionData { VortexCopyGlobalState(unique_ptr ffi_global) : ffi_global(std::move(ffi_global)) { } unique_ptr ffi_global; + // Non-owning; null when the plan does not request statistics. + CopyFunctionFileStatistics *written_stats = nullptr; }; struct VortexCopyPreparedBatchData final : PreparedBatchData { diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index f591a00f46a..954cae2dec4 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -51,6 +51,30 @@ typedef struct { duckdb_logical_type type; } duckdb_column_statistics; +// File-level statistics of a written Vortex file, for the DuckLake +// WRITTEN_FILE_STATISTICS return path. Filled by Rust from the WriteSummary. +typedef struct { + uint64_t row_count; + uint64_t file_size_bytes; + uint64_t footer_size_bytes; + uint64_t num_columns; +} duckdb_vx_written_file_statistics; + +// Per-column statistics of a written Vortex file. +typedef struct { + // Owned values, null if absent; the caller must destroy them. + duckdb_value min; + duckdb_value max; + bool has_null_count; + uint64_t null_count; + uint64_t num_values; + bool has_column_size; + uint64_t column_size_bytes; + // Whether a NaN-count statistic was available (float columns), and whether it saw any NaN. + bool has_nan_stat; + bool contains_nan; +} duckdb_vx_written_column_statistics; + duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db); typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input; diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index e096926d454..daf971c53ac 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -124,6 +124,16 @@ void duckdb_copy_function_flush_batch(const void *global, const void *batch, duckdb_vx_error *error); +extern +bool duckdb_copy_function_get_written_file_statistics(const void *global_data, + duckdb_vx_written_file_statistics *out); + +extern +bool duckdb_copy_function_get_written_column_statistics(const void *global_data, + size_t column_index, + duckdb_vx_written_column_statistics *out, + duckdb_vx_error *error_out); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/vortex-duckdb/src/copy.rs b/vortex-duckdb/src/copy.rs index 688999eb160..6d3dbbbd952 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -23,6 +23,8 @@ use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; +use vortex::expr::stats::Precision; +use vortex::expr::stats::Stat; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::file::multi::parse_uri_or_path; @@ -32,14 +34,18 @@ use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::scalar::Scalar; +use vortex::scalar::ScalarValue; use crate::REGISTRY; use crate::RUNTIME; use crate::SESSION; use crate::convert::FromLogicalType; +use crate::convert::ToDuckDBScalar; use crate::convert::data_chunk_to_vortex; use crate::duckdb::DataChunkRef; use crate::duckdb::LogicalTypeRef; +use crate::duckdb::Value; #[derive(Clone)] pub struct CopyFunctionBind { @@ -48,6 +54,13 @@ pub struct CopyFunctionBind { } assert_impl_all!(CopyFunctionBind: Send, Clone); +/// The per-column compressed sizes are computed once here rather than per column, since DuckDB +/// queries statistics one column at a time. +struct FinishedWrite { + summary: WriteSummary, + column_sizes: Vec, +} + /// Write to a file has two phases, writing data chunks and then closing the file. /// We use a spawned tokio task to actually compress arrays and write it to disk. /// Each chunk is pushed into the sink and read from the task. @@ -55,6 +68,7 @@ assert_impl_all!(CopyFunctionBind: Send, Clone); /// flushed to disk. pub struct CopyFunctionGlobal { write_task: Mutex>>>, + finished: Mutex>, sink: Option>>, } assert_impl_all!(CopyFunctionGlobal: Send, Sync); @@ -144,11 +158,129 @@ pub fn copy_to_finalize(init_global: &mut CopyFunctionGlobal) -> VortexResult<() .lock() .take() .vortex_expect("no file to close"); - task.await?; + // Keep the write summary (footer + size) so DuckLake can read per-file statistics back + // without re-opening the file. Compute the per-column compressed sizes once, up front. + let summary = task.await?; + let column_sizes = summary.compressed_column_sizes().unwrap_or_default(); + *init_global.finished.lock() = Some(FinishedWrite { + summary, + column_sizes, + }); Ok(()) }) } +/// File-level statistics of the written Vortex file, for the WRITTEN_FILE_STATISTICS return path. +pub(crate) struct WrittenFileStats { + pub row_count: u64, + pub file_size_bytes: u64, + pub footer_size_bytes: u64, + pub num_columns: usize, +} + +/// Per-column statistics of the written Vortex file. `min`/`max` are DuckDB values converted from +/// the Vortex scalar; every field is optional and omitted when the statistic is not available. +pub(crate) struct WrittenColumnStats { + pub min: Option, + pub max: Option, + pub null_count: Option, + pub num_values: u64, + pub column_size_bytes: Option, + pub has_nan: Option, +} + +/// Read file-level statistics back from the finished write. `None` before finalize. +pub(crate) fn written_file_stats(global: &CopyFunctionGlobal) -> Option { + let guard = global.finished.lock(); + Some(file_stats_from_summary(&guard.as_ref()?.summary)) +} + +/// Read per-column statistics for `column_index` from the finished write. `Ok(None)` if the file is +/// not finalized. +pub(crate) fn written_column_stats( + global: &CopyFunctionGlobal, + column_index: usize, +) -> VortexResult> { + let guard = global.finished.lock(); + let Some(finished) = guard.as_ref() else { + return Ok(None); + }; + column_stats_from_summary(&finished.summary, column_index, &finished.column_sizes).map(Some) +} + +fn file_stats_from_summary(summary: &WriteSummary) -> WrittenFileStats { + let num_columns = summary + .footer() + .statistics() + .map_or(0, |s| s.stats_sets().len()); + WrittenFileStats { + row_count: summary.row_count(), + file_size_bytes: summary.size(), + // Vortex has no separate footer-size hint; 0 means "read the footer normally". + footer_size_bytes: 0, + num_columns, + } +} + +/// Per-column statistics from a finished write's summary and its precomputed compressed sizes +/// (`column_sizes`, indexed the same as the footer's stats sets). +/// +/// Only top-level columns are covered: the footer exposes one statistics set per top-level field, +/// so nested struct/list leaf columns are not reported (parquet, by contrast, recurses to leaf +/// paths). Flat tables - the common DuckLake case - are fully covered. +fn column_stats_from_summary( + summary: &WriteSummary, + column_index: usize, + column_sizes: &[u64], +) -> VortexResult { + let file_stats = summary + .footer() + .statistics() + .ok_or_else(|| vortex_err!("written file has no statistics"))?; + let stats_sets = file_stats.stats_sets(); + if column_index >= stats_sets.len() { + vortex_bail!( + "column index {column_index} out of range for {} statistics sets", + stats_sets.len() + ); + } + let stats = &stats_sets[column_index]; + let dtype = &file_stats.dtypes()[column_index]; + + Ok(WrittenColumnStats { + min: exact_scalar_to_duckdb(stats.get(Stat::Min), dtype)?, + max: exact_scalar_to_duckdb(stats.get(Stat::Max), dtype)?, + null_count: exact_u64(stats.get(Stat::NullCount)), + // NaNCount is exact only for float columns, so this is emitted just for them (as in parquet). + has_nan: exact_u64(stats.get(Stat::NaNCount)).map(|count| count > 0), + num_values: summary.row_count(), + // On-disk compressed size; excludes bytes not attributable to a column (e.g. struct validity). + column_size_bytes: column_sizes.get(column_index).copied(), + }) +} + +/// Convert an exact scalar statistic to a DuckDB value, propagating a conversion failure rather than +/// dropping it. `Ok(None)` when the statistic is not exactly known. +fn exact_scalar_to_duckdb( + stat: Precision, + dtype: &DType, +) -> VortexResult> { + match stat { + Precision::Exact(value) => Ok(Some( + Scalar::try_new(dtype.clone(), Some(value))?.try_to_duckdb_scalar()?, + )), + _ => Ok(None), + } +} + +/// Extract an exact `u64` statistic (e.g. a count), or `None` if not exactly known. +fn exact_u64(stat: Precision) -> Option { + match stat { + Precision::Exact(value) => value.as_primitive().as_u64(), + _ => None, + } +} + pub fn copy_to_initialize_global( bind_data: &CopyFunctionBind, file_path: String, @@ -191,6 +323,48 @@ pub fn copy_to_initialize_global( Ok(CopyFunctionGlobal { write_task: Mutex::new(Some(write_task)), + finished: Mutex::new(None), sink: Some(sink), }) } + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::arrays::StructArray; + use vortex::array::stats::PRUNING_STATS; + use vortex::buffer::ByteBufferMut; + use vortex::buffer::buffer; + + use super::*; + + /// Writes a one-column file and returns its summary, with `file_statistics` controlling which + /// statistics the footer carries (empty means none at all). + fn write_summary(file_statistics: Vec) -> WriteSummary { + RUNTIME.block_on(async { + let array = StructArray::from_fields(&[("i", buffer![1u32, 2, 3].into_array())]) + .unwrap() + .into_array(); + let mut buf = ByteBufferMut::empty(); + let mut writer = SESSION + .write_options() + .with_file_statistics(file_statistics) + .writer(&mut buf, array.dtype().clone()); + writer.push(array).await.unwrap(); + writer.finish().await.unwrap() + }) + } + + #[test] + fn column_stats_out_of_range_is_an_error() { + let summary = write_summary(PRUNING_STATS.to_vec()); + assert!(column_stats_from_summary(&summary, 0, &[]).is_ok()); + assert!(column_stats_from_summary(&summary, 1, &[]).is_err()); + } + + #[test] + fn column_stats_without_file_statistics_is_an_error() { + let summary = write_summary(vec![]); + assert!(column_stats_from_summary(&summary, 0, &[]).is_err()); + } +} diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 160f0217d71..ec5b5d890d6 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -20,6 +20,8 @@ use crate::copy::copy_to_initialize_global; use crate::copy::copy_to_sink; use crate::copy::flush_batch; use crate::copy::prepare_batch_push; +use crate::copy::written_column_stats; +use crate::copy::written_file_stats; use crate::cpp; use crate::duckdb::AggregatePushdownInput; use crate::duckdb::BindResult; @@ -407,3 +409,48 @@ pub unsafe extern "C-unwind" fn duckdb_copy_function_flush_batch( let batch = unsafe { batch.cast::().as_ref() }.vortex_expect("null pointer"); try_or(error, || flush_batch(global, batch)) } + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_copy_function_get_written_file_statistics( + global_data: *const c_void, + out: *mut cpp::duckdb_vx_written_file_statistics, +) -> bool { + let global_data = unsafe { global_data.cast::().as_ref() } + .vortex_expect("global_data null pointer"); + let Some(stats) = written_file_stats(global_data) else { + return false; + }; + let out = unsafe { &mut *out }; + out.row_count = stats.row_count; + out.file_size_bytes = stats.file_size_bytes; + out.footer_size_bytes = stats.footer_size_bytes; + out.num_columns = stats.num_columns as u64; + true +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_copy_function_get_written_column_statistics( + global_data: *const c_void, + column_index: usize, + out: *mut cpp::duckdb_vx_written_column_statistics, + error_out: *mut cpp::duckdb_vx_error, +) -> bool { + let global_data = unsafe { global_data.cast::().as_ref() } + .vortex_expect("global_data null pointer"); + try_or(error_out, || { + let Some(stats) = written_column_stats(global_data, column_index)? else { + return Ok(false); + }; + let out = unsafe { &mut *out }; + out.min = stats.min.map_or(ptr::null_mut(), |v| v.into_ptr()); + out.max = stats.max.map_or(ptr::null_mut(), |v| v.into_ptr()); + out.has_null_count = stats.null_count.is_some(); + out.null_count = stats.null_count.unwrap_or(0); + out.num_values = stats.num_values; + out.has_column_size = stats.column_size_bytes.is_some(); + out.column_size_bytes = stats.column_size_bytes.unwrap_or(0); + out.has_nan_stat = stats.has_nan.is_some(); + out.contains_nan = stats.has_nan.unwrap_or(false); + Ok(true) + }) +} diff --git a/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt b/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt new file mode 100644 index 00000000000..ceeaa0ac9ce --- /dev/null +++ b/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +include ../setup.slt.no + +query TIII?? +COPY (SELECT * FROM (VALUES (1, 'a', 1.5), (2, 'b', NULL)) t(i, s, d)) +TO '${WORK_DIR}/copy_return_stats.vortex' (FORMAT vortex, RETURN_STATS); +---- +${WORK_DIR}/copy_return_stats.vortex 2 4324 0 {'"d"'={column_size_bytes=580, max=1.5, min=1.5, null_count=1, num_values=2}, '"i"'={column_size_bytes=440, max=2, min=1, null_count=0, num_values=2}, '"s"'={column_size_bytes=724, max=b, min=a, null_count=0, num_values=2}} NULL + +# Vortex reports statistics per top-level field, so nested columns are covered by the file-level +# counts only; the COPY must still succeed. +query TIII?? +COPY (SELECT {'a': 1, 'b': 2} AS st, [1, 2, 3] AS lst, 7 AS i) +TO '${WORK_DIR}/copy_return_stats_nested.vortex' (FORMAT vortex, RETURN_STATS); +---- +${WORK_DIR}/copy_return_stats_nested.vortex 1 4632 0 {'"i"'={column_size_bytes=380, max=7, min=7, null_count=0, num_values=1}, '"lst"'={column_size_bytes=560, null_count=0, num_values=1}, '"st"'={column_size_bytes=872, null_count=0, num_values=1}} NULL