re-add column pos

Author: myrrc

vortex-duckdb/cpp/table_function.cpp

@@ -21,6 +21,7 @@ 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) {
@@ -391,13 +392,14 @@ extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const d
     tf.statistics = c_statistics;
 
     tf.get_row_id_columns = [](auto &, auto) -> vector<column_t> {
-        return {COLUMN_IDENTIFIER_FILE_INDEX};
+        return {COLUMN_IDENTIFIER_FILE_INDEX, COLUMN_IDENTIFIER_FILE_ROW_NUMBER};
     };
 
     tf.get_virtual_columns = [](auto &, auto) -> virtual_column_map_t {
         return {
             {COLUMN_IDENTIFIER_EMPTY, {"", LogicalTypeId::BOOLEAN}},
             {COLUMN_IDENTIFIER_FILE_INDEX, {"file_index", LogicalType::UBIGINT}},
+            {COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::BIGINT}},
         };
     };
 

vortex-duckdb/src/datasource.rs

@@ -75,8 +75,10 @@ use crate::duckdb::Value;
 use crate::exporter::ArrayExporter;
 use crate::exporter::ConversionCache;
 
-/// File index virtual column, may be requested either by user or optimizer.
+/// "file_index" virtual column, may be requested either by user or optimizer.
 static FILE_INDEX_COLUMN_IDX: u64 = 9223372036854775810;
+/// "file_row_number" virtual column, may be requested either by user or optimizer.
+static FILE_ROW_NUMBER_COLUMN_IDX: u64 = 9223372036854775809;
 
 /// See duckdb/src/common/constants.cpp
 fn is_virtual_column(id: u64) -> bool {
@@ -146,6 +148,7 @@ pub struct DataSourceGlobal {
     bytes_total: Arc<AtomicU64>,
     bytes_read: AtomicU64,
     file_index_column_pos: Option<usize>,
+    file_row_number_column_pos: Option<usize>,
 }
 
 /// Per-thread local scan state.
@@ -278,8 +281,11 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         let column_ids = init_input.column_ids();
         let projection_ids = init_input.projection_ids();
 
-        let (projection_expr, file_idx_pos) =
-            extract_projection_expr(projection_ids, column_ids, &bind_data.column_fields);
+        let ProjectionWithVirtualColumns {
+            projection,
+            file_index_column_pos,
+            file_row_number_column_pos,
+        } = extract_projection_expr(projection_ids, column_ids, &bind_data.column_fields);
         let filter_expr = extract_table_filter_expr(
             init_input.table_filter_set(),
             column_ids,
@@ -291,10 +297,10 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         let filter_expr_str = filter_expr
             .as_ref()
             .map_or_else(|| "true".to_string(), |f| f.to_string());
-        debug!("Global init Vortex scan SELECT {projection_expr} WHERE {filter_expr_str}");
+        debug!("Global init Vortex scan SELECT {projection} WHERE {filter_expr_str}");
 
         let request = ScanRequest {
-            projection: projection_expr,
+            projection,
             filter: filter_expr,
             ordered: false,
             ..Default::default()
@@ -357,7 +363,8 @@ impl<T: DataSourceTableFunction> TableFunction for T {
             batch_id: AtomicU64::new(0),
             bytes_total: Arc::new(AtomicU64::new(0)),
             bytes_read: AtomicU64::new(0),
-            file_index_column_pos: file_idx_pos,
+            file_index_column_pos,
+            file_row_number_column_pos,
         })
     }
 
@@ -568,6 +575,12 @@ fn extract_schema_from_dtype(dtype: &DType) -> VortexResult<Vec<DuckdbField>> {
     Ok(fields)
 }
 
+struct ProjectionWithVirtualColumns {
+    projection: Expression,
+    file_index_column_pos: Option<usize>,
+    file_row_number_column_pos: Option<usize>,
+}
+
 /// 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
@@ -576,15 +589,16 @@ fn extract_projection_expr(
     projection_ids: Option<&[u64]>,
     column_ids: &[u64],
     column_fields: &[DuckdbField],
-) -> (Expression, Option<usize>) {
+) -> ProjectionWithVirtualColumns {
     // If projection ids are empty, use column_ids.
     // See duckdb/src/planner/operator/logical_get.cpp#L168
     let (ids, has_projection_ids) = match projection_ids {
         Some(ids) => (ids, true),
         None => (column_ids, false),
     };
 
-    let mut file_idx_pos = None;
+    let mut file_index_column_pos = None;
+    let mut file_row_number_column_pos = None;
 
     #[expect(clippy::cast_possible_truncation)]
     let names = ids
@@ -598,7 +612,10 @@ fn extract_projection_expr(
             };
 
             if column_id == FILE_INDEX_COLUMN_IDX {
-                file_idx_pos = Some(column_pos);
+                file_index_column_pos = Some(column_pos);
+            }
+            if column_id == FILE_ROW_NUMBER_COLUMN_IDX {
+                file_row_number_column_pos = Some(column_pos);
             }
 
             column_id
@@ -607,7 +624,11 @@ fn extract_projection_expr(
         .map(|col_id| Arc::from(column_fields[col_id as usize].name.as_str()))
         .collect::<FieldNames>();
 
-    (select(names, root()), file_idx_pos)
+    ProjectionWithVirtualColumns {
+        projection: select(names, root()),
+        file_index_column_pos,
+        file_row_number_column_pos,
+    }
 }
 
 /// Creates a table filter expression from the table filter set, column metadata, additional