diff --git a/vortex-duckdb/cpp/copy_function.cpp b/vortex-duckdb/cpp/copy_function.cpp index 4740b92bbc3..3076e3d938a 100644 --- a/vortex-duckdb/cpp/copy_function.cpp +++ b/vortex-duckdb/cpp/copy_function.cpp @@ -9,6 +9,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" using namespace duckdb; @@ -17,6 +18,8 @@ struct CopyBindData final : TableFunctionData { CopyBindData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { } unique_ptr ffi_data; + // Column names in write order, used to key WRITTEN_FILE_STATISTICS. + vector column_names; }; struct CopyGlobalData final : GlobalFunctionData { @@ -24,6 +27,9 @@ struct CopyGlobalData final : GlobalFunctionData { } unique_ptr ffi_data; + // Non-owning; set in copy_to_get_written_statistics (before the write) and filled in + // copy_to_finalize (after the write). Null when the plan does not request statistics. + CopyFunctionFileStatistics *written_stats = nullptr; }; unique_ptr copy_to_bind(ClientContext &, @@ -53,7 +59,9 @@ 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)); + auto bind = make_uniq(std::move(cdata)); + bind->column_names = column_names; + return bind; } unique_ptr @@ -86,13 +94,77 @@ void copy_to_sink(ExecutionContext &, } } -void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstate) { - void *const ffi_global = gstate.Cast().ffi_data->DataPtr(); +// Called before the write (right after global init) when the plan requests +// WRITTEN_FILE_STATISTICS. We stash the target struct and fill it at finalize, once the +// footer statistics exist - mirroring the parquet writer's store-pointer-then-fill pattern. +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_data->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; + } + // Fill the statistics captured in copy_to_get_written_statistics from the finished write's + // footer. min/max come back as owned duckdb values; DuckLake expects their string form. + 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"); + } + 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 < names.size(); 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)); + } + // No statistics for this column (e.g. a type without min/max); skip it. + continue; + } + 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)); + } } extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) { @@ -108,6 +180,7 @@ extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) }; fn.copy_to_sink = copy_to_sink; fn.copy_to_finalize = copy_to_finalize; + fn.copy_to_get_written_statistics = copy_to_get_written_statistics; fn.extension = "vortex"; // TODO(joe): expose this via c our api diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 65a1a3f1dd3..52414b96d08 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -76,6 +76,30 @@ typedef struct { bool has_null; } 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. `min`/`max` are owned +// duckdb_value handles (null if absent) that the caller must destroy. +typedef struct { + 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; + const idx_t INVALID_IDX = UINT64_MAX; typedef struct { diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 96f5804d3b4..64c9f3f03c7 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -94,6 +94,16 @@ void duckdb_copy_function_copy_to_sink(const void *bind_data, extern void duckdb_copy_function_copy_to_finalize(void *global_data, duckdb_vx_error *error_out); +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 f158b2e45c5..5ed8e98a2f7 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -22,6 +22,8 @@ use vortex::dtype::StructFields; use vortex::error::VortexExpect; use vortex::error::VortexResult; 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::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::runtime::current::CurrentThreadWorkerPool; 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,14 @@ pub struct CopyFunctionBind { } assert_impl_all!(CopyFunctionBind: Send, Clone); +/// Captured in `copy_to_finalize` and read back by the WRITTEN_FILE_STATISTICS path. 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 +69,7 @@ assert_impl_all!(CopyFunctionBind: Send, Clone); /// flushed to disk. pub struct CopyFunctionGlobal { write_task: Mutex>>>, + finished: Mutex>, sink: Option>>, // Pool of background workers helping to drive the write task. // Note that this is optional and without it, we would only drive the task when DuckDB calls @@ -113,11 +128,125 @@ 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 or the column has no statistics; `Err` if a scalar could not be converted. +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) +} + +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 Some(file_stats) = summary.footer().statistics() else { + return Ok(None); + }; + let stats_sets = file_stats.stats_sets(); + if column_index >= stats_sets.len() { + return Ok(None); + } + let stats = &stats_sets[column_index]; + let dtype = &file_stats.dtypes()[column_index]; + + Ok(Some(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, @@ -163,6 +292,7 @@ pub fn copy_to_initialize_global( Ok(CopyFunctionGlobal { worker_pool, write_task: Mutex::new(Some(write_task)), + finished: Mutex::new(None), sink: Some(sink), }) } diff --git a/vortex-duckdb/src/e2e_test/copy_statistics_test.rs b/vortex-duckdb/src/e2e_test/copy_statistics_test.rs new file mode 100644 index 00000000000..cd2779e8f5f --- /dev/null +++ b/vortex-duckdb/src/e2e_test/copy_statistics_test.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! End-to-end tests for the vortex COPY function's WRITTEN_FILE_STATISTICS support, driven through +//! DuckDB with `COPY … (FORMAT vortex, RETURN_STATS)`. + +use num_traits::AsPrimitive; +use tempfile::NamedTempFile; + +use crate::duckdb::Connection; +use crate::duckdb::Database; + +fn database_connection() -> Connection { + let db = Database::open_in_memory().unwrap(); + crate::initialize(&db).unwrap(); + db.connect().unwrap() +} + +/// `RETURN_STATS` binds only because the vortex copy function now implements +/// `copy_to_get_written_statistics`; running it exercises the whole path through DuckDB +/// (bind, the C++ fill trampoline, and the Rust getters) and returns the six-column +/// WRITTEN_FILE_STATISTICS schema. The nested per-column statistics map is validated against +/// DuckLake in the duckdb-vortex integration tests; here we assert the file-level statistics. +#[test] +fn copy_return_stats_reports_file_statistics() { + let conn = database_connection(); + let file = NamedTempFile::with_suffix(".vortex").unwrap(); + let path = file.path().to_string_lossy(); + + let result = conn + .query(&format!( + "COPY (SELECT * FROM (VALUES (1, 'a', 1.5), (2, 'b', NULL)) t(i, s, d)) \ + TO '{path}' (FORMAT vortex, RETURN_STATS)" + )) + .unwrap(); + + // filename, count, file_size_bytes, footer_size, column_statistics, partition_keys + assert_eq!(result.column_count(), 6); + + let chunk = result.into_iter().next().unwrap(); + let len = chunk.len().as_(); + // count and file_size_bytes are UBIGINT. + let count = chunk.get_vector(1).as_slice_with_len::(len)[0]; + let file_size = chunk.get_vector(2).as_slice_with_len::(len)[0]; + + assert_eq!(count, 2); + assert!(file_size > 0); +} + +/// A table with nested (struct/list) columns must not crash the statistics hook. Vortex reports +/// statistics for top-level fields only, so this asserts the COPY succeeds and returns rather than +/// checking per-leaf statistics. +#[test] +fn copy_return_stats_handles_nested_columns() { + let conn = database_connection(); + let file = NamedTempFile::with_suffix(".vortex").unwrap(); + let path = file.path().to_string_lossy(); + + let result = conn + .query(&format!( + "COPY (SELECT {{'a': 1, 'b': 2}} AS st, [1, 2, 3] AS lst, 7 AS i) \ + TO '{path}' (FORMAT vortex, RETURN_STATS)" + )) + .unwrap(); + let chunk = result.into_iter().next().unwrap(); + let count = chunk + .get_vector(1) + .as_slice_with_len::(chunk.len().as_())[0]; + assert_eq!(count, 1); +} + +/// Without `RETURN_STATS` the statistics hook is never invoked; a plain vortex COPY must still +/// succeed unchanged. +#[test] +fn copy_without_return_stats_still_works() { + let conn = database_connection(); + let file = NamedTempFile::with_suffix(".vortex").unwrap(); + let path = file.path().to_string_lossy(); + + conn.query(&format!( + "COPY (SELECT 1 AS i, 'a' AS s) TO '{path}' (FORMAT vortex)" + )) + .unwrap(); +} diff --git a/vortex-duckdb/src/e2e_test/mod.rs b/vortex-duckdb/src/e2e_test/mod.rs index 3bba671e79b..ae057a48fec 100644 --- a/vortex-duckdb/src/e2e_test/mod.rs +++ b/vortex-duckdb/src/e2e_test/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(test)] +mod copy_statistics_test; #[cfg(test)] mod spatial_pushdown_test; #[cfg(test)] diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 2a5e9316434..b00b1b7ea3b 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -16,6 +16,8 @@ use crate::copy::copy_to_bind; use crate::copy::copy_to_finalize; use crate::copy::copy_to_initialize_global; use crate::copy::copy_to_sink; +use crate::copy::written_column_stats; +use crate::copy::written_file_stats; use crate::cpp; use crate::duckdb::AggregatePushdownInput; use crate::duckdb::BindInput; @@ -333,3 +335,55 @@ pub unsafe extern "C-unwind" fn duckdb_copy_function_copy_to_finalize( .vortex_expect("bind_data null pointer"); try_or(error_out, || copy_to_finalize(global_data)) } + +/// Fill file-level statistics of the just-written Vortex file. Returns `false` if the +/// file has not been finalized (no statistics available), leaving `out` untouched. +#[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 +} + +/// Fill per-column statistics for `column_index` of the just-written Vortex file. +/// `min`/`max` are owned duckdb values the caller must destroy. Returns `false` if +/// no statistics are available for that column. +#[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"); + // Converting a Vortex scalar to a DuckDB value can fail; surface it through error_out like the + // rest of ffi.rs rather than swallowing it as "no statistics". + 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) + }) +}