better

Author: myrrc

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

@@ -85,7 +85,10 @@ typedef struct {
 
 typedef struct {
     idx_t batch_index;
-    size_t file_index;
+    size_t file_index_column_pos;
+    size_t file_row_number_column_pos;
+    bool has_file_index_column_pos;
+    bool has_file_row_number_column_pos;
 } duckdb_vx_partition_data;
 
 // vtable mimicking subset of TableFunction.

vortex-duckdb/cpp/table_function.cpp

@@ -338,9 +338,6 @@ extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_resu
 }
 
 OperatorPartitionData c_get_partition_data(ClientContext &, TableFunctionGetPartitionInput &input) {
-    if (input.partition_info.RequiresPartitionColumns()) {
-        throw InternalException("TableScan::GetPartitionData: partition columns not supported");
-    }
     auto &bind = input.bind_data->Cast<CTableBindData>();
     void *const ffi_bind = bind.ffi_data->DataPtr();
     void *const ffi_global = input.global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
@@ -350,7 +347,16 @@ OperatorPartitionData c_get_partition_data(ClientContext &, TableFunctionGetPart
     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))};
+    for (idx_t col : input.partition_info.partition_columns) {
+        if (partition_data.has_file_index_column_pos && col == partition_data.file_index_column_pos) {
+            out.partition_data.emplace_back(Value::UBIGINT(partition_data.file_index_column_pos));
+        } else if (partition_data.has_file_row_number_column_pos &&
+                   col == partition_data.file_row_number_column_pos) {
+            out.partition_data.emplace_back(Value::BIGINT(partition_data.file_row_number_column_pos));
+        } else {
+            throw InternalException("get_partition_data - did not find constant for the given partition");
+        }
+    }
     return out;
 }
 
@@ -391,6 +397,8 @@ 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;
 
+    // What columns uniquely identify a row for deferred re-fetch in a
+    // multi-file scan: (file index, row index in file).
     tf.get_row_id_columns = [](auto &, auto) -> vector<column_t> {
         return {COLUMN_IDENTIFIER_FILE_INDEX, COLUMN_IDENTIFIER_FILE_ROW_NUMBER};
     };

vortex-duckdb/src/datasource.rs

@@ -28,12 +28,16 @@ use vortex::array::optimizer::ArrayOptimizer;
 use vortex::array::stats::StatsSet;
 use vortex::dtype::DType;
 use vortex::dtype::FieldNames;
+use vortex::dtype::PType;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_err;
 use vortex::expr::Expression;
 use vortex::expr::and_collect;
+use vortex::expr::cast;
 use vortex::expr::col;
+use vortex::expr::merge;
+use vortex::expr::pack;
 use vortex::expr::root;
 use vortex::expr::select;
 use vortex::expr::stats::Precision;
@@ -42,6 +46,7 @@ use vortex::file::v2::FileStatsLayoutReader;
 use vortex::io::kanal_ext::KanalExt;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::ThreadSafeIterator;
+use vortex::layout::layouts::row_idx::row_idx;
 use vortex::layout::scan::multi::MultiLayoutChild;
 use vortex::layout::scan::multi::MultiLayoutDataSource;
 use vortex::metrics::tracing::get_global_labels;
@@ -75,9 +80,11 @@ use crate::duckdb::Value;
 use crate::exporter::ArrayExporter;
 use crate::exporter::ConversionCache;
 
-/// "file_index" virtual column, may be requested either by user or optimizer.
+// See MultiFileReader for constants
+
+/// "file_index" virtual column
 static FILE_INDEX_COLUMN_IDX: u64 = 9223372036854775810;
-/// "file_row_number" virtual column, may be requested either by user or optimizer.
+/// "file_row_number" virtual column
 static FILE_ROW_NUMBER_COLUMN_IDX: u64 = 9223372036854775809;
 
 /// See duckdb/src/common/constants.cpp
@@ -156,7 +163,7 @@ pub struct DataSourceLocal {
     iterator: DataSourceIterator,
     exporter: Option<ArrayExporter>,
     batch_id: u64,
-    file_idx: usize,
+    file_index: usize,
 }
 
 /// Returns scan progress as a percentage (0.0–100.0).
@@ -302,7 +309,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         let request = ScanRequest {
             projection,
             filter: filter_expr,
-            ordered: false,
+            ordered: file_row_number_column_pos.is_some(),
             ..Default::default()
         };
 
@@ -321,11 +328,11 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         let stream = scan
             .partitions()
             .enumerate()
-            .map(move |(file_idx, partition)| {
+            .map(move |(file_index, 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 {
-                    file_idx,
+                    file_index,
                     ..Default::default()
                 });
                 let tx = tx.clone();
@@ -391,7 +398,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
             iterator: global.iterator.clone(),
             exporter: None,
             batch_id: 0,
-            file_idx: 0,
+            file_index: 0,
         })
     }
 
@@ -409,8 +416,8 @@ impl<T: DataSourceTableFunction> TableFunction for T {
                     return Ok(());
                 };
                 let (array_result, conversion_cache) = result?;
-                local_state.file_idx = conversion_cache.file_idx;
                 let array_result = array_result.optimize_recursive(ctx.session())?;
+                local_state.file_index = conversion_cache.file_index;
 
                 let array_result: StructArray = if let Some(array) = array_result.as_opt::<Struct>()
                 {
@@ -441,7 +448,11 @@ impl<T: DataSourceTableFunction> TableFunction for T {
                 .exporter
                 .as_mut()
                 .vortex_expect("error: exporter missing");
-            let has_more_data = exporter.export(chunk, global_state.file_index_column_pos)?;
+            let has_more_data = exporter.export(
+                chunk,
+                global_state.file_index_column_pos,
+                global_state.file_row_number_column_pos,
+            )?;
 
             global_state
                 .bytes_read
@@ -461,7 +472,7 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         if let Some(pos) = global_state.file_index_column_pos {
             chunk
                 .get_vector_mut(pos)
-                .reference_value(&Value::from(local_state.file_idx as u64));
+                .reference_value(&Value::from(local_state.file_index as u64));
         }
 
         Ok(())
@@ -533,12 +544,13 @@ impl<T: DataSourceTableFunction> TableFunction for T {
 
     fn partition_data(
         _bind_data: &Self::BindData,
-        _global_init_data: &Self::GlobalState,
+        global_init_data: &Self::GlobalState,
         local_init_data: &mut Self::LocalState,
     ) -> PartitionData {
         PartitionData {
             batch_index: local_init_data.batch_id,
-            file_index: local_init_data.file_idx,
+            file_index_column_pos: global_init_data.file_index_column_pos,
+            file_row_number_column_pos: global_init_data.file_row_number_column_pos,
         }
     }
 
@@ -583,8 +595,6 @@ struct ProjectionWithVirtualColumns {
 
 /// 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],
@@ -624,8 +634,21 @@ fn extract_projection_expr(
         .map(|col_id| Arc::from(column_fields[col_id as usize].name.as_str()))
         .collect::<FieldNames>();
 
+    // file_index column will be filled later when exporting the chunk.
+
+    let select = select(names, root());
+    let projection = if file_row_number_column_pos.is_some() {
+        // Append row_idx to end - it will be rearranged to correct position
+        // in scan()
+        let row_idx = cast(row_idx(), DType::Primitive(PType::I64, false.into()));
+        let row_idx_struct = pack([("file_row_number", row_idx)], false.into());
+        merge([select, row_idx_struct])
+    } else {
+        select
+    };
+
     ProjectionWithVirtualColumns {
-        projection: select(names, root()),
+        projection,
         file_index_column_pos,
         file_row_number_column_pos,
     }

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

@@ -38,7 +38,8 @@ use crate::duckdb_try;
 
 pub struct PartitionData {
     pub batch_index: u64,
-    pub file_index: usize,
+    pub file_index_column_pos: Option<usize>,
+    pub file_row_number_column_pos: Option<usize>,
 }
 
 #[derive(Debug, Default)]

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

@@ -24,5 +24,18 @@ pub(crate) unsafe extern "C-unwind" fn get_partition_data_callback<T: TableFunct
     let data = T::partition_data(bind_data, global_init_data, local_init_data);
     let out = unsafe { &mut *partition_data_out };
     out.batch_index = data.batch_index;
-    out.file_index = data.file_index;
+
+    if let Some(pos) = data.file_index_column_pos {
+        out.has_file_index_column_pos = true;
+        out.file_index_column_pos = pos;
+    } else {
+        out.has_file_index_column_pos = false;
+    }
+
+    if let Some(pos) = data.file_row_number_column_pos {
+        out.has_file_row_number_column_pos = true;
+        out.file_row_number_column_pos = pos;
+    } else {
+        out.has_file_row_number_column_pos = false;
+    }
 }

vortex-duckdb/src/e2e_test/object_cache_test.rs

@@ -114,7 +114,7 @@ impl TableFunction for TestTableFunction {
     ) -> PartitionData {
         PartitionData {
             batch_index: 0,
-            file_index: 0,
+            file_index_column_pos: 0,
         }
     }
 

vortex-duckdb/src/exporter/cache.rs

@@ -21,5 +21,5 @@ pub struct ConversionCache {
     pub dict_cache: DashMap<usize, (ArrayRef, ReusableDict)>,
     pub values_cache: DashMap<usize, (ArrayRef, Arc<Mutex<Vector>>)>,
     pub canonical_cache: DashMap<usize, (ArrayRef, Canonical)>,
-    pub file_idx: usize,
+    pub file_index: usize,
 }

vortex-duckdb/src/exporter/mod.rs

@@ -78,15 +78,19 @@ impl ArrayExporter {
     pub fn export(
         &mut self,
         chunk: &mut DataChunkRef,
-        file_idx_pos: Option<usize>,
+        file_index_column_pos: Option<usize>,
+        file_row_number_column_pos: Option<usize>,
     ) -> VortexResult<bool> {
         chunk.reset();
         if self.remaining == 0 {
             return Ok(false);
         }
 
         let zero_projection = self.fields.is_empty();
-        let expected_cols = self.fields.len() + file_idx_pos.is_some() as usize;
+
+        // file_row_number column is already populated in scan construction
+        let expected_cols = self.fields.len() + file_index_column_pos.is_some() as usize;
+
         let chunk_cols = chunk.column_count();
         if !zero_projection && chunk_cols != expected_cols {
             vortex_bail!("Expected {expected_cols} columns in output chunk, got {chunk_cols}");
@@ -110,13 +114,27 @@ impl ArrayExporter {
             return Ok(true);
         }
 
-        let mut fields = self.fields.iter_mut();
+        let mut fields = self.fields.iter();
         for i in 0..chunk_cols {
-            if let Some(file_idx_pos) = file_idx_pos
-                && i == file_idx_pos
+            // file_index column: skip this index - it will be filled after
+            // chunk export.
+            if let Some(pos) = file_index_column_pos
+                && i == pos
             {
                 continue;
             }
+
+            // file_row_number column: populate the column at index with
+            // contents of last column - row index column is appended at end in
+            // scan construction
+            if let Some(pos) = file_row_number_column_pos
+                && i == pos
+            {
+                let field = &self.fields[chunk_cols - 1];
+                field.export(position, chunk_len, chunk.get_vector_mut(i), &mut self.ctx)?;
+                continue;
+            }
+
             let field = fields.next().vortex_expect("field count mismatch");
             field.export(position, chunk_len, chunk.get_vector_mut(i), &mut self.ctx)?;
         }

vortex-file/src/multi/mod.rs

@@ -16,6 +16,7 @@ use vortex_error::vortex_bail;
 use vortex_io::filesystem::FileListing;
 use vortex_io::filesystem::FileSystemRef;
 use vortex_layout::LayoutReaderRef;
+use vortex_layout::layouts::row_idx::RowIdxLayoutReader;
 use vortex_layout::scan::multi::LayoutReaderFactory;
 use vortex_layout::scan::multi::MultiLayoutDataSource;
 use vortex_scan::DataSource;
@@ -142,7 +143,8 @@ impl MultiFileDataSource {
         let (first_file_listing, first_fs) = &all_files[0];
         let open_fn = self.open_options_fn.as_ref();
         let first_file = open_file(first_fs, first_file_listing, &self.session, open_fn).await?;
-        let first_reader = layout_reader_with_stats(&first_file)?;
+        let first_reader =
+            layout_reader_with_stats_and_row_idxs(&first_file, self.session.clone())?;
 
         let factories: Vec<Arc<dyn LayoutReaderFactory>> = all_files[1..]
             .iter()
@@ -228,8 +230,14 @@ async fn open_file(
 
 /// Creates a layout reader from a VortexFile, wrapping with `FileStatsLayoutReader` when
 /// file-level statistics are available.
-fn layout_reader_with_stats(file: &crate::VortexFile) -> VortexResult<LayoutReaderRef> {
+/// Also wraps with RowIdxLayoutReader
+fn layout_reader_with_stats_and_row_idxs(
+    file: &crate::VortexFile,
+    session: VortexSession,
+) -> VortexResult<LayoutReaderRef> {
     let mut reader = file.layout_reader()?;
+    reader = Arc::new(RowIdxLayoutReader::new(0, reader, session));
+
     if let Some(stats) = file.file_stats().cloned() {
         reader = Arc::new(FileStatsLayoutReader::new(
             reader,
@@ -258,6 +266,9 @@ impl LayoutReaderFactory for VortexFileReaderFactory {
             self.open_options_fn.as_ref(),
         )
         .await?;
-        Ok(Some(layout_reader_with_stats(&file)?))
+        Ok(Some(layout_reader_with_stats_and_row_idxs(
+            &file,
+            self.session.clone(),
+        )?))
     }
 }