[WIP] feat(duckdb): read pushed-down struct extracts

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.

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<idx_t> extract_indexes;
+    vector<size_t> extract_offsets;
+    vector<duckdb_logical_type> 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<duckdb_logical_type>(
+                const_cast<LogicalType *>(&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<duckdb_vx_table_filter_set>(input.filters.get()),
         .client_context = reinterpret_cast<duckdb_client_context>(&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;

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<VortexReaderInterface>::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<BaseStatistics>();
+        }
+        return MultiFileFunction<VortexReaderInterface>::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<MultiFileBindData>();
+        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;
 

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<StructExtract<'a>> {
+        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,
+                )
+            },
+        }
+    }
 }

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::<i32>(len));
+        xb.extend_from_slice(chunk.get_vector(1).as_slice_with_len::<i32>(len));
+        xa.extend_from_slice(chunk.get_vector(2).as_slice_with_len::<i32>(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.