From 6689e8f00288f75a0d5aba5ad2cf5ccfada0d98a Mon Sep 17 00:00:00 2001 From: Paolo Valletta Date: Fri, 9 Oct 2026 08:14:36 +0200 Subject: [PATCH] feat(duckdb): read pushed-down struct extracts DuckDB can ask a table function for a struct field instead of the whole column (`ColumnIndex::IsPushdownExtract()`), which is what lets `SELECT s.a.b` avoid reading `s`. Answer the callback for struct columns -- but not when the scan has aggregates bound, since the aggregate pushdown binds column indexes itself and cannot express an extracted path -- and swap `statistics` for `statistics_extended`, which is what makes DuckDB consult that callback at all. The extended callback returns no statistics for an extracted path, as Vortex has no nested statistics. `InitializeGlobalState` now passes the flattened child paths, their offsets and the type DuckDB expects each one to be emitted as, and `Projection::new` turns a path into a chain of `get_item`s with a cast when the requested type differs from the field's own. The packed output names include the path, because DuckDB binds one output column per extracted path. `test_vortex_scan_struct_extract_projection` is ignored: DuckDB's `MultiFileColumnMapper` still rebuilds the struct for a `PUSHDOWN_EXTRACT` column, so `MultiFileReader::FinalizeChunk` executes a struct expression against a vector typed as the extracted field and raises an internal error. Supporting this for a MultiFileReader-based scan needs a DuckDB-side patch. Part of #10297 Signed-off-by: Paolo Valletta --- vortex-duckdb/cpp/include/table_function.h | 26 +++ vortex-duckdb/cpp/multi_file_reader.cpp | 28 +++ vortex-duckdb/cpp/table_function.cpp | 26 ++- vortex-duckdb/src/duckdb/table_init_input.rs | 78 ++++++++ .../src/e2e_test/vortex_scan_test.rs | 48 +++++ vortex-duckdb/src/projection.rs | 181 +++++++++++++++--- vortex-duckdb/src/table_function.rs | 2 +- 7 files changed, 365 insertions(+), 24 deletions(-) diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 5924a209fbc..12ef236e966 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -50,6 +50,32 @@ typedef struct { duckdb_vx_table_filter_set filters; duckdb_client_context client_context; + + /** + * Struct extract paths pushed down by DuckDB for the entries in + * `column_ids` where DuckDB only needs a struct field + * (`ColumnIndex::IsPushdownExtract()`). + * + * `column_extract_indexes` holds every path concatenated, in the order of + * `column_ids`, and `column_extract_offsets` has `column_ids_count + 1` + * entries giving each column the slice `[offsets[i], offsets[i + 1])`. + * A column that is read whole has an empty slice. The indexes are struct + * child positions, outermost first: `SELECT s.a.b` on `s = {a: {b: …}}` + * yields `[0, 0]`. + */ + const idx_t *column_extract_indexes; + size_t column_extract_indexes_count; + const size_t *column_extract_offsets; + + /** + * For each entry in `column_ids`, the type DuckDB expects the scan to emit + * (`ColumnIndex::GetScanType()`, which is the cast target when the query + * casts the extracted field), or NULL when the column is read whole. + * The types are borrowed from the bind data and only valid for the duration + * of the init_global call. + */ + const duckdb_logical_type *column_extract_types; + size_t column_extract_types_count; } duckdb_vx_tfunc_init_input; // Result data returned from the cardinality callback. diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp index 7de57ff907d..d4140013a1a 100644 --- a/vortex-duckdb/cpp/multi_file_reader.cpp +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -152,6 +152,29 @@ VortexReaderInterface::InitializeGlobalState(ClientContext &context, column_ids[i] = storage_index; } + // DuckDB asks for a struct field instead of the whole column when it knows + // the scan can extract it. Pass the path and the type it expects back, so + // Rust can build the extract expression: one flattened path per column, with + // offsets, and the leaf type to cast to when the query asked for one. + vector extract_indexes; + vector extract_offsets; + vector extract_types(input.column_indexes.size(), nullptr); + extract_offsets.reserve(input.column_indexes.size() + 1); + extract_offsets.push_back(0); + for (size_t i = 0; i < input.column_indexes.size(); ++i) { + const ColumnIndex *level = &input.column_indexes[i]; + if (level->IsPushdownExtract()) { + while (level->HasChildren()) { + level = &level->GetChildIndex(0); + extract_indexes.push_back(level->GetPrimaryIndex()); + } + // Borrowed from the bind data, valid for this call. + extract_types[i] = reinterpret_cast( + const_cast(&level->GetScanType())); + } + extract_offsets.push_back(extract_indexes.size()); + } + // MultiFileGlobalState projection_ids are filled only when this call // returns. Take these from a physical operator. const idx_t *projection_ids = nullptr; @@ -173,6 +196,11 @@ VortexReaderInterface::InitializeGlobalState(ClientContext &context, .projection_ids_count = projection_ids_count, .filters = reinterpret_cast(input.filters.get()), .client_context = reinterpret_cast(&context), + .column_extract_indexes = extract_indexes.empty() ? nullptr : extract_indexes.data(), + .column_extract_indexes_count = extract_indexes.size(), + .column_extract_offsets = extract_offsets.data(), + .column_extract_types = extract_types.data(), + .column_extract_types_count = extract_types.size(), }; duckdb_vx_error error_out = nullptr; diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 0f6b35f3c9c..9d62f7e7909 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -209,7 +209,31 @@ duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter return {COLUMN_IDENTIFIER_FILE_INDEX, COLUMN_IDENTIFIER_FILE_ROW_NUMBER}; }; - fn.statistics = MultiFileFunction::MultiFileScanStats; + // DuckDB only pushes a struct extract into a scan when the table function has + // no `statistics` callback (RemoveUnusedColumns::CheckPushdownExtract), so use + // the extended variant: it sees the whole column index, and can report no + // statistics for an extracted path, since Vortex has no nested statistics. + fn.statistics_extended = [](ClientContext &context, TableFunctionGetStatisticsInput &input) { + if (input.column_index.IsPushdownExtract()) { + return unique_ptr(); + } + return MultiFileFunction::MultiFileScanStats( + context, input.bind_data.get(), input.column_index.GetPrimaryIndex()); + }; + + // Only a struct column can be read as an extracted path, and never when + // DuckDB bound aggregates to this scan: the aggregate pushdown binds column + // indexes itself and cannot express an extracted path. + fn.supports_pushdown_extract = [](const FunctionData &bind_data, const LogicalIndex &col_idx) { + if (duckdb_reader_is_aggregate(get_ffi_bind(&bind_data))) { + return false; + } + const auto &bind = bind_data.Cast(); + if (col_idx.index >= bind.types.size()) { + return false; + } + return bind.types[col_idx.index].id() == LogicalTypeId::STRUCT; + }; fn.get_partition_stats = get_partition_stats; fn.get_multi_file_reader = get_multi_file_reader; diff --git a/vortex-duckdb/src/duckdb/table_init_input.rs b/vortex-duckdb/src/duckdb/table_init_input.rs index f5098c3a347..434f1ac6335 100644 --- a/vortex-duckdb/src/duckdb/table_init_input.rs +++ b/vortex-duckdb/src/duckdb/table_init_input.rs @@ -6,9 +6,56 @@ use std::fmt::Formatter; use std::fmt::Result; use crate::cpp; +use crate::duckdb::LogicalType; +use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableFilterSet; use crate::duckdb::TableFilterSetRef; +/// A struct field DuckDB asked the scan to read instead of the whole column. +#[derive(Debug, Clone, Copy)] +pub struct StructExtract<'a> { + /// Struct child indexes, outermost first. + pub indexes: &'a [u64], + /// The type DuckDB expects the scan to emit for this field, which is the + /// cast target when the query casts the extracted field itself. + pub datatype: &'a LogicalTypeRef, +} + +/// The struct extracts DuckDB pushed into the scan, one entry per column id. +#[derive(Debug, Clone, Copy, Default)] +pub struct StructExtracts<'a> { + pub(crate) indexes: &'a [u64], + pub(crate) offsets: &'a [usize], + pub(crate) datatypes: &'a [cpp::duckdb_logical_type], +} + +impl<'a> StructExtracts<'a> { + /// The extract for the column at `column`, if DuckDB asked for one. + pub fn get(&self, column: usize) -> Option> { + let start = usize::try_from(*self.offsets.get(column)?).ok()?; + let end = usize::try_from(*self.offsets.get(column + 1)?).ok()?; + if start == end { + return None; + } + + let indexes = self.indexes.get(start..end)?; + let datatype = unsafe { LogicalType::borrow(*self.datatypes.get(column)?) }; + Some(StructExtract { indexes, datatype }) + } +} + +/// Borrows a C++ array as a slice, reading a null pointer as an empty array. +/// +/// # Safety +/// +/// `ptr` must point to `len` initialized values that stay alive for `'a`. +unsafe fn borrow_array<'a, T>(ptr: *const T, len: usize) -> &'a [T] { + if ptr.is_null() { + return &[]; + } + unsafe { std::slice::from_raw_parts(ptr, len) } +} + pub struct TableInitInput<'a> { pub input: &'a cpp::duckdb_vx_tfunc_init_input, } @@ -52,4 +99,35 @@ impl<'a> TableInitInput<'a> { Some(unsafe { TableFilterSet::borrow(ptr) }) } } + + /// The struct extracts DuckDB pushed into the scan, one per column id. + /// + /// The paths and types are borrowed from the bind data and only valid for + /// the duration of the `init_global` call. + pub fn struct_extracts(&self) -> StructExtracts<'_> { + // The offsets array carries one entry per column id plus the end offset. + let offsets_len = if self.input.column_extract_offsets.is_null() { + 0 + } else { + self.input.column_ids_count + 1 + }; + + StructExtracts { + indexes: unsafe { + borrow_array( + self.input.column_extract_indexes, + self.input.column_extract_indexes_count, + ) + }, + offsets: unsafe { + borrow_array(self.input.column_extract_offsets, offsets_len) + }, + datatypes: unsafe { + borrow_array( + self.input.column_extract_types, + self.input.column_extract_types_count, + ) + }, + } + } } diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index e8c40346d24..391dfeb79af 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -1034,6 +1034,54 @@ fn test_geometry() { assert_eq!(area, 1000.0); } +/// `SELECT s.x.a` is pushed down as a struct extract: DuckDB asks the scan to read +/// the field at that child path instead of the whole `s` struct, and the scan +/// emits the extracted field directly. +/// +/// Ignored until DuckDB's `MultiFileColumnMapper` stops rebuilding the struct for a +/// `PUSHDOWN_EXTRACT` column: `MultiFileReader::FinalizeChunk` currently applies a +/// struct expression to a vector DuckDB typed as the extracted field. +#[test] +#[ignore = "needs a DuckDB patch in MultiFileColumnMapper"] +fn test_vortex_scan_struct_extract_projection() { + let file = RUNTIME.block_on(async { + let inner = StructArray::try_from_iter([ + ("a", PrimitiveArray::from_iter([1i32, 2, 3]).into_array()), + ("b", PrimitiveArray::from_iter([4i32, 5, 6]).into_array()), + ]) + .unwrap(); + let top = StructArray::try_from_iter([ + ("x", inner.into_array()), + ("y", PrimitiveArray::from_iter([7i32, 8, 9]).into_array()), + ]) + .unwrap(); + + write_single_column_vortex_file("s", top).await + }); + + let conn = database_connection(); + let file_path = file.path().to_string_lossy(); + + // One level down, and two levels down, both requested by child index. + let result = conn + .query(&format!("SELECT s.y, s.x.b, s.x.a FROM '{file_path}'")) + .unwrap(); + + let mut y = Vec::new(); + let mut xb = Vec::new(); + let mut xa = Vec::new(); + for chunk in result { + let len = chunk.len().as_(); + y.extend_from_slice(chunk.get_vector(0).as_slice_with_len::(len)); + xb.extend_from_slice(chunk.get_vector(1).as_slice_with_len::(len)); + xa.extend_from_slice(chunk.get_vector(2).as_slice_with_len::(len)); + } + + assert_eq!(y, vec![7, 8, 9], "s.y mismatch"); + assert_eq!(xb, vec![4, 5, 6], "s.x.b mismatch"); + assert_eq!(xa, vec![1, 2, 3], "s.x.a mismatch"); +} + /// `SELECT array_length(list)` / `len(list)` / `length(list)` should push the list-length /// computation into the Vortex scan (computed from offsets, without materializing the list /// elements) and return the per-row element counts. diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index adecd2dbbdf..2ce289eb9bf 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -8,10 +8,12 @@ use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::error::VortexExpect; use vortex::error::VortexResult; +use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex::expr::BoundExpression; use vortex::expr::Expression; use vortex::expr::and_collect; +use vortex::expr::cast; use vortex::expr::col; use vortex::expr::get_item; use vortex::expr::lit; @@ -23,9 +25,12 @@ use vortex::scalar::Scalar; use vortex::scan::selection::Selection; use vortex_utils::aliases::hash_set::HashSet; +use crate::convert::FromLogicalType; use crate::convert::try_from_table_filter; use crate::convert::try_from_virtual_column_filter; use crate::duckdb::LogicalType; +use crate::duckdb::StructExtract; +use crate::duckdb::StructExtracts; use crate::duckdb::TableFilterClass; use crate::duckdb::TableFilterSetRef; use crate::table_function::ColumnAggregate; @@ -71,14 +76,14 @@ pub struct ProjectionInput<'a> { } impl Projection { - pub fn new(input: ProjectionInput) -> Self { + pub fn new(input: ProjectionInput, struct_extracts: &StructExtracts<'_>) -> VortexResult { let projection_ids: HashSet = input.projection_ids.iter().copied().collect(); // If projection ids are empty, use column_ids. // See duckdb/src/planner/operator/logical_get.cpp#L168 let is_projected = |pos: usize| projection_ids.is_empty() || projection_ids.contains(&(pos as u64)); - let mut exprs = Vec::with_capacity(input.column_ids.len() + 1); + let mut exprs: Vec<(String, Expression)> = Vec::with_capacity(input.column_ids.len() + 1); let mut file_row_number_column_pos = None; let mut is_star = true; let mut real_column_count = 0; @@ -93,7 +98,7 @@ impl Projection { // filter-only column needs to be emitted only for output // vector position match, it will never be read let dtype = DType::Primitive(PType::U64, Nullability::Nullable); - exprs.push(("file_row_number", lit(Scalar::null(dtype)))); + exprs.push(("file_row_number".to_owned(), lit(Scalar::null(dtype)))); } continue; } @@ -110,7 +115,10 @@ impl Projection { is_star = false; // filter-only column needs to be emitted only for output // vector position match, it will never be read - exprs.push((name, lit(Scalar::null(column_field.dtype.as_nullable())))); + exprs.push(( + name.to_owned(), + lit(Scalar::null(column_field.dtype.as_nullable())), + )); continue; } @@ -121,12 +129,21 @@ impl Projection { // Example: if we SELECT len(str), we can't use root() as we try to // pushdown scalar functions. - let expr = match &column_field.projection_expr { - None => get_item(name, root()), - Some(func) => { + let (name, expr) = match (&column_field.projection_expr, struct_extracts.get(column_pos)) { + (Some(_), Some(_)) => vortex_bail!( + "column '{name}' has both a scalar function and a struct extract pushed into the scan" + ), + (Some(func), None) => { is_star = false; - func.clone() + (name.to_owned(), func.clone()) + } + // DuckDB asked for a field of this struct column only, so read + // and emit that field instead of the whole struct. + (None, Some(extract)) => { + is_star = false; + struct_extract_expression(name, &extract, &column_field.dtype)? } + (None, None) => (name.to_owned(), get_item(name, root())), }; exprs.push((name, expr)); real_column_count += 1; @@ -136,19 +153,19 @@ impl Projection { is_star &= real_column_count == input.column_fields.len() as u64; if is_star { - return Projection { + return Ok(Projection { projection: root(), file_row_number_column_pos: None, - }; + }); } if file_row_number_column_pos.is_some() { // row_idx will be moved to correct position in scan(), prepend here - exprs.insert(0, ("file_row_number", row_idx())); + exprs.insert(0, ("file_row_number".to_owned(), row_idx())); } - Self { + Ok(Self { projection: pack(exprs, false.into()), file_row_number_column_pos, - } + }) } // Create a projection for aggregate scan @@ -187,6 +204,49 @@ impl Projection { } } +/// Builds the expression reading a struct field DuckDB pushed into the scan, and +/// the name to give its output column. +/// +/// The names of the output columns are not read back by DuckDB, but they have to +/// stay distinct: DuckDB binds one output column per extracted path, so two +/// extracts of the same column would otherwise pack two fields with one name. +fn struct_extract_expression( + column_name: &str, + extract: &StructExtract<'_>, + dtype: &DType, +) -> VortexResult<(String, Expression)> { + let mut expr = col(column_name); + let mut dtype = dtype.clone(); + let mut name = column_name.to_owned(); + + for index in extract.indexes { + let index = usize::try_from(*index) + .map_err(|_| vortex_err!("struct extract index {index} does not fit in usize"))?; + let fields = dtype + .as_struct_fields_opt() + .ok_or_else(|| vortex_err!("'{name}' is not a struct, cannot extract field {index}"))?; + let field_name = fields + .field_name(index) + .ok_or_else(|| vortex_err!("'{name}' has no field at index {index}"))? + .clone(); + dtype = fields + .field_by_index(index) + .ok_or_else(|| vortex_err!("'{name}' has no field at index {index}"))?; + + expr = get_item(field_name.clone(), expr); + name = format!("{name}.{field_name}"); + } + + // DuckDB passes the type it expects the scan to emit for this field, which is + // the cast target when the query casts the extracted field itself. + let expected = DType::from_logical_type(extract.datatype, Nullability::Nullable)?; + if expected != dtype.as_nullable() { + expr = cast(expr, expected); + } + + Ok((name, expr)) +} + pub struct Filter { pub filter: Option, pub row_selection: Selection, @@ -307,12 +367,12 @@ mod tests { projection_ids: &[], column_fields: &fields, }; - assert_eq!(Projection::new(input.clone()).projection, root()); + assert_eq!(Projection::new(input.clone(), &StructExtracts::default()).unwrap().projection, root()); // file_row_number turns star into an explicit pack with row_idx first let ids = [FILE_ROW_NUMBER_COLUMN_IDX, 0, 1, 2]; input.column_ids = &ids; - let result = Projection::new(input.clone()); + let result = Projection::new(input.clone(), &StructExtracts::default()).unwrap(); let expected = pack( [ ("file_row_number", row_idx()), @@ -326,13 +386,13 @@ mod tests { assert_eq!(result.file_row_number_column_pos, Some(0)); input.column_ids = &[0, 1]; - assert_ne!(Projection::new(input.clone()).projection, root()); + assert_ne!(Projection::new(input.clone(), &StructExtracts::default()).unwrap().projection, root()); input.column_ids = &[0, 2, 2]; - assert_ne!(Projection::new(input.clone()).projection, root()); + assert_ne!(Projection::new(input.clone(), &StructExtracts::default()).unwrap().projection, root()); input.column_ids = &[2, 1, 0]; - assert_ne!(Projection::new(input.clone()).projection, root()); + assert_ne!(Projection::new(input.clone(), &StructExtracts::default()).unwrap().projection, root()); // If any column has a projection expression, we can't use SELECT * let mut fields = [field("a"), field("b"), field("c")]; @@ -342,7 +402,7 @@ mod tests { projection_ids: &[], column_fields: &fields, }; - assert_ne!(Projection::new(input).projection, root()); + assert_ne!(Projection::new(input, &StructExtracts::default()).unwrap().projection, root()); } #[test] @@ -354,7 +414,7 @@ mod tests { projection_ids: &[0, 2], column_fields: &fields, }; - let projection = Projection::new(input).projection; + let projection = Projection::new(input, &StructExtracts::default()).unwrap().projection; let expected = pack( [ ("a", get_item("a", root())), @@ -370,7 +430,7 @@ mod tests { projection_ids: &[1], column_fields: &fields, }; - let result = Projection::new(input); + let result = Projection::new(input, &StructExtracts::default()).unwrap(); let frn_dtype = DType::Primitive(PType::U64, Nullability::Nullable); let expected = pack( [ @@ -387,7 +447,7 @@ mod tests { projection_ids: &[0, 1], column_fields: &fields, }; - let result = Projection::new(input); + let result = Projection::new(input, &StructExtracts::default()).unwrap(); let expected = pack( [ ("file_row_number", row_idx()), @@ -400,6 +460,83 @@ mod tests { assert_eq!(result.file_row_number_column_pos, Some(1)); } + #[test] + fn test_struct_extract() { + let fields = [DuckdbField { + name: "s".to_owned(), + logical_type: LogicalType::null(), + dtype: DType::struct_( + [("a", DType::Primitive(PType::I32, Nullability::NonNullable))], + Nullability::NonNullable, + ), + projection_expr: None, + }]; + + // DuckDB hands us the type it expects for the extracted field. + let expected_type = + LogicalType::try_from(&DType::Primitive(PType::I32, Nullability::Nullable)) + .vortex_expect("logical type"); + let datatypes = [expected_type.as_ptr()]; + let extracts = StructExtracts { + indexes: &[0], + offsets: &[0, 1], + datatypes: &datatypes, + }; + + let input = ProjectionInput { + column_ids: &[0], + projection_ids: &[], + column_fields: &fields, + }; + let result = Projection::new(input, &extracts).vortex_expect("projection"); + let expected = pack( + [("s.a", get_item("a", get_item("s", root())))], + false.into(), + ); + assert_eq!(result.projection, expected); + + // A cast target is emitted as a cast expression. + let cast_type = LogicalType::try_from(&DType::Primitive(PType::I64, Nullability::Nullable)) + .vortex_expect("logical type"); + let datatypes = [cast_type.as_ptr()]; + let extracts = StructExtracts { + indexes: &[0], + offsets: &[0, 1], + datatypes: &datatypes, + }; + let input = ProjectionInput { + column_ids: &[0], + projection_ids: &[], + column_fields: &fields, + }; + let result = Projection::new(input, &extracts).vortex_expect("projection"); + let expected = pack( + [( + "s.a", + cast( + get_item("a", get_item("s", root())), + DType::Primitive(PType::I64, Nullability::Nullable), + ), + )], + false.into(), + ); + assert_eq!(result.projection, expected); + + // Without an extract the whole column is read, and the extract still has + // to be an error when the path does not exist. + let input = ProjectionInput { + column_ids: &[0], + projection_ids: &[], + column_fields: &fields, + }; + let bad = StructExtracts { + indexes: &[1], + offsets: &[0, 1], + datatypes: &datatypes, + }; + assert!(Projection::new(input, &bad).is_err()); + } + #[test] fn test_push_filter_expr_preserves_order() { let first = col("first"); diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index ced5fb27e59..92d666135ae 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -262,7 +262,7 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult { projection_ids, column_fields: &bind_data.columns, }; - Projection::new(input) + Projection::new(input, &init_input.struct_extracts())? } else { Projection::new_aggregate(&bind_data.aggregates, &bind_data.columns) };