vcftools: address remaining review comments

psuszyns · psuszyns · commit a2d036378c01 · 2026-03-12T11:41:31.000+01:00
diff --git a/datafusion/bio-function-vcftools/src/logical/optimizer_rule.rs b/datafusion/bio-function-vcftools/src/logical/optimizer_rule.rs
@@ -216,7 +216,13 @@ struct TraversalResult<'a> {
 ///
 /// This handles transformation CTEs by traversing multiple Projection layers and
 /// composing expressions.
-fn traverse_to_unnest(plan: &LogicalPlan) -> Option<TraversalResult<'_>> {
+///
+/// # Returns
+///
+/// - `Ok(Some(result))` if the pattern matches and traversal succeeded
+/// - `Ok(None)` if the pattern doesn't match (no Unnest found)
+/// - `Err(e)` if traversal failed due to expression resolution errors
+fn traverse_to_unnest(plan: &LogicalPlan) -> Result<Option<TraversalResult<'_>>> {
     // Skip SubqueryAlias wrappers
     let plan = skip_wrappers(plan);
 
@@ -266,14 +272,16 @@ fn traverse_to_unnest(plan: &LogicalPlan) -> Option<TraversalResult<'_>> {
                 });
             }
 
-            Some(TraversalResult {
+            Ok(Some(TraversalResult {
                 unnest,
                 column_definitions,
-            })
+            }))
         }
         LogicalPlan::Projection(projection) => {
             // Recurse into child
-            let child_result = traverse_to_unnest(projection.input.as_ref())?;
+            let Some(child_result) = traverse_to_unnest(projection.input.as_ref())? else {
+                return Ok(None);
+            };
 
             // Build new column definitions by resolving projection expressions
             let mut new_definitions = HashMap::new();
@@ -289,28 +297,33 @@ fn traverse_to_unnest(plan: &LogicalPlan) -> Option<TraversalResult<'_>> {
                 };
 
                 // Resolve the expression against child definitions
-                let resolved = resolve_expr(&inner_expr, &child_result.column_definitions);
+                let resolved = resolve_expr(&inner_expr, &child_result.column_definitions)?;
                 new_definitions.insert(alias, resolved);
             }
 
-            Some(TraversalResult {
+            Ok(Some(TraversalResult {
                 unnest: child_result.unnest,
                 column_definitions: new_definitions,
-            })
+            }))
         }
         _ => {
             trace!(
                 "traverse_to_unnest: expected Projection or Unnest, got {}",
                 plan_type_name(plan)
             );
-            None
+            Ok(None)
         }
     }
 }
 
 /// Resolve an expression by substituting column references with their definitions.
 /// Uses DataFusion's `transform` to recursively traverse all expression variants.
-fn resolve_expr(expr: &Expr, definitions: &HashMap<String, Expr>) -> Expr {
+///
+/// # Errors
+///
+/// Returns an error if the expression tree traversal fails (e.g., due to an
+/// unexpected expression variant or internal DataFusion error).
+fn resolve_expr(expr: &Expr, definitions: &HashMap<String, Expr>) -> Result<Expr> {
     expr.clone()
         .transform(|e| {
             if let Expr::Column(col) = &e {
@@ -321,7 +334,6 @@ fn resolve_expr(expr: &Expr, definitions: &HashMap<String, Expr>) -> Expr {
             Ok(Transformed::no(e))
         })
         .map(|t| t.data)
-        .unwrap_or_else(|_| expr.clone())
 }
 
 /// Attempt to detect and optimize the pattern.
@@ -355,9 +367,16 @@ fn try_optimize(plan: &LogicalPlan) -> Option<Result<LogicalPlan>> {
     };
 
     // Traverse to find Unnest while collecting column definitions
-    let Some(traversal) = traverse_to_unnest(aggregate.input.as_ref()) else {
-        trace!("traverse_to_unnest returned None");
-        return None;
+    let traversal = match traverse_to_unnest(aggregate.input.as_ref()) {
+        Ok(Some(t)) => t,
+        Ok(None) => {
+            trace!("traverse_to_unnest returned None");
+            return None;
+        }
+        Err(e) => {
+            // Pattern matched but traversal failed - propagate error
+            return Some(Err(e));
+        }
     };
     let unnest_plan = traversal.unnest;
     let column_definitions = traversal.column_definitions;
@@ -410,7 +429,10 @@ fn try_optimize(plan: &LogicalPlan) -> Option<Result<LogicalPlan>> {
         if let Expr::AggregateFunction(AggregateFunction { params, .. }) = expr {
             if let Some(arg) = params.args.first() {
                 // Resolve the argument through the column definitions
-                let resolved = resolve_expr(arg, &column_definitions);
+                let resolved = match resolve_expr(arg, &column_definitions) {
+                    Ok(r) => r,
+                    Err(e) => return Some(Err(e)),
+                };
                 trace!("Resolved array_agg argument: {arg:?} -> {resolved:?}");
                 transform_exprs.push(resolved);
             } else {
diff --git a/datafusion/bio-function-vcftools/src/physical/fused_array_transform_exec.rs b/datafusion/bio-function-vcftools/src/physical/fused_array_transform_exec.rs
@@ -323,7 +323,9 @@ impl FusedArrayTransformStream {
             for col_name in &self.array_columns {
                 let idx = input_schema.index_of(col_name)?;
                 let col = batch.column(idx);
-                let output_array = self.apply_identity_transform(col, col_name)?;
+
+                let filtered_col = datafusion::arrow::compute::filter(col.as_ref(), &bool_mask)?;
+                let output_array = self.apply_identity_transform(&filtered_col, col_name)?;
                 output_columns.push(output_array);
             }
         } else {
@@ -627,28 +629,35 @@ mod tests {
     use datafusion::arrow::datatypes::{Field, Schema};
     use datafusion::physical_plan::test::TestMemoryExec;
 
-    fn create_test_batch() -> RecordBatch {
+    fn create_test_batch(
+        row0_lista: Vec<f64>,
+        row0_listb: Vec<f64>,
+        row1_lista: Vec<f64>,
+        row1_listb: Vec<f64>
+    ) -> RecordBatch {
         let mut list_builder_a = ListBuilder::new(Float64Builder::new());
         let mut list_builder_b = ListBuilder::new(Float64Builder::new());
 
-        // Row 0: [1.0, 2.0, 3.0], [10.0, 20.0, 30.0]
-        list_builder_a.values().append_value(1.0);
-        list_builder_a.values().append_value(2.0);
-        list_builder_a.values().append_value(3.0);
+        // Row 0
+        for val in row0_lista {
+            list_builder_a.values().append_value(val);
+        }
         list_builder_a.append(true);
 
-        list_builder_b.values().append_value(10.0);
-        list_builder_b.values().append_value(20.0);
-        list_builder_b.values().append_value(30.0);
+        for val in row0_listb {
+            list_builder_b.values().append_value(val);
+        }
         list_builder_b.append(true);
 
-        // Row 1: [4.0, 5.0], [40.0, 50.0]
-        list_builder_a.values().append_value(4.0);
-        list_builder_a.values().append_value(5.0);
+        // Row 1
+        for val in row1_lista {
+            list_builder_a.values().append_value(val);
+        }
         list_builder_a.append(true);
 
-        list_builder_b.values().append_value(40.0);
-        list_builder_b.values().append_value(50.0);
+        for val in row1_listb {
+            list_builder_b.values().append_value(val);
+        }
         list_builder_b.append(true);
 
         let arr_a = list_builder_a.finish();
@@ -677,9 +686,18 @@ mod tests {
         .unwrap()
     }
 
+    macro_rules! create_test_batch {
+        ($row0_lista: expr, $row0_listb: expr, $row1_lista: expr, $row1_listb: expr) => {
+            create_test_batch($row0_lista, $row0_listb, $row1_lista, $row1_listb)
+        };
+        () => {
+            create_test_batch(vec![1.0, 2.0, 3.0], vec![10.0, 20.0, 30.0], vec![4.0, 5.0], vec![40.0, 50.0])
+        };
+    }
+
     #[tokio::test]
     async fn test_identity_transform() {
-        let batch = create_test_batch();
+        let batch = create_test_batch!();
         let schema = batch.schema();
 
         let mem_exec = TestMemoryExec::try_new(&[vec![batch.clone()]], schema, None).unwrap();
@@ -695,11 +713,41 @@ mod tests {
 
         // Schema should have 2 fields: metadata + values_a_out
         assert_eq!(fused.schema().fields().len(), 2);
+
+        let ctx = Arc::new(TaskContext::default());
+        let mut stream = fused.execute(0, ctx).unwrap();
+        let result_batch = stream.next().await.unwrap().unwrap();
+        assert_eq!(result_batch.num_rows(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_identity_transform_with_empty_array() {
+        let batch = create_test_batch!(vec![], vec![], vec![4.0, 5.0], vec![40.0, 50.0]);
+        let schema = batch.schema();
+
+        let mem_exec = TestMemoryExec::try_new(&[vec![batch.clone()]], schema, None).unwrap();
+
+        let fused = FusedArrayTransformExec::try_new(
+            Arc::new(mem_exec),
+            vec!["values_a".to_string(), "values_b".to_string()],
+            vec!["metadata".to_string()],
+            vec!["values_a_out".to_string(), "values_b_out".to_string()],
+            vec![],
+        )
+        .unwrap();
+
+        // Schema should have 3 fields: metadata + values_a_out + values_b_out
+        assert_eq!(fused.schema().fields().len(), 3);
+        // row0 should be filtered out due to empty array, so only row1 remains
+        let ctx = Arc::new(TaskContext::default());
+        let mut stream = fused.execute(0, ctx).unwrap();
+        let result_batch = stream.next().await.unwrap().unwrap();
+        assert_eq!(result_batch.num_rows(), 1);
     }
 
     #[tokio::test]
     async fn test_execution() {
-        let batch = create_test_batch();
+        let batch = create_test_batch!();
         let schema = batch.schema();
 
         let mem_exec = TestMemoryExec::try_new(&[vec![batch.clone()]], schema, None).unwrap();
diff --git a/datafusion/bio-function-vcftools/tests/integration_test.rs b/datafusion/bio-function-vcftools/tests/integration_test.rs
@@ -959,7 +959,7 @@ async fn test_mismatched_array_lengths() {
         .to_string();
     assert!(
         plan_str_optimized.contains("FusedArrayTransform"),
-        "FusedArrayTransform optimization was NOT applied for empty arrays case! Physical plan:\n{plan_str_optimized}"
+        "FusedArrayTransform optimization was NOT applied for mismatched lengths case! Physical plan:\n{plan_str_optimized}"
     );
     let df_optimized2 = ctx_optimized.sql(sql).await.unwrap();
     let optimized_results = df_optimized2.collect().await.unwrap();
@@ -973,7 +973,7 @@ async fn test_mismatched_array_lengths() {
         .to_string();
     assert_eq!(
         baseline_str, optimized_str,
-        "Baseline and optimized results differ for empty arrays case"
+        "Baseline and optimized results differ for mismatched lengths case"
     );
 }
 
