add coverage

Author: friendlymatthew

datafusion/datasource-parquet/src/row_filter.rs

@@ -723,6 +723,8 @@ fn build_filter_schema(
     regular_indices: &[usize],
     struct_field_accesses: &[StructFieldAccess],
 ) -> SchemaRef {
+    let regular_set: BTreeSet<usize> = regular_indices.iter().copied().collect();
+
     let all_indices = regular_indices
         .iter()
         .copied()
@@ -738,6 +740,15 @@ fn build_filter_schema(
         .map(|&idx| {
             let field = file_schema.field(idx);
 
+            // if this column appears as a regular (whole-column) reference,
+            // keep the full type
+            //
+            // Pruning is only valid when the column is accessed exclusively
+            // through struct field accesses
+            if regular_set.contains(&idx) {
+                return Arc::new(field.clone());
+            }
+
             // collect all field paths that access this root struct column
             let field_paths = struct_field_accesses
                 .iter()
@@ -752,7 +763,6 @@ fn build_filter_schema(
                 .collect::<Vec<_>>();
 
             if field_paths.is_empty() {
-                // its a regular column - use the full type
                 return Arc::new(field.clone());
             }
 
@@ -1027,6 +1037,8 @@ mod test {
     use parquet::file::reader::{FileReader, SerializedFileReader};
     use tempfile::NamedTempFile;
 
+    use datafusion_physical_expr::expressions::Column as PhysicalColumn;
+
     // List predicates used by the decoder should be accepted for pushdown
     #[test]
     fn test_filter_candidate_builder_supports_list_types() {
@@ -1883,6 +1895,86 @@ mod test {
         assert_eq!(file_metrics.pushdown_rows_matched.value(), 2);
     }
 
+    #[test]
+    fn projection_read_plan_preserves_full_struct() {
+        // Schema: id (Int32), s (Struct{value: Int32, label: Utf8})
+        // Parquet leaves: id=0, s.value=1, s.label=2
+        let struct_fields: Fields = vec![
+            Arc::new(Field::new("value", DataType::Int32, false)),
+            Arc::new(Field::new("label", DataType::Utf8, false)),
+        ]
+        .into();
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("s", DataType::Struct(struct_fields.clone()), false),
+        ]));
+
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3])),
+                Arc::new(StructArray::new(
+                    struct_fields,
+                    vec![
+                        Arc::new(Int32Array::from(vec![10, 20, 30])) as _,
+                        Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
+                    ],
+                    None,
+                )),
+            ],
+        )
+        .unwrap();
+
+        let file = NamedTempFile::new().expect("temp file");
+        let mut writer =
+            ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None)
+                .expect("writer");
+        writer.write(&batch).expect("write batch");
+        writer.close().expect("close writer");
+
+        let reader_file = file.reopen().expect("reopen file");
+        let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file)
+            .expect("reader builder");
+        let metadata = builder.metadata().clone();
+        let file_schema = builder.schema().clone();
+        let schema_descr = metadata.file_metadata().schema_descr();
+
+        // Simulate SELECT * output projection: Column("id") and Column("s")
+        // Plus a get_field(s, 'value') expression from the pushed-down filter
+        let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
+            Arc::new(PhysicalColumn::new("id", 0)),
+            Arc::new(PhysicalColumn::new("s", 1)),
+            logical2physical(
+                &get_field().call(vec![
+                    col("s"),
+                    Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None),
+                ]),
+                &file_schema,
+            ),
+        ];
+
+        let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr);
+
+        // The projected schema must have the FULL struct type because Column("s")
+        // is in the projection. It should NOT be narrowed to Struct{value: Int32}.
+        let s_field = read_plan.projected_schema.field_with_name("s").unwrap();
+        assert_eq!(
+            s_field.data_type(),
+            &DataType::Struct(
+                vec![
+                    Arc::new(Field::new("value", DataType::Int32, false)),
+                    Arc::new(Field::new("label", DataType::Utf8, false)),
+                ]
+                .into()
+            ),
+        );
+
+        // all3 Parquet leaves should be in the projection mask
+        let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]);
+        assert_eq!(read_plan.projection_mask, expected_mask,);
+    }
+
     /// Sanity check that the given expression could be evaluated against the given schema without any errors.
     /// This will fail if the expression references columns that are not in the schema or if the types of the columns are incompatible, etc.
     fn check_expression_can_evaluate_against_schema(

datafusion/sqllogictest/test_files/projection_pushdown.slt

@@ -1994,9 +1994,83 @@ WHERE COALESCE(get_field(s, 'f1'), get_field(s, 'f2')) = 1;
 ----
 1
 
+#####################
+# Section 8: SELECT * with struct field filter
+#####################
+
+# When SELECT * includes the full struct but the filter only accesses a
+# sub-field (e.g. s['id']), the leaf-level projection must not narrow the
+# struct schema in the output. Previously build_projection_read_plan would
+# produce a schema with Struct("id": Int32) while the data still contained
+# Struct("id": Int32, "value": Utf8), causing an ArrowError.
+
+# 8.1: SELECT * with equality filter on struct sub-field
+query I?
+SELECT * FROM simple_struct WHERE s['value'] = 100;
+----
+1 {value: 100, label: alpha}
+
+# 8.2: Explicit SELECT of whole struct with struct sub-field filter
+query ?
+SELECT s FROM simple_struct WHERE s['value'] = 100;
+----
+{value: 100, label: alpha}
+
+# 8.3: Whole struct + sub-field projection + sub-field filter
+query I?I
+SELECT s['value'], s, id FROM simple_struct WHERE s['value'] = 100;
+----
+100 {value: 100, label: alpha} 1
+
+# 8.4: Whole struct in output, filter on a different sub-field than projected
+query ?T
+SELECT s, s['label'] FROM simple_struct WHERE s['value'] > 200;
+----
+{value: 300, label: delta} delta
+{value: 250, label: epsilon} epsilon
+
+# 8.5: Filter references both sub-fields, output includes whole struct
+query I?
+SELECT id, s FROM simple_struct WHERE s['value'] > 100 AND s['label'] = 'beta';
+----
+2 {value: 200, label: beta}
+
+# 8.6: Only sub-field projection with sub-field filter (no whole struct — should prune)
+query II
+SELECT id, s['value'] FROM simple_struct WHERE s['value'] = 100;
+----
+1 100
+
+# 8.7: Nested struct — whole struct output with deeply nested field filter
+query I?
+SELECT * FROM nested_struct WHERE nested['outer']['inner'] > 15;
+----
+2 {outer: {inner: 20, name: two}, extra: y}
+3 {outer: {inner: 30, name: three}, extra: z}
+
+# 8.8: Nested struct — explicit whole struct select with sibling field filter
+query ?
+SELECT nested FROM nested_struct WHERE nested['extra'] = 'y';
+----
+{outer: {inner: 20, name: two}, extra: y}
+
+# 8.9: Nullable struct — whole struct output with sub-field filter
+query ?
+SELECT s FROM nullable_struct WHERE s['value'] > 100;
+----
+{value: 150, label: gamma}
+{value: 250, label: epsilon}
+
+# 8.10: Struct sub-field filter combined with top-level column filter
+query ?I
+SELECT s, id FROM simple_struct WHERE s['value'] > 100 AND id < 4;
+----
+{value: 200, label: beta} 2
+{value: 150, label: gamma} 3
+
 # Config reset
 
-# The SLT runner sets `target_partitions` to 4 instead of using the default, so 
+# The SLT runner sets `target_partitions` to 4 instead of using the default, so
 # reset it explicitly.
 statement ok
 SET datafusion.execution.target_partitions = 4;