initial

Author: myrrc

vortex-duckdb/cpp/include/duckdb_vx/table_function.h

@@ -83,6 +83,11 @@ typedef struct {
     bool has_null;
 } duckdb_column_statistics;
 
+typedef struct {
+    idx_t batch_index;
+    size_t file_index;
+} duckdb_vx_partition_data;
+
 // vtable mimicking subset of TableFunction.
 // See duckdb/include/function/tfunc.hpp
 typedef struct {
@@ -123,10 +128,10 @@ typedef struct {
 
     double (*table_scan_progress)(duckdb_client_context ctx, void *bind_data, void *global_state);
 
-    idx_t (*get_partition_data)(const void *bind_data,
-                                void *init_global_data,
-                                void *init_local_data,
-                                duckdb_vx_error *error_out);
+    void (*get_partition_data)(const void *bind_data,
+                               void *init_global_data,
+                               void *init_local_data,
+                               duckdb_vx_partition_data *partition_data_out);
 } duckdb_vx_tfunc_vtab_t;
 
 // A single function for configuring the DuckDB table function vtable.

vortex-duckdb/cpp/table_function.cpp

@@ -10,6 +10,7 @@ DUCKDB_INCLUDES_BEGIN
 #include "duckdb.h"
 #include "duckdb/catalog/catalog.hpp"
 #include "duckdb/common/insertion_order_preserving_map.hpp"
+#include "duckdb/common/multi_file/multi_file_reader.hpp"
 #include "duckdb/function/table_function.hpp"
 #include "duckdb/main/capi/capi_internal.hpp"
 #include "duckdb/main/connection.hpp"
@@ -19,6 +20,8 @@ DUCKDB_INCLUDES_END
 using namespace duckdb;
 using vortex::CData;
 using vortex::IntoErrString;
+constexpr column_t COLUMN_IDENTIFIER_FILE_INDEX = MultiFileReader::COLUMN_IDENTIFIER_FILE_INDEX;
+// constexpr column_t COLUMN_IDENTIFIER_FILE_ROW_NUMBER = MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER;
 
 struct CTableFunctionInfo final : TableFunctionInfo {
     explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {
@@ -343,12 +346,12 @@ OperatorPartitionData c_get_partition_data(ClientContext &, TableFunctionGetPart
     void *const ffi_global = input.global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
     void *const ffi_local = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
 
-    duckdb_vx_error error_out = nullptr;
-    const idx_t batch_index = bind.info.vtab.get_partition_data(ffi_bind, ffi_global, ffi_local, &error_out);
-    if (error_out) {
-        throw InvalidInputException(IntoErrString(error_out));
-    }
-    return OperatorPartitionData(batch_index);
+    duckdb_vx_partition_data partition_data;
+    bind.info.vtab.get_partition_data(ffi_bind, ffi_global, ffi_local, &partition_data);
+
+    OperatorPartitionData out(partition_data.batch_index);
+    out.partition_data = {ColumnPartitionData(Value::UBIGINT(partition_data.file_index))};
+    return out;
 }
 
 extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, const char *value) {
@@ -380,8 +383,8 @@ extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const d
     // We can prune out filter columns that are unused in the remainder of the query plan.
     // e.g. in "SELECT i FROM tbl WHERE j = 42" j does not leave Vortex table function.
     tf.filter_prune = true;
+    tf.late_materialization = true;
     tf.sampling_pushdown = false;
-    tf.late_materialization = false;
 
     tf.pushdown_complex_filter = c_pushdown_complex_filter;
     tf.cardinality = c_cardinality;
@@ -390,8 +393,19 @@ extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const d
     tf.table_scan_progress = c_table_scan_progress;
     tf.statistics = c_statistics;
 
+    tf.get_row_id_columns = [](auto &, auto) -> vector<column_t> {
+        return {
+            COLUMN_IDENTIFIER_FILE_INDEX,
+            // COLUMN_IDENTIFIER_FILE_ROW_NUMBER
+        };
+    };
+
     tf.get_virtual_columns = [](auto &, auto) -> virtual_column_map_t {
-        return {{COLUMN_IDENTIFIER_EMPTY, TableColumn("", LogicalTypeId::BOOLEAN)}};
+        return {
+            {COLUMN_IDENTIFIER_EMPTY, {"", LogicalTypeId::BOOLEAN}},
+            {COLUMN_IDENTIFIER_FILE_INDEX, {"file_index", LogicalType::UBIGINT}},
+            //{COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::BIGINT}}
+        };
     };
 
     tf.arguments.resize(vtab->parameter_count);

vortex-duckdb/src/datasource.rs

@@ -67,25 +67,21 @@ use crate::duckdb::DataChunkRef;
 use crate::duckdb::DuckdbStringMapRef;
 use crate::duckdb::ExpressionRef;
 use crate::duckdb::LogicalType;
+use crate::duckdb::PartitionData;
 use crate::duckdb::TableFilterSetRef;
 use crate::duckdb::TableFunction;
 use crate::duckdb::TableInitInput;
+use crate::duckdb::Value;
 use crate::exporter::ArrayExporter;
 use crate::exporter::ConversionCache;
 
-/// Taken from
-/// https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/include/duckdb/common/constants.hpp#L44
-///
-/// If DuckDB requests a zero-column projection from read_vortex like count(*),
-/// its planner tries to get any column:
-/// https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/planner/operator/logical_get.cpp#L149
-///
-/// If you define COLUMN_IDENTIFIER_EMPTY, planner takes it, otherwise the
-/// first column. As we don't want to fill the output chunk and we can leave
-/// it uninitialized in this case, we define COLUMN_IDENTIFIER_EMPTY as a
-/// virtual column.
-/// See virtual_columns in vortex-duckdb/cpp/table_function.cpp
-static EMPTY_COLUMN_IDX: u64 = 18446744073709551614;
+/// File index virtual column, may be requested either by user or optimizer.
+static FILE_INDEX_COLUMN_IDX: u64 = 9223372036854775810;
+
+/// See duckdb/src/common/constants.cpp
+fn is_virtual_column(id: u64) -> bool {
+    return id >= 9223372036854775808u64;
+}
 
 /// A trait for table functions that resolve to a [`DataSourceRef`].
 ///
@@ -141,22 +137,24 @@ impl Debug for DataSourceBindData {
     }
 }
 
-type DataSourceIterator = ThreadSafeIterator<VortexResult<(ArrayRef, Arc<ConversionCache>)>>;
+type DataSourceIterator = ThreadSafeIterator<VortexResult<(ArrayRef, Arc<ConversionCache>, usize)>>;
 
 /// Global scan state for driving a `DataSource` scan through DuckDB.
 pub struct DataSourceGlobal {
     iterator: DataSourceIterator,
     batch_id: AtomicU64,
     bytes_total: Arc<AtomicU64>,
     bytes_read: AtomicU64,
+    /// "file_index" column position if requested
+    file_idx_pos: Option<usize>,
 }
 
 /// Per-thread local scan state.
 pub struct DataSourceLocal {
     iterator: DataSourceIterator,
     exporter: Option<ArrayExporter>,
-    /// The unique batch id of the last chunk exported via scan().
-    batch_id: Option<u64>,
+    batch_id: u64,
+    file_idx: usize,
 }
 
 /// Returns scan progress as a percentage (0.0–100.0).
@@ -281,7 +279,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         let column_ids = init_input.column_ids();
         let projection_ids = init_input.projection_ids();
 
-        let projection_expr =
+        let (projection_expr, file_idx_pos) =
             extract_projection_expr(projection_ids, column_ids, &bind_data.column_fields);
         let filter_expr = extract_table_filter_expr(
             init_input.table_filter_set(),
@@ -317,7 +315,8 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         // first available array chunk.
         let stream = scan
             .partitions()
-            .map(move |partition| {
+            .enumerate()
+            .map(move |(file_idx, partition)| {
                 // We create a new conversion cache scoped to the partition, since there's no point
                 // caching anything across partitions.
                 let cache = Arc::new(ConversionCache::default());
@@ -333,7 +332,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
                     };
                     while let Some(item) = stream.next().await {
                         if tx
-                            .send(item.map(|a| (a, Arc::clone(&cache))))
+                            .send(item.map(|a| (a, Arc::clone(&cache), file_idx)))
                             .await
                             .is_err()
                         {
@@ -356,6 +355,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
             batch_id: AtomicU64::new(0),
             bytes_total: Arc::new(AtomicU64::new(0)),
             bytes_read: AtomicU64::new(0),
+            file_idx_pos,
         })
     }
 
@@ -381,7 +381,8 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         Ok(DataSourceLocal {
             iterator: global.iterator.clone(),
             exporter: None,
-            batch_id: None,
+            batch_id: 0,
+            file_idx: 0,
         })
     }
 
@@ -398,7 +399,8 @@ impl<T: DataSourceTableFunction> TableFunction for T {
                 let Some(result) = local_state.iterator.next() else {
                     return Ok(());
                 };
-                let (array_result, conversion_cache) = result?;
+                let (array_result, conversion_cache, file_idx) = result?;
+                local_state.file_idx = file_idx;
                 let array_result = array_result.optimize_recursive()?;
 
                 let array_result: StructArray = if let Some(array) = array_result.as_opt::<Struct>()
@@ -423,30 +425,36 @@ impl<T: DataSourceTableFunction> TableFunction for T {
                     ctx,
                 )?);
                 // Relaxed since there is no intra-instruction ordering required.
-                local_state.batch_id = Some(global_state.batch_id.fetch_add(1, Ordering::Relaxed));
+                local_state.batch_id = global_state.batch_id.fetch_add(1, Ordering::Relaxed);
             }
 
             let exporter = local_state
                 .exporter
                 .as_mut()
                 .vortex_expect("error: exporter missing");
+            let has_more_data = exporter.export(chunk, global_state.file_idx_pos)?;
 
-            let has_more_data = exporter.export(chunk)?;
             global_state
                 .bytes_read
                 .fetch_add(chunk.len(), Ordering::Relaxed);
 
             if !has_more_data {
                 // This exporter is fully consumed.
                 local_state.exporter = None;
-                local_state.batch_id = None;
+                local_state.batch_id = 0;
             } else {
                 break;
             }
         }
 
         assert!(!chunk.is_empty());
 
+        if let Some(pos) = global_state.file_idx_pos {
+            chunk
+                .get_vector_mut(pos)
+                .reference_value(&Value::new_u64(local_state.file_idx as u64));
+        }
+
         Ok(())
     }
 
@@ -518,10 +526,11 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         _bind_data: &Self::BindData,
         _global_init_data: &Self::GlobalState,
         local_init_data: &mut Self::LocalState,
-    ) -> VortexResult<u64> {
-        local_init_data
-            .batch_id
-            .ok_or_else(|| vortex_err!("batch id missing, no batches exported"))
+    ) -> PartitionData {
+        PartitionData {
+            batch_index: local_init_data.batch_id,
+            file_index: local_init_data.file_idx,
+        }
     }
 
     fn to_string(bind_data: &Self::BindData, map: &mut DuckdbStringMapRef) {
@@ -557,39 +566,46 @@ fn extract_schema_from_dtype(dtype: &DType) -> VortexResult<Vec<DuckdbField>> {
     Ok(fields)
 }
 
-/// Creates a projection expression from raw projection/column ID slices and column names.
+/// Creates a projection expression from raw projection/column ID slices and
+/// column names.
+/// If FILE_INDEX_COLUMN_IDX is present, returns its position as second
+/// parameter
 fn extract_projection_expr(
     projection_ids: Option<&[u64]>,
     column_ids: &[u64],
     column_fields: &[DuckdbField],
-) -> Expression {
-    // Projection ids may be empty, in which case you need to use projection_ids
+) -> (Expression, Option<usize>) {
+    // If projection ids are empty, use column_ids.
     // See duckdb/src/planner/operator/logical_get.cpp#L168
-    let (projection_ids, has_projection_ids) = match projection_ids {
+    let (ids, has_projection_ids) = match projection_ids {
         Some(ids) => (ids, true),
         None => (column_ids, false),
     };
 
-    // duckdb index is u64 (size_t) but in Rust u64 and usize are different things.
+    let mut file_idx_pos = None;
+
     #[expect(clippy::cast_possible_truncation)]
-    let names = projection_ids
+    let names = ids
         .iter()
-        .filter(|p| **p != EMPTY_COLUMN_IDX)
-        .map(|mut idx| {
-            if has_projection_ids {
-                idx = &column_ids[*idx as usize];
+        .enumerate()
+        .map(|(column_pos, &column_id)| {
+            let column_id = if has_projection_ids {
+                column_ids[column_id as usize]
+            } else {
+                column_id
+            };
+
+            if column_id == FILE_INDEX_COLUMN_IDX {
+                file_idx_pos = Some(column_pos);
             }
 
-            #[expect(clippy::cast_possible_truncation)]
-            &column_fields
-                .get(*idx as usize)
-                .vortex_expect("prune idx in column names")
-                .name
+            column_id
         })
-        .map(|s| Arc::from(s.as_str()))
+        .filter(|&col_id| !is_virtual_column(col_id))
+        .map(|col_id| Arc::from(column_fields[col_id as usize].name.as_str()))
         .collect::<FieldNames>();
 
-    select(names, root())
+    (select(names, root()), file_idx_pos)
 }
 
 /// Creates a table filter expression from the table filter set, column metadata, additional
@@ -604,6 +620,10 @@ fn extract_table_filter_expr(
     let mut table_filter_exprs: HashSet<Expression> = if let Some(filter) = table_filter_set {
         filter
             .into_iter()
+            .filter(|(idx, _)| {
+                let idx_u: usize = idx.as_();
+                !is_virtual_column(column_ids[idx_u].as_())
+            })
             .map(|(idx, ex)| {
                 let idx_u: usize = idx.as_();
                 let col_idx: usize = column_ids[idx_u].as_();

vortex-duckdb/src/duckdb/table_function/mod.rs

@@ -36,6 +36,11 @@ use crate::duckdb::table_function::statistics::statistics;
 use crate::duckdb::table_function::table_scan_progress::table_scan_progress_callback;
 use crate::duckdb_try;
 
+pub struct PartitionData {
+    pub batch_index: u64,
+    pub file_index: usize,
+}
+
 #[derive(Debug, Default)]
 pub struct ColumnStatistics {
     pub min: Option<Value>,
@@ -137,10 +142,10 @@ pub trait TableFunction: Sized + Debug {
     /// Returns the idx of the current partition being processed by a local threa.
     /// This *must* be globally unique.
     fn partition_data(
-        _bind_data: &Self::BindData,
-        _global_init_data: &Self::GlobalState,
-        _local_init_data: &mut Self::LocalState,
-    ) -> VortexResult<u64>;
+        bind_data: &Self::BindData,
+        global_init_data: &Self::GlobalState,
+        local_init_data: &mut Self::LocalState,
+    ) -> PartitionData;
 
     /// Returns a vector of key-value pairs for EXPLAIN output
     fn to_string(bind_data: &Self::BindData, map: &mut DuckdbStringMapRef);