duckdb: factor out fns from init_global (#3557)

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

vortex-duckdb-ext/src/scan.rs

@@ -51,7 +51,7 @@ pub struct VortexGlobalData {
     file_paths: SegQueue<PathBuf>,
     _is_first_file_processed: std::sync::atomic::AtomicBool,
     filter_expr: ExprRef,
-    project_expr: ExprRef,
+    projection_expr: ExprRef,
 }
 
 pub struct VortexLocalData {
@@ -86,6 +86,56 @@ fn extract_schema_from_vortex_file(
     Ok((column_names, column_types))
 }
 
+/// Creates a projection expression based on the table initialization input.
+fn create_projection_expr(init: &TableInitInput<VortexTableFunction>) -> ExprRef {
+    let projection_ids = init.projection_ids().unwrap_or(&[]);
+    let column_ids = init.column_ids();
+
+    let projected_ids = projection_ids.iter().map(|p| column_ids[p.as_usize()]);
+    select(
+        projected_ids
+            .map(|idx| {
+                init.bind_data()
+                    .column_names
+                    .get(idx.as_usize())
+                    .vortex_expect("prune idx in column names")
+            })
+            .map(|s| Arc::from(s.clone()))
+            .collect::<FieldNames>(),
+        root(),
+    )
+}
+
+/// Creates a table filter expression from the table filter set.
+fn create_table_filter_expr(
+    init: &TableInitInput<VortexTableFunction>,
+) -> VortexResult<Option<ExprRef>> {
+    init.table_filter_set()
+        .and_then(|filter| {
+            filter
+                .into_iter()
+                .map(|(idx, ex)| {
+                    let name = init
+                        .bind_data()
+                        .column_names
+                        .get(idx.as_usize())
+                        .vortex_expect("exists");
+                    try_from_table_filter(&ex, name)
+                })
+                .reduce(|l, r| Ok(and(l?, r?)))
+        })
+        .transpose()
+}
+
+/// Creates a SegQueue populated with file paths from bind data.
+fn create_file_paths_queue(bind_data: &VortexBindData) -> SegQueue<PathBuf> {
+    let file_paths = SegQueue::new();
+    for path in bind_data.file_paths.iter() {
+        file_paths.push(path.clone());
+    }
+    file_paths
+}
+
 impl TableFunction for VortexTableFunction {
     type BindData = VortexBindData;
     type GlobalState = VortexGlobalData;
@@ -156,7 +206,7 @@ impl TableFunction for VortexTableFunction {
 
                 let array_iter = file
                     .scan()?
-                    .with_projection(global_state.project_expr.clone())
+                    .with_projection(global_state.projection_expr.clone())
                     .with_filter(global_state.filter_expr.clone())
                     .into_array_iter()
                     .map_err(|e| vortex_err!("Failed to create array iterator: {}", e))?;
@@ -184,62 +234,24 @@ impl TableFunction for VortexTableFunction {
         Ok(())
     }
 
-    fn init_global(init: &TableInitInput<Self>) -> VortexResult<Self::GlobalState> {
-        let bind_data = init.bind_data();
-        let file_paths = SegQueue::new();
-
-        // Skip the first file path, as the file is opened during bind.
-        for path in bind_data.file_paths.iter() {
-            file_paths.push(path.clone());
-        }
-
-        let complex_filter = and_collect(init.bind_data().filter_exprs.clone());
-
-        let projection_ids = init.projection_ids().unwrap_or(&[]);
-        let column_ids = init.column_ids();
-
-        let projected_ids = projection_ids.iter().map(|p| column_ids[p.as_usize()]);
-        let project_expr = select(
-            projected_ids
-                .map(|idx| {
-                    init.bind_data()
-                        .column_names
-                        .get(idx.as_usize())
-                        .vortex_expect("prune idx in column names")
-                })
-                .map(|s| Arc::from(s.clone()))
-                .collect::<FieldNames>(),
-            root(),
-        );
-
-        let filter = init
-            .table_filter_set()
-            .and_then(|filter| {
-                filter
-                    .into_iter()
-                    .map(|(idx, ex)| {
-                        let name = init
-                            .bind_data()
-                            .column_names
-                            .get(idx.as_usize())
-                            .vortex_expect("exists");
-                        try_from_table_filter(&ex, name)
-                    })
-                    .reduce(|l, r| Ok(and(l?, r?)))
-            })
-            .transpose()?;
+    fn init_global(init_input: &TableInitInput<Self>) -> VortexResult<Self::GlobalState> {
+        let bind_data = init_input.bind_data();
+        let file_paths = create_file_paths_queue(bind_data);
+        let projection_expr = create_projection_expr(init_input);
+        let filter_expr = create_table_filter_expr(init_input)?;
 
-        let filter_expr = complex_filter
+        let complex_filter_expr = and_collect(init_input.bind_data().filter_exprs.clone());
+        let filter_expr = complex_filter_expr
             .into_iter()
-            .chain(filter)
+            .chain(filter_expr)
             .reduce(and)
             .unwrap_or_else(|| lit(true));
 
         Ok(VortexGlobalData {
             file_paths,
             _is_first_file_processed: std::sync::atomic::AtomicBool::new(false),
             filter_expr,
-            project_expr,
+            projection_expr,
         })
     }