adriangb commented on code in PR #24090:
URL: https://github.com/apache/datafusion/pull/24090#discussion_r3718149700
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -94,6 +99,13 @@ pub(crate) struct PushdownChecker<'schema> {
required_columns: Vec<usize>,
/// Struct field accesses via `get_field`.
struct_field_accesses: Vec<StructFieldAccess>,
+ /// Whole-column casts to a narrower nested type
+ /// (`CAST(col AS narrower_struct)`). Only collected when
+ /// [`Self::with_cast_collection`] enables it (projection analysis);
+ /// filter pushdown leaves this off.
Review Comment:
Is it intentional that this not work for filter push down? I don't
immediately see any reason why it couldn't work.
##########
datafusion/core/tests/parquet/expr_adapter.rs:
##########
@@ -1204,3 +1204,671 @@ async fn
test_physical_expr_adapter_factory_reuse_across_tables() {
];
assert_batches_eq!(expected, &batches);
}
+
+// ---------------------------------------------------------------------------
+// Nested projection pruning: when the table schema declares a nested column
+// narrower than the physical Parquet file, the scan should only read the
+// leaves the declared schema names, instead of reading the whole column and
+// discarding the extra subfields in the adapter-inserted cast.
+//
+// Each test registers two tables against the *same* physical file: `t_narrow`
+// (the declared schema under test) and `t_full` (the file's own physical
+// schema, so no cast is inserted and the scan always reads every leaf). That
+// gives a same-context upper bound to compare `bytes_scanned` against,
+// without needing a config flag to disable pruning.
Review Comment:
Good idea 👍🏻
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -419,6 +484,160 @@ pub(crate) fn build_projection_read_plan(
read_plan
}
+/// Builds a [`ParquetReadPlan`] when at least one projected root column is
+/// consumed through a cast to a narrower nested type.
+///
+/// Per root, in ascending root-index order:
+/// - roots referenced as whole columns keep every leaf and their full
+/// physical field (whole-column reads take precedence; cast accesses on
+/// such roots were already dropped by the caller);
+/// - roots consumed through a cast, and not also through a `get_field`
+/// access on the same root, keep only the leaves the cast target names
+/// (see `crate::nested_schema_pruning`);
+/// - roots consumed only through `get_field` accesses keep the union of the
+/// leaves those accesses reach, as before;
+/// - any other referenced root, a cast that can't be safely clipped (see
+/// `nested_schema_pruning::clip_for_cast`), or a root reached by both a
+/// cast and a `get_field` access (not produced by
+/// `DefaultPhysicalExprAdapter`, which always routes a `get_field` over a
+/// narrowed column through the same cast rather than a separate access,
+/// but a custom `PhysicalExprAdapter` could in principle inject both),
+/// falls back to a full read of that root.
+fn build_read_plan_with_cast_clipping(
Review Comment:
This seems like a possible source of fragility, e.g. because of the indirect
coupling to `DefaultPhysicalExprAdapter`. I think the failure modes are okay
(loss of optimization) and we can fix bugs / edge cases if they are added in
the future. For now though it's worth measuring coverage on this function and
making sure it's well tested.
##########
datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt:
##########
@@ -0,0 +1,83 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+##########
+# Nested projection pruning: a table whose declared nested type is narrower
+# than the Parquet file's physical type reads only the declared leaves.
+# The bytes-scanned assertions live in the Rust tests
+# (datafusion/core/tests/parquet/expr_adapter.rs); this file covers the
+# end-to-end SQL correctness path.
Review Comment:
Thank you. I'll try to check that we have full coverage via SLT tests and if
there is any low hanging fruit etsting we can add as SLTs.
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -419,6 +484,160 @@ pub(crate) fn build_projection_read_plan(
read_plan
}
+/// Builds a [`ParquetReadPlan`] when at least one projected root column is
+/// consumed through a cast to a narrower nested type.
+///
+/// Per root, in ascending root-index order:
+/// - roots referenced as whole columns keep every leaf and their full
+/// physical field (whole-column reads take precedence; cast accesses on
+/// such roots were already dropped by the caller);
+/// - roots consumed through a cast, and not also through a `get_field`
+/// access on the same root, keep only the leaves the cast target names
+/// (see `crate::nested_schema_pruning`);
+/// - roots consumed only through `get_field` accesses keep the union of the
+/// leaves those accesses reach, as before;
+/// - any other referenced root, a cast that can't be safely clipped (see
+/// `nested_schema_pruning::clip_for_cast`), or a root reached by both a
+/// cast and a `get_field` access (not produced by
+/// `DefaultPhysicalExprAdapter`, which always routes a `get_field` over a
+/// narrowed column through the same cast rather than a separate access,
+/// but a custom `PhysicalExprAdapter` could in principle inject both),
+/// falls back to a full read of that root.
+fn build_read_plan_with_cast_clipping(
Review Comment:
I also want to check that all operations are O(projected_columns) or
O(logical_plan_column) and ideally not O(schema) in some way, or at least not
e.g. cuadratic with schema or struct column size.
##########
datafusion/core/tests/parquet/expr_adapter.rs:
##########
@@ -1204,3 +1204,671 @@ async fn
test_physical_expr_adapter_factory_reuse_across_tables() {
];
assert_batches_eq!(expected, &batches);
}
+
+// ---------------------------------------------------------------------------
+// Nested projection pruning: when the table schema declares a nested column
+// narrower than the physical Parquet file, the scan should only read the
+// leaves the declared schema names, instead of reading the whole column and
+// discarding the extra subfields in the adapter-inserted cast.
+//
+// Each test registers two tables against the *same* physical file: `t_narrow`
+// (the declared schema under test) and `t_full` (the file's own physical
+// schema, so no cast is inserted and the scan always reads every leaf). That
+// gives a same-context upper bound to compare `bytes_scanned` against,
+// without needing a config flag to disable pruning.
+// ---------------------------------------------------------------------------
+
+mod nested_projection_pruning {
+ use super::*;
+ use arrow::buffer::NullBuffer;
+ use datafusion::physical_plan::collect;
+ use datafusion_physical_plan::metrics::MetricsSet;
+
+ use crate::parquet::utils::MetricsFinder;
+
+ const NUM_ELEMENTS: usize = 64;
+ const PAD_LEN: usize = 2048;
+
+ /// Physical item struct written to the file: the narrow fields plus fat
+ /// pads the narrow table schema will not mention. `x` is Int32 in the
+ /// file (the narrow schema declares Int64 to also exercise leaf
+ /// promotion).
+ fn wide_item_fields() -> Fields {
+ Fields::from(vec![
+ Field::new("x", DataType::Int32, false),
+ Field::new("y", DataType::Utf8, true),
+ Field::new("pad_a", DataType::Utf8, false),
+ Field::new("pad_b", DataType::Utf8, false),
+ Field::new("pad_c", DataType::Utf8, false),
+ ])
+ }
+
+ /// The narrow item struct one table declares: a subset of the physical
+ /// fields in a different order, a promoted leaf type for `x`, plus `z`
+ /// which does not exist in the file (null-filled by the cast).
+ fn narrow_item_fields() -> Fields {
+ Fields::from(vec![
+ Field::new("y", DataType::Utf8, true),
+ Field::new("x", DataType::Int64, true),
+ Field::new("z", DataType::Int64, true),
+ ])
+ }
+
+ fn wide_struct_values(validity: Option<NullBuffer>) -> StructArray {
+ let pad = |seed: usize| {
+ let base = "x".repeat(PAD_LEN);
+ Arc::new(StringArray::from_iter_values(
+ (0..NUM_ELEMENTS).map(|i| format!("{}{base}", seed + i)),
+ )) as ArrayRef
+ };
+ StructArray::new(
+ wide_item_fields(),
+ vec![
+ Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)),
+ Arc::new(StringArray::from_iter_values(
+ (0..NUM_ELEMENTS).map(|i| format!("y-{i}")),
+ )),
+ pad(1000),
+ pad(2000),
+ pad(3000),
+ ],
+ validity,
+ )
+ }
+
+ fn wide_list_schema() -> SchemaRef {
+ let item = Arc::new(Field::new(
+ "item",
+ DataType::Struct(wide_item_fields()),
+ true,
+ ));
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("events", DataType::List(item), true),
+ ]))
+ }
+
+ /// File batch: `id Int32`, `events List<wide struct>` (one element per
row).
+ fn wide_list_batch() -> RecordBatch {
+ let schema = wide_list_schema();
+ let item = match schema.field(1).data_type() {
+ DataType::List(item) => Arc::clone(item),
+ other => unreachable!("expected List, got {other:?}"),
+ };
+ let events = ListArray::new(
+ item,
+ OffsetBuffer::from_lengths(std::iter::repeat_n(1, NUM_ELEMENTS)),
+ Arc::new(wide_struct_values(None)),
+ None,
+ );
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)),
+ Arc::new(events),
+ ],
+ )
+ .unwrap()
+ }
+
+ fn narrow_list_table_schema() -> SchemaRef {
+ let item = Arc::new(Field::new(
+ "item",
+ DataType::Struct(narrow_item_fields()),
+ true,
+ ));
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("events", DataType::List(item), true),
+ ]))
+ }
+
+ fn wide_struct_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("s", DataType::Struct(wide_item_fields()), true),
+ ]))
+ }
+
+ /// File batch: `id Int32`, `s <wide struct>`, with per-row struct
+ /// validity so struct-level nullability can be asserted.
+ fn wide_struct_batch() -> RecordBatch {
+ // rows 0, 10, 20, ... have a NULL struct
+ let validity =
+ NullBuffer::from((0..NUM_ELEMENTS).map(|i| i % 10 !=
0).collect::<Vec<_>>());
+ RecordBatch::try_new(
+ wide_struct_schema(),
+ vec![
+ Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)),
+ Arc::new(wide_struct_values(Some(validity))),
+ ],
+ )
+ .unwrap()
+ }
+
+ fn narrow_struct_table_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("s", DataType::Struct(narrow_item_fields()), true),
+ ]))
+ }
+
+ /// Registers `t_narrow` (the schema under test) and `t_full` (the file's
+ /// own physical schema, so no cast is inserted) against the same store.
+ async fn register_narrow_and_full(
+ ctx: &SessionContext,
+ store: Arc<dyn ObjectStore>,
+ narrow_schema: SchemaRef,
+ full_schema: SchemaRef,
+ ) {
+ let store_url = ObjectStoreUrl::parse("memory://").unwrap();
+ ctx.register_object_store(store_url.as_ref(), store);
+
+ for (name, schema) in [("t_narrow", narrow_schema), ("t_full",
full_schema)] {
+ let config = ListingTableConfig::new(
+ ListingTableUrl::parse("memory:///data/").unwrap(),
+ )
+ .infer_options(&ctx.state())
+ .await
+ .unwrap()
+ .with_schema(schema)
+
.with_expr_adapter_factory(Arc::new(DefaultPhysicalExprAdapterFactory));
+ let table = ListingTable::try_new(config).unwrap();
+ ctx.register_table(name, Arc::new(table)).unwrap();
+ }
+ }
+
+ async fn setup_with_config(
+ batches: Vec<(&str, RecordBatch)>,
+ narrow_schema: SchemaRef,
+ full_schema: SchemaRef,
+ cfg: SessionConfig,
+ ) -> SessionContext {
+ let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+ for (name, batch) in batches {
+ write_parquet(batch, Arc::clone(&store),
&format!("data/{name}")).await;
+ }
+ let ctx = SessionContext::new_with_config(cfg);
+ register_narrow_and_full(&ctx, store, narrow_schema,
full_schema).await;
+ ctx
+ }
+
+ async fn setup(
+ batches: Vec<(&str, RecordBatch)>,
+ narrow_schema: SchemaRef,
+ full_schema: SchemaRef,
+ ) -> SessionContext {
+ setup_with_config(
+ batches,
+ narrow_schema,
+ full_schema,
+ SessionConfig::new().with_collect_statistics(false),
+ )
+ .await
+ }
+
+ async fn run(ctx: &SessionContext, sql: &str) -> (Vec<RecordBatch>,
MetricsSet) {
+ let df = ctx.sql(sql).await.unwrap();
+ let (state, logical) = df.into_parts();
+ let plan = state.create_physical_plan(&logical).await.unwrap();
+ let batches = collect(Arc::clone(&plan),
state.task_ctx()).await.unwrap();
+ let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap();
+ (batches, metrics)
+ }
+
+ fn bytes_scanned(metrics: &MetricsSet) -> usize {
+ metrics
+ .sum(|m| m.value().name() == "bytes_scanned")
+ .map(|v| v.as_usize())
+ .expect("bytes_scanned metric")
+ }
+
+ /// Run `narrow_sql` against `t_narrow` and `full_sql` against `t_full`;
+ /// assert the narrow scan read strictly less than half of the full
+ /// scan's bytes (the pads dominate the file), and return the narrow
+ /// scan's results for correctness assertions.
+ ///
+ /// The two SQL strings need not have the same shape: a `get_field` over
+ /// a narrowed struct clips to the *cast target*, not further down to the
+ /// specific field accessed (see `prunes_get_field_on_narrowed_struct`),
+ /// so comparing against the same `get_field` query on `t_full` would
+ /// unfairly compare this clip against `get_field`'s own, more precise,
+ /// single-leaf pruning (which only applies when there is no cast in the
+ /// way). Callers that aren't in that situation can just pass the same
+ /// query shape with the table name substituted.
+ async fn assert_prunes(
+ batches: Vec<(&str, RecordBatch)>,
+ narrow_schema: SchemaRef,
+ full_schema: SchemaRef,
+ narrow_sql: &str,
+ full_sql: &str,
+ ) -> Vec<RecordBatch> {
+ let ctx = setup(batches, narrow_schema, full_schema).await;
+
+ let (result_narrow, metrics_narrow) = run(&ctx, narrow_sql).await;
+ let (_result_full, metrics_full) = run(&ctx, full_sql).await;
+
+ let (narrow_bytes, full_bytes) =
+ (bytes_scanned(&metrics_narrow), bytes_scanned(&metrics_full));
+ assert!(
+ narrow_bytes * 2 < full_bytes,
+ "expected pruned scan to read less than half of {full_bytes}
bytes, \
+ read {narrow_bytes}: {narrow_sql}"
+ );
+ result_narrow
+ }
+
+ #[tokio::test]
+ async fn prunes_list_of_struct() {
+ // Narrow schema over the wide file: subset of fields, reordered,
+ // promoted leaf (x: Int32 -> Int64), missing subfield z null-filled.
+ let batches = assert_prunes(
+ vec![("wide.parquet", wide_list_batch())],
+ narrow_list_table_schema(),
+ wide_list_schema(),
+ "SELECT events FROM t_narrow ORDER BY id",
+ "SELECT events FROM t_full ORDER BY id",
+ )
+ .await;
+
+ let events = batches[0].column(0);
+ let list = events.as_any().downcast_ref::<ListArray>().unwrap();
+ let items = list
+ .values()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(items.fields().len(), 3);
+ let x = items
+ .column_by_name("x")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(x.value(5), 5);
+ let z = items.column_by_name("z").unwrap();
+ assert_eq!(z.null_count(), z.len(), "z is not in the file");
+ }
+
+ #[tokio::test]
+ async fn prunes_top_level_struct() {
+ assert_prunes(
+ vec![("wide.parquet", wide_struct_batch())],
+ narrow_struct_table_schema(),
+ wide_struct_schema(),
+ "SELECT s FROM t_narrow ORDER BY id",
+ "SELECT s FROM t_full ORDER BY id",
+ )
+ .await;
+ }
+
+ /// Struct-level nullability must survive the clip: rows where the struct
+ /// itself is NULL stay NULL (not `{y: NULL, x: NULL, z: NULL}`).
+ #[tokio::test]
+ async fn preserves_struct_nullability() {
+ let batches = assert_prunes(
+ vec![("wide.parquet", wide_struct_batch())],
+ narrow_struct_table_schema(),
+ wide_struct_schema(),
+ "SELECT id, s IS NULL AS s_null, s FROM t_narrow ORDER BY id",
+ "SELECT id, s IS NULL AS s_null, s FROM t_full ORDER BY id",
+ )
+ .await;
+
+ let combined = concat_batches(&batches[0].schema(), &batches).unwrap();
+ let s_null = combined
+ .column(1)
+ .as_any()
+ .downcast_ref::<BooleanArray>()
+ .unwrap();
+ for i in 0..NUM_ELEMENTS {
+ assert_eq!(s_null.value(i), i % 10 == 0, "row {i}");
+ }
+ }
+
+ /// `get_field` on a schema-narrowed struct becomes
+ /// `get_field(CAST(s), 'x')`; the read clips to the cast target (every
+ /// field the *narrow* schema declares), not further down to just `x`.
+ /// The fair "no clipping happened" baseline is therefore reading every
+ /// physical leaf of `s` (`SELECT s FROM t_full`), not the same
+ /// `get_field` query against `t_full`. That query needs no cast at all
+ /// and takes `get_field`'s own, more precise, single-leaf pushdown path.
+ #[tokio::test]
+ async fn prunes_get_field_on_narrowed_struct() {
+ let batches = assert_prunes(
+ vec![("wide.parquet", wide_struct_batch())],
+ narrow_struct_table_schema(),
+ wide_struct_schema(),
+ "SELECT s['x'] AS x FROM t_narrow ORDER BY id",
+ "SELECT s FROM t_full ORDER BY id",
+ )
+ .await;
+ let combined = concat_batches(&batches[0].schema(), &batches).unwrap();
+ let x = combined
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(x.value(5), 5);
+ assert_eq!(x.value(NUM_ELEMENTS - 1), NUM_ELEMENTS as i64 - 1);
+ }
+
+ /// Mixed access: the whole (narrowed) column and a subfield of it.
+ #[tokio::test]
+ async fn prunes_mixed_struct_and_subfield_access() {
+ assert_prunes(
+ vec![("wide.parquet", wide_struct_batch())],
+ narrow_struct_table_schema(),
+ wide_struct_schema(),
+ "SELECT s, s['y'] AS y FROM t_narrow ORDER BY id",
+ "SELECT s, s['y'] AS y FROM t_full ORDER BY id",
+ )
+ .await;
+ }
+
+ /// Predicate on a primitive column with filter pushdown enabled while
+ /// the projected nested column is clipped.
+ #[tokio::test]
+ async fn prunes_with_filter_pushdown() {
+ let mut cfg = SessionConfig::new().with_collect_statistics(false);
+ cfg.options_mut().execution.parquet.pushdown_filters = true;
+ let ctx = setup_with_config(
+ vec![("wide.parquet", wide_list_batch())],
+ narrow_list_table_schema(),
+ wide_list_schema(),
+ cfg,
+ )
+ .await;
+
+ let filter = "WHERE id >= 32 ORDER BY id";
+ let (result_narrow, metrics_narrow) =
+ run(&ctx, &format!("SELECT events FROM t_narrow {filter}")).await;
+ let (_result_full, metrics_full) =
+ run(&ctx, &format!("SELECT events FROM t_full {filter}")).await;
+
+ let combined =
+ concat_batches(&result_narrow[0].schema(),
&result_narrow).unwrap();
+ assert_eq!(combined.num_rows(), NUM_ELEMENTS / 2);
+ assert!(bytes_scanned(&metrics_narrow) * 2 <
bytes_scanned(&metrics_full));
+ }
+
+ /// A scan over two files where one matches the table schema exactly (no
+ /// cast is inserted) and one is wider (clipped): both must be read
+ /// correctly in the same scan.
+ #[tokio::test]
+ async fn mixed_files_narrow_and_wide() {
+ // The physically-narrow file has exactly the table's item struct.
+ let narrow_item = narrow_item_fields();
+ let item = Arc::new(Field::new(
+ "item",
+ DataType::Struct(narrow_item.clone()),
+ true,
+ ));
+ let events = ListArray::new(
+ Arc::clone(&item),
+ OffsetBuffer::from_lengths([1]),
+ Arc::new(StructArray::new(
+ narrow_item,
+ vec![
+ Arc::new(StringArray::from(vec![Some("y-narrow")])) as
ArrayRef,
+ Arc::new(Int64Array::from(vec![Some(4242)])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
+ ],
+ None,
+ )),
+ None,
+ );
+ let narrow_batch = RecordBatch::try_new(
+ narrow_list_table_schema(),
+ vec![
+ Arc::new(Int32Array::from(vec![NUM_ELEMENTS as i32])),
+ Arc::new(events),
+ ],
+ )
+ .unwrap();
+
+ let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+ write_parquet(wide_list_batch(), Arc::clone(&store),
"data/wide.parquet").await;
+ write_parquet(narrow_batch, Arc::clone(&store),
"data/narrow.parquet").await;
+
+ let ctx = test_context();
+ register_memory_listing_table(
+ &ctx,
+ store,
+ "memory:///data/",
+ narrow_list_table_schema(),
+ )
+ .await;
+
+ let (batches, _) = run(
+ &ctx,
+ "SELECT id, e['x'] AS x, e['z'] AS z \
+ FROM (SELECT id, unnest(events) AS e FROM t) ORDER BY id",
+ )
+ .await;
+ let combined = concat_batches(&batches[0].schema(), &batches).unwrap();
+ assert_eq!(combined.num_rows(), NUM_ELEMENTS + 1);
+ let x = combined
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(x.value(NUM_ELEMENTS), 4242, "row from the narrow file");
+ let z = combined.column(2);
+ // z is null-filled for the wide file, present in the narrow file
+ assert_eq!(z.null_count(), NUM_ELEMENTS);
+ }
+
+ /// Regression test for the exact shape reported in
+ /// `datafusion-comet#4859`: a two-level `array<struct<...,
Review Comment:
Can we make this a real link?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
