adriangb commented on code in PR #15057:
URL: https://github.com/apache/datafusion/pull/15057#discussion_r1988140038
##########
datafusion/datasource-parquet/src/opener.rs:
##########
@@ -111,18 +109,18 @@ impl FileOpener for ParquetOpener {
.schema_adapter_factory
.create(projected_schema, Arc::clone(&self.table_schema));
let predicate = self.predicate.clone();
- let pruning_predicate = self.pruning_predicate.clone();
- let page_pruning_predicate = self.page_pruning_predicate.clone();
let table_schema = Arc::clone(&self.table_schema);
+ let filter_expression_rewriter_factory =
+ self.filter_expression_rewriter_factory.clone();
let reorder_predicates = self.reorder_filters;
let pushdown_filters = self.pushdown_filters;
- let enable_page_index = should_enable_page_index(
- self.enable_page_index,
- &self.page_pruning_predicate,
- );
+ let enable_page_index = self.enable_page_index &&
self.predicate.is_some();
Review Comment:
We need to decide upfront if we need the page index or not.
Previously the pruning predicate got calculated from the "file schema" that
users pass in (which is obviously _not_ the schema of the physical file we are
reading, it is expected to be the table schema - partition columns; I think
that should be renamed / reworked but that's for another day). This seems wrong
to me: I think it should be using the physical file schema since that's what
the rg/page stats it will be loading match. So *I think* moving it here
actually improves some edge cases. _But_ we now have to start loading the file
before we know if we have any page index pruning predicate, so we may load the
page index then realize we don't have any applicable filters. That should only
happen if we have a predicate and can't build a page pruning predicate from the
transformed predicate, I don't think that will be all that common.
##########
datafusion-examples/examples/struct_field_rewrite.rs:
##########
@@ -0,0 +1,353 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray,
StructArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::Fields;
+use async_trait::async_trait;
+
+use datafusion::assert_batches_eq;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRewriter,
+};
+use datafusion::common::{DFSchema, Result};
+use datafusion::datasource::file_expr_rewriter::FileExpressionRewriter;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource};
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use datafusion::logical_expr::utils::conjunction;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
+use datafusion::parquet::arrow::ArrowWriter;
+use datafusion::parquet::file::properties::WriterProperties;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::{expressions, ScalarFunctionExpr};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::lit;
+use futures::StreamExt;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::{ObjectStore, PutPayload};
+
+// Example showing how to implement custom filter rewriting for struct fields.
+//
+// In this example, we have a table with a struct column like:
+// struct_col: {"a": 1, "b": "foo"}
+//
+// Our custom TableProvider will use a FilterExpressionRewriter to rewrite
+// expressions like `struct_col['a'] = 10` to use a flattened column name
+// `_struct_col.a` if it exists in the file schema.
+#[tokio::main]
+async fn main() -> Result<()> {
+ println!("=== Creating example data with structs and flattened fields
===");
+ println!("NOTE: This example demonstrates filter expression rewriting for
struct field access.");
+ println!(" We deliberately create different values in the struct
field vs. flattened column");
+ println!(" to prove that the rewriting system is actually using the
flattened column.");
+
+ // Create sample data with both struct columns and flattened fields
+ let (table_schema, batch) = create_sample_data();
+
+ let store = InMemory::new();
+ let buf = {
+ let mut buf = vec![];
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(1)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props))
+ .expect("creating writer");
+
+ writer.write(&batch).expect("Writing batch");
+ writer.close().unwrap();
+ buf
+ };
+ let path = Path::from("example.parquet");
+ let payload = PutPayload::from_bytes(buf.into());
+ store.put(&path, payload).await?;
+
+ // Create a custom table provider that rewrites struct field access
+ let table_provider = Arc::new(ExampleTableProvider::new(table_schema));
+
+ // Set up query execution
+ let ctx = SessionContext::new();
+
+ // Register our table
+ ctx.register_table("structs", table_provider)?;
+
+ ctx.runtime_env().register_object_store(
+ ObjectStoreUrl::parse("memory://")?.as_ref(),
+ Arc::new(store),
+ );
+
+ println!("\n=== Showing all data ===");
+ let batches = ctx.sql("SELECT * FROM structs").await?.collect().await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ println!("\n=== Running query with struct field access and filter < 30
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("SELECT * FROM structs WHERE user_info['age'] > 30")
+ .await?
+ .collect()
+ .await?;
+
+ #[rustfmt::skip]
+ let expected = vec![
+ "+--------------------------+",
+ "| user_info |",
+ "+--------------------------+",
+ "| {name: Charlie, age: 35} |",
+ "+--------------------------+",
+ ];
+ assert_batches_eq!(expected, &batches);
+
+ println!("\n=== Running query with struct field access and filter > 130
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("explain analyze SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30")
+ .await?
+ .collect()
+ .await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ Ok(())
+}
+
+/// Create the example data with both struct fields and flattened fields
+fn create_sample_data() -> (SchemaRef, RecordBatch) {
+ // Create a schema with a struct column
+ let user_info_fields = Fields::from(vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("age", DataType::Int32, false),
+ ]);
+
+ let file_schema = Schema::new(vec![
+ Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ ),
+ // Include flattened fields (in real scenarios these might be in some
files but not others)
+ Field::new("_user_info.name", DataType::Utf8, true),
+ Field::new("_user_info.age", DataType::Int32, true),
+ ]);
+
+ let table_schema = Schema::new(vec![Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ )]);
+
+ // Create struct array for user_info
+ let names = StringArray::from(vec!["Alice", "Bob", "Charlie", "Dave"]);
+ let ages = Int32Array::from(vec![30, 25, 35, 22]);
+
+ let user_info = StructArray::from(vec![
+ (
+ Arc::new(Field::new("name", DataType::Utf8, false)),
+ Arc::new(names.clone()) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("age", DataType::Int32, false)),
+ Arc::new(ages.clone()) as ArrayRef,
+ ),
+ ]);
+
+ // Create a record batch with the data
+ let batch = RecordBatch::try_new(
+ Arc::new(file_schema.clone()),
+ vec![
+ Arc::new(user_info),
+ Arc::new(names), // Shredded name field
+ Arc::new(ages), // Shredded age field
+ ],
+ )
+ .unwrap();
+
+ (Arc::new(table_schema), batch)
+}
+
+/// Custom TableProvider that uses a StructFieldRewriter
+#[derive(Debug)]
+struct ExampleTableProvider {
+ schema: SchemaRef,
+}
+
+impl ExampleTableProvider {
+ fn new(schema: SchemaRef) -> Self {
+ Self { schema }
+ }
+}
+
+#[async_trait]
+impl TableProvider for ExampleTableProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ // Implementers can choose to mark these filters as exact or inexact.
+ // If marked as exact they cannot have false positives and must always
be applied.
+ // If marked as Inexact they can have false positives and at runtime
the rewriter
+ // can decide to not rewrite / ignore some filters since they will be
re-evaluated upstream.
+ Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
+ }
Review Comment:
I'm not sure how we're thinking about these.
Making them inexact is a pitty because it means we might still have to do
all of the work of pulling in the large column and re-filtering without
pushdown (at least until predicate pushdown is done).
Making the exact is dangerous because it's possible that they're not
evaluated completely as pushed-down filters:
https://github.com/apache/datafusion/blob/9382add72b929c553ca4976d1423d8ebbc80889d/datafusion/datasource-parquet/src/row_filter.rs#L333-L336
@alamb @tustvold any recollection from
https://github.com/apache/datafusion/pull/3380 as to why these specific
expressions do not get evaluated here? It seems to me that it'd be nice if this
layer of evaluation can do everything that a `FilterExec` can do so that it's
safer to mark filters as exact.
##########
datafusion-examples/examples/struct_field_rewrite.rs:
##########
@@ -0,0 +1,353 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray,
StructArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::Fields;
+use async_trait::async_trait;
+
+use datafusion::assert_batches_eq;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRewriter,
+};
+use datafusion::common::{DFSchema, Result};
+use datafusion::datasource::file_expr_rewriter::FileExpressionRewriter;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource};
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use datafusion::logical_expr::utils::conjunction;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
+use datafusion::parquet::arrow::ArrowWriter;
+use datafusion::parquet::file::properties::WriterProperties;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::{expressions, ScalarFunctionExpr};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::lit;
+use futures::StreamExt;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::{ObjectStore, PutPayload};
+
+// Example showing how to implement custom filter rewriting for struct fields.
+//
+// In this example, we have a table with a struct column like:
+// struct_col: {"a": 1, "b": "foo"}
+//
+// Our custom TableProvider will use a FilterExpressionRewriter to rewrite
+// expressions like `struct_col['a'] = 10` to use a flattened column name
+// `_struct_col.a` if it exists in the file schema.
+#[tokio::main]
+async fn main() -> Result<()> {
+ println!("=== Creating example data with structs and flattened fields
===");
+ println!("NOTE: This example demonstrates filter expression rewriting for
struct field access.");
+ println!(" We deliberately create different values in the struct
field vs. flattened column");
+ println!(" to prove that the rewriting system is actually using the
flattened column.");
+
+ // Create sample data with both struct columns and flattened fields
+ let (table_schema, batch) = create_sample_data();
+
+ let store = InMemory::new();
+ let buf = {
+ let mut buf = vec![];
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(1)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props))
+ .expect("creating writer");
+
+ writer.write(&batch).expect("Writing batch");
+ writer.close().unwrap();
+ buf
+ };
+ let path = Path::from("example.parquet");
+ let payload = PutPayload::from_bytes(buf.into());
+ store.put(&path, payload).await?;
+
+ // Create a custom table provider that rewrites struct field access
+ let table_provider = Arc::new(ExampleTableProvider::new(table_schema));
+
+ // Set up query execution
+ let ctx = SessionContext::new();
+
+ // Register our table
+ ctx.register_table("structs", table_provider)?;
+
+ ctx.runtime_env().register_object_store(
+ ObjectStoreUrl::parse("memory://")?.as_ref(),
+ Arc::new(store),
+ );
+
+ println!("\n=== Showing all data ===");
+ let batches = ctx.sql("SELECT * FROM structs").await?.collect().await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ println!("\n=== Running query with struct field access and filter < 30
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("SELECT * FROM structs WHERE user_info['age'] > 30")
+ .await?
+ .collect()
+ .await?;
+
+ #[rustfmt::skip]
+ let expected = vec![
+ "+--------------------------+",
+ "| user_info |",
+ "+--------------------------+",
+ "| {name: Charlie, age: 35} |",
+ "+--------------------------+",
+ ];
+ assert_batches_eq!(expected, &batches);
+
+ println!("\n=== Running query with struct field access and filter > 130
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("explain analyze SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30")
+ .await?
+ .collect()
+ .await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ Ok(())
+}
+
+/// Create the example data with both struct fields and flattened fields
+fn create_sample_data() -> (SchemaRef, RecordBatch) {
+ // Create a schema with a struct column
+ let user_info_fields = Fields::from(vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("age", DataType::Int32, false),
+ ]);
+
+ let file_schema = Schema::new(vec![
+ Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ ),
+ // Include flattened fields (in real scenarios these might be in some
files but not others)
+ Field::new("_user_info.name", DataType::Utf8, true),
+ Field::new("_user_info.age", DataType::Int32, true),
+ ]);
+
+ let table_schema = Schema::new(vec![Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ )]);
+
+ // Create struct array for user_info
+ let names = StringArray::from(vec!["Alice", "Bob", "Charlie", "Dave"]);
+ let ages = Int32Array::from(vec![30, 25, 35, 22]);
+
+ let user_info = StructArray::from(vec![
+ (
+ Arc::new(Field::new("name", DataType::Utf8, false)),
+ Arc::new(names.clone()) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("age", DataType::Int32, false)),
+ Arc::new(ages.clone()) as ArrayRef,
+ ),
+ ]);
+
+ // Create a record batch with the data
+ let batch = RecordBatch::try_new(
+ Arc::new(file_schema.clone()),
+ vec![
+ Arc::new(user_info),
+ Arc::new(names), // Shredded name field
+ Arc::new(ages), // Shredded age field
+ ],
+ )
+ .unwrap();
+
+ (Arc::new(table_schema), batch)
+}
+
+/// Custom TableProvider that uses a StructFieldRewriter
+#[derive(Debug)]
+struct ExampleTableProvider {
+ schema: SchemaRef,
+}
+
+impl ExampleTableProvider {
+ fn new(schema: SchemaRef) -> Self {
+ Self { schema }
+ }
+}
+
+#[async_trait]
+impl TableProvider for ExampleTableProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ // Implementers can choose to mark these filters as exact or inexact.
+ // If marked as exact they cannot have false positives and must always
be applied.
+ // If marked as Inexact they can have false positives and at runtime
the rewriter
+ // can decide to not rewrite / ignore some filters since they will be
re-evaluated upstream.
+ Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ filters: &[Expr],
+ limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = self.schema.clone();
+ let df_schema = DFSchema::try_from(schema.clone())?;
+ let filter = state.create_physical_expr(
+ conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)),
+ &df_schema,
+ )?;
+
+ let parquet_source = ParquetSource::default()
+ .with_predicate(self.schema.clone(), filter)
+ .with_pushdown_filters(true)
+ .with_filter_expression_rewriter(Arc::new(StructFieldRewriter) as
_);
Review Comment:
This is the API for users to attach this rewriter to their plan
##########
datafusion/datasource-parquet/src/opener.rs:
##########
@@ -164,6 +162,53 @@ impl FileOpener for ParquetOpener {
adapted_projections.iter().cloned(),
);
+ let mut predicate = predicate;
+ let mut pruning_predicate = None;
+ let mut page_pruning_predicate = None;
+
+ // Filter pushdown: evaluate predicates during scan
+ if let Some(original_predicate) = &predicate {
+ // First check if any filter rewriters are available
+ let mut rewritten_predicate = original_predicate.clone();
+ if let Some(filter_rewriter) =
filter_expression_rewriter_factory {
+ // Try to rewrite the predicate using our filter
expression rewriter
+ // This MUST happen BEFORE the pushdown checking in
build_row_filter
+ // so that struct field accesses can be rewritten to flat
column references
+ if let Ok(rewritten) = filter_rewriter.rewrite(
+ Arc::clone(&table_schema),
+ Arc::clone(&file_schema),
+ original_predicate.clone(),
+ ) {
+ predicate = Some(rewritten.clone());
Review Comment:
This is a big gross and has several (cheap) clones. Would like to rework.
##########
datafusion/datasource-parquet/src/source.rs:
##########
@@ -295,46 +296,16 @@ impl ParquetSource {
self
}
- fn with_metrics(&self, metrics: ExecutionPlanMetricsSet) -> Self {
- let mut conf = self.clone();
- conf.metrics = metrics;
- conf
- }
-
/// Set predicate information, also sets pruning_predicate and
page_pruning_predicate attributes
pub fn with_predicate(
&self,
- file_schema: Arc<Schema>,
+ _file_schema: Arc<Schema>,
predicate: Arc<dyn PhysicalExpr>,
) -> Self {
let mut conf = self.clone();
- let metrics = ExecutionPlanMetricsSet::new();
- let predicate_creation_errors =
-
MetricBuilder::new(&metrics).global_counter("num_predicate_creation_errors");
-
- conf.with_metrics(metrics);
conf.predicate = Some(Arc::clone(&predicate));
- match PruningPredicate::try_new(Arc::clone(&predicate),
Arc::clone(&file_schema))
- {
- Ok(pruning_predicate) => {
- if !pruning_predicate.always_true() {
- conf.pruning_predicate = Some(Arc::new(pruning_predicate));
- }
- }
- Err(e) => {
- debug!("Could not create pruning predicate for: {e}");
- predicate_creation_errors.add(1);
- }
- };
-
- let page_pruning_predicate = Arc::new(PagePruningAccessPlanFilter::new(
- &predicate,
- Arc::clone(&file_schema),
- ));
- conf.page_pruning_predicate = Some(page_pruning_predicate);
Review Comment:
We could continue to generate these here and then just not use them later,
but that seems worse? If we do remove them maybe we should completely do so
(remove the methods, remove the struct fields, make it a breaking change).
##########
datafusion/datasource-parquet/src/source.rs:
##########
@@ -559,24 +556,8 @@ impl FileSource for ParquetSource {
.predicate()
.map(|p| format!(", predicate={p}"))
.unwrap_or_default();
- let pruning_predicate_string = self
- .pruning_predicate()
- .map(|pre| {
- let mut guarantees = pre
- .literal_guarantees()
- .iter()
- .map(|item| format!("{}", item))
- .collect_vec();
- guarantees.sort();
- format!(
- ", pruning_predicate={}, required_guarantees=[{}]",
- pre.predicate_expr(),
- guarantees.join(", ")
- )
- })
- .unwrap_or_default();
- write!(f, "{}{}", predicate_string, pruning_predicate_string)
Review Comment:
This is what's causing tests to fail. Tests assert against the formatting of
a ParquetSource and accessing it's pruning predicate method.
##########
datafusion-examples/examples/struct_field_rewrite.rs:
##########
@@ -0,0 +1,353 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray,
StructArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::Fields;
+use async_trait::async_trait;
+
+use datafusion::assert_batches_eq;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRewriter,
+};
+use datafusion::common::{DFSchema, Result};
+use datafusion::datasource::file_expr_rewriter::FileExpressionRewriter;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource};
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use datafusion::logical_expr::utils::conjunction;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
+use datafusion::parquet::arrow::ArrowWriter;
+use datafusion::parquet::file::properties::WriterProperties;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::{expressions, ScalarFunctionExpr};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::lit;
+use futures::StreamExt;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::{ObjectStore, PutPayload};
+
+// Example showing how to implement custom filter rewriting for struct fields.
+//
+// In this example, we have a table with a struct column like:
+// struct_col: {"a": 1, "b": "foo"}
+//
+// Our custom TableProvider will use a FilterExpressionRewriter to rewrite
+// expressions like `struct_col['a'] = 10` to use a flattened column name
+// `_struct_col.a` if it exists in the file schema.
+#[tokio::main]
+async fn main() -> Result<()> {
+ println!("=== Creating example data with structs and flattened fields
===");
+ println!("NOTE: This example demonstrates filter expression rewriting for
struct field access.");
+ println!(" We deliberately create different values in the struct
field vs. flattened column");
+ println!(" to prove that the rewriting system is actually using the
flattened column.");
+
+ // Create sample data with both struct columns and flattened fields
+ let (table_schema, batch) = create_sample_data();
+
+ let store = InMemory::new();
+ let buf = {
+ let mut buf = vec![];
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(1)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props))
+ .expect("creating writer");
+
+ writer.write(&batch).expect("Writing batch");
+ writer.close().unwrap();
+ buf
+ };
+ let path = Path::from("example.parquet");
+ let payload = PutPayload::from_bytes(buf.into());
+ store.put(&path, payload).await?;
+
+ // Create a custom table provider that rewrites struct field access
+ let table_provider = Arc::new(ExampleTableProvider::new(table_schema));
+
+ // Set up query execution
+ let ctx = SessionContext::new();
+
+ // Register our table
+ ctx.register_table("structs", table_provider)?;
+
+ ctx.runtime_env().register_object_store(
+ ObjectStoreUrl::parse("memory://")?.as_ref(),
+ Arc::new(store),
+ );
+
+ println!("\n=== Showing all data ===");
+ let batches = ctx.sql("SELECT * FROM structs").await?.collect().await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ println!("\n=== Running query with struct field access and filter < 30
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("SELECT * FROM structs WHERE user_info['age'] > 30")
+ .await?
+ .collect()
+ .await?;
+
+ #[rustfmt::skip]
+ let expected = vec![
+ "+--------------------------+",
+ "| user_info |",
+ "+--------------------------+",
+ "| {name: Charlie, age: 35} |",
+ "+--------------------------+",
+ ];
+ assert_batches_eq!(expected, &batches);
+
+ println!("\n=== Running query with struct field access and filter > 130
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("explain analyze SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30")
+ .await?
+ .collect()
+ .await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ Ok(())
+}
+
+/// Create the example data with both struct fields and flattened fields
+fn create_sample_data() -> (SchemaRef, RecordBatch) {
+ // Create a schema with a struct column
+ let user_info_fields = Fields::from(vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("age", DataType::Int32, false),
+ ]);
+
+ let file_schema = Schema::new(vec![
+ Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ ),
+ // Include flattened fields (in real scenarios these might be in some
files but not others)
+ Field::new("_user_info.name", DataType::Utf8, true),
+ Field::new("_user_info.age", DataType::Int32, true),
+ ]);
+
+ let table_schema = Schema::new(vec![Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ )]);
+
+ // Create struct array for user_info
+ let names = StringArray::from(vec!["Alice", "Bob", "Charlie", "Dave"]);
+ let ages = Int32Array::from(vec![30, 25, 35, 22]);
+
+ let user_info = StructArray::from(vec![
+ (
+ Arc::new(Field::new("name", DataType::Utf8, false)),
+ Arc::new(names.clone()) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("age", DataType::Int32, false)),
+ Arc::new(ages.clone()) as ArrayRef,
+ ),
+ ]);
+
+ // Create a record batch with the data
+ let batch = RecordBatch::try_new(
+ Arc::new(file_schema.clone()),
+ vec![
+ Arc::new(user_info),
+ Arc::new(names), // Shredded name field
+ Arc::new(ages), // Shredded age field
+ ],
+ )
+ .unwrap();
+
+ (Arc::new(table_schema), batch)
+}
+
+/// Custom TableProvider that uses a StructFieldRewriter
+#[derive(Debug)]
+struct ExampleTableProvider {
+ schema: SchemaRef,
+}
+
+impl ExampleTableProvider {
+ fn new(schema: SchemaRef) -> Self {
+ Self { schema }
+ }
+}
+
+#[async_trait]
+impl TableProvider for ExampleTableProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ // Implementers can choose to mark these filters as exact or inexact.
+ // If marked as exact they cannot have false positives and must always
be applied.
+ // If marked as Inexact they can have false positives and at runtime
the rewriter
+ // can decide to not rewrite / ignore some filters since they will be
re-evaluated upstream.
+ Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ filters: &[Expr],
+ limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = self.schema.clone();
+ let df_schema = DFSchema::try_from(schema.clone())?;
+ let filter = state.create_physical_expr(
+ conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)),
+ &df_schema,
+ )?;
+
+ let parquet_source = ParquetSource::default()
+ .with_predicate(self.schema.clone(), filter)
+ .with_pushdown_filters(true)
+ .with_filter_expression_rewriter(Arc::new(StructFieldRewriter) as
_);
+
+ let object_store_url = ObjectStoreUrl::parse("memory://")?;
+
+ let store = state.runtime_env().object_store(object_store_url)?;
+
+ let mut files = vec![];
+ let mut listing = store.list(None);
+ while let Some(file) = listing.next().await {
+ if let Ok(file) = file {
+ files.push(file);
+ }
+ }
+
+ let file_group = files
+ .iter()
+ .map(|file| {
+ PartitionedFile::new(
+ file.location.clone(),
+ u64::try_from(file.size.clone()).expect("fits in a u64"),
+ )
+ })
+ .collect();
+
+ let file_scan_config = FileScanConfig::new(
+ ObjectStoreUrl::parse("memory://")?,
+ schema,
+ Arc::new(parquet_source),
+ )
+ .with_projection(projection.cloned())
+ .with_limit(limit)
+ .with_file_group(file_group);
+
+ Ok(file_scan_config.build())
+ }
+}
+
+/// Rewriter that converts struct field access to flattened column references
+#[derive(Debug)]
+struct StructFieldRewriter;
+
+impl FileExpressionRewriter for StructFieldRewriter {
+ fn rewrite(
+ &self,
+ _table_schema: SchemaRef,
+ file_schema: SchemaRef,
+ expr: Arc<dyn PhysicalExpr>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let mut rewrite = StructFieldRewriterImpl { file_schema };
+ expr.rewrite(&mut rewrite).data()
+ }
+}
+
+struct StructFieldRewriterImpl {
+ file_schema: SchemaRef,
+}
+
+impl TreeNodeRewriter for StructFieldRewriterImpl {
+ type Node = Arc<dyn PhysicalExpr>;
+
+ fn f_down(
+ &mut self,
+ expr: Arc<dyn PhysicalExpr>,
+ ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ if let Some(scalar_function) =
expr.as_any().downcast_ref::<ScalarFunctionExpr>()
+ {
+ if scalar_function.name() == "get_field" {
+ if scalar_function.args().len() == 2 {
+ // First argument is the column, second argument is the
field name
+ let column = scalar_function.args()[0].clone();
+ let field_name = scalar_function.args()[1].clone();
+ if let Some(literal) =
+
field_name.as_any().downcast_ref::<expressions::Literal>()
+ {
+ if let Some(field_name) =
literal.value().try_as_str().flatten() {
+ if let Some(column) =
+
column.as_any().downcast_ref::<expressions::Column>()
+ {
+ let column_name = column.name();
+ let source_field =
+
self.file_schema.field_with_name(column_name)?;
+ let expected_flattened_column_name =
+ format!("_{}.{}", column_name, field_name);
+ // Check if the flattened column exists in the
file schema and has the same type
+ if let Ok(shredded_field) = self
+ .file_schema
+
.field_with_name(&expected_flattened_column_name)
+ {
+ if source_field.data_type()
+ == shredded_field.data_type()
+ {
+ // Rewrite the expression to use the
flattened column
+ let rewritten_expr = expressions::col(
+ &expected_flattened_column_name,
+ &self.file_schema,
+ )?;
+ return
Ok(Transformed::yes(rewritten_expr));
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Ok(Transformed::no(expr))
+ }
+}
Review Comment:
Example implementation of a rewriter
##########
datafusion-examples/examples/struct_field_rewrite.rs:
##########
@@ -0,0 +1,353 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray,
StructArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::Fields;
+use async_trait::async_trait;
+
+use datafusion::assert_batches_eq;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRewriter,
+};
+use datafusion::common::{DFSchema, Result};
+use datafusion::datasource::file_expr_rewriter::FileExpressionRewriter;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource};
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use datafusion::logical_expr::utils::conjunction;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
+use datafusion::parquet::arrow::ArrowWriter;
+use datafusion::parquet::file::properties::WriterProperties;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::{expressions, ScalarFunctionExpr};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::lit;
+use futures::StreamExt;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::{ObjectStore, PutPayload};
+
+// Example showing how to implement custom filter rewriting for struct fields.
+//
+// In this example, we have a table with a struct column like:
+// struct_col: {"a": 1, "b": "foo"}
+//
+// Our custom TableProvider will use a FilterExpressionRewriter to rewrite
+// expressions like `struct_col['a'] = 10` to use a flattened column name
+// `_struct_col.a` if it exists in the file schema.
+#[tokio::main]
+async fn main() -> Result<()> {
+ println!("=== Creating example data with structs and flattened fields
===");
+ println!("NOTE: This example demonstrates filter expression rewriting for
struct field access.");
+ println!(" We deliberately create different values in the struct
field vs. flattened column");
+ println!(" to prove that the rewriting system is actually using the
flattened column.");
+
+ // Create sample data with both struct columns and flattened fields
+ let (table_schema, batch) = create_sample_data();
+
+ let store = InMemory::new();
+ let buf = {
+ let mut buf = vec![];
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(1)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props))
+ .expect("creating writer");
+
+ writer.write(&batch).expect("Writing batch");
+ writer.close().unwrap();
+ buf
+ };
+ let path = Path::from("example.parquet");
+ let payload = PutPayload::from_bytes(buf.into());
+ store.put(&path, payload).await?;
+
+ // Create a custom table provider that rewrites struct field access
+ let table_provider = Arc::new(ExampleTableProvider::new(table_schema));
+
+ // Set up query execution
+ let ctx = SessionContext::new();
+
+ // Register our table
+ ctx.register_table("structs", table_provider)?;
+
+ ctx.runtime_env().register_object_store(
+ ObjectStoreUrl::parse("memory://")?.as_ref(),
+ Arc::new(store),
+ );
+
+ println!("\n=== Showing all data ===");
+ let batches = ctx.sql("SELECT * FROM structs").await?.collect().await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ println!("\n=== Running query with struct field access and filter < 30
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("SELECT * FROM structs WHERE user_info['age'] > 30")
+ .await?
+ .collect()
+ .await?;
+
+ #[rustfmt::skip]
+ let expected = vec![
+ "+--------------------------+",
+ "| user_info |",
+ "+--------------------------+",
+ "| {name: Charlie, age: 35} |",
+ "+--------------------------+",
+ ];
+ assert_batches_eq!(expected, &batches);
+
+ println!("\n=== Running query with struct field access and filter > 130
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("explain analyze SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30")
+ .await?
+ .collect()
+ .await?;
+ arrow::util::pretty::print_batches(&batches)?;
Review Comment:
Need a better assertion here but the point is that it shows
`row_groups_pruned_statistics=2` (without the rewrite it's
`row_groups_pruned_statistics=0`).
##########
datafusion-examples/examples/struct_field_rewrite.rs:
##########
@@ -0,0 +1,353 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray,
StructArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::Fields;
+use async_trait::async_trait;
+
+use datafusion::assert_batches_eq;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRewriter,
+};
+use datafusion::common::{DFSchema, Result};
+use datafusion::datasource::file_expr_rewriter::FileExpressionRewriter;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource};
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::object_store::ObjectStoreUrl;
+use datafusion::logical_expr::utils::conjunction;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
+use datafusion::parquet::arrow::ArrowWriter;
+use datafusion::parquet::file::properties::WriterProperties;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::{expressions, ScalarFunctionExpr};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::lit;
+use futures::StreamExt;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::{ObjectStore, PutPayload};
+
+// Example showing how to implement custom filter rewriting for struct fields.
+//
+// In this example, we have a table with a struct column like:
+// struct_col: {"a": 1, "b": "foo"}
+//
+// Our custom TableProvider will use a FilterExpressionRewriter to rewrite
+// expressions like `struct_col['a'] = 10` to use a flattened column name
+// `_struct_col.a` if it exists in the file schema.
+#[tokio::main]
+async fn main() -> Result<()> {
+ println!("=== Creating example data with structs and flattened fields
===");
+ println!("NOTE: This example demonstrates filter expression rewriting for
struct field access.");
+ println!(" We deliberately create different values in the struct
field vs. flattened column");
+ println!(" to prove that the rewriting system is actually using the
flattened column.");
+
+ // Create sample data with both struct columns and flattened fields
+ let (table_schema, batch) = create_sample_data();
+
+ let store = InMemory::new();
+ let buf = {
+ let mut buf = vec![];
+
+ let props = WriterProperties::builder()
+ .set_max_row_group_size(1)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props))
+ .expect("creating writer");
+
+ writer.write(&batch).expect("Writing batch");
+ writer.close().unwrap();
+ buf
+ };
+ let path = Path::from("example.parquet");
+ let payload = PutPayload::from_bytes(buf.into());
+ store.put(&path, payload).await?;
+
+ // Create a custom table provider that rewrites struct field access
+ let table_provider = Arc::new(ExampleTableProvider::new(table_schema));
+
+ // Set up query execution
+ let ctx = SessionContext::new();
+
+ // Register our table
+ ctx.register_table("structs", table_provider)?;
+
+ ctx.runtime_env().register_object_store(
+ ObjectStoreUrl::parse("memory://")?.as_ref(),
+ Arc::new(store),
+ );
+
+ println!("\n=== Showing all data ===");
+ let batches = ctx.sql("SELECT * FROM structs").await?.collect().await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ println!("\n=== Running query with struct field access and filter < 30
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("SELECT * FROM structs WHERE user_info['age'] > 30")
+ .await?
+ .collect()
+ .await?;
+
+ #[rustfmt::skip]
+ let expected = vec![
+ "+--------------------------+",
+ "| user_info |",
+ "+--------------------------+",
+ "| {name: Charlie, age: 35} |",
+ "+--------------------------+",
+ ];
+ assert_batches_eq!(expected, &batches);
+
+ println!("\n=== Running query with struct field access and filter > 130
===");
+ println!("Query: SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30");
+
+ let batches = ctx
+ .sql("explain analyze SELECT user_info['name'] FROM structs WHERE
user_info['age'] < 30")
+ .await?
+ .collect()
+ .await?;
+ arrow::util::pretty::print_batches(&batches)?;
+
+ Ok(())
+}
+
+/// Create the example data with both struct fields and flattened fields
+fn create_sample_data() -> (SchemaRef, RecordBatch) {
+ // Create a schema with a struct column
+ let user_info_fields = Fields::from(vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("age", DataType::Int32, false),
+ ]);
+
+ let file_schema = Schema::new(vec![
+ Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ ),
+ // Include flattened fields (in real scenarios these might be in some
files but not others)
+ Field::new("_user_info.name", DataType::Utf8, true),
+ Field::new("_user_info.age", DataType::Int32, true),
+ ]);
+
+ let table_schema = Schema::new(vec![Field::new(
+ "user_info",
+ DataType::Struct(user_info_fields.clone()),
+ false,
+ )]);
+
+ // Create struct array for user_info
+ let names = StringArray::from(vec!["Alice", "Bob", "Charlie", "Dave"]);
+ let ages = Int32Array::from(vec![30, 25, 35, 22]);
+
+ let user_info = StructArray::from(vec![
+ (
+ Arc::new(Field::new("name", DataType::Utf8, false)),
+ Arc::new(names.clone()) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("age", DataType::Int32, false)),
+ Arc::new(ages.clone()) as ArrayRef,
+ ),
+ ]);
+
+ // Create a record batch with the data
+ let batch = RecordBatch::try_new(
+ Arc::new(file_schema.clone()),
+ vec![
+ Arc::new(user_info),
+ Arc::new(names), // Shredded name field
+ Arc::new(ages), // Shredded age field
+ ],
+ )
+ .unwrap();
+
+ (Arc::new(table_schema), batch)
+}
+
+/// Custom TableProvider that uses a StructFieldRewriter
+#[derive(Debug)]
+struct ExampleTableProvider {
+ schema: SchemaRef,
+}
+
+impl ExampleTableProvider {
+ fn new(schema: SchemaRef) -> Self {
+ Self { schema }
+ }
+}
+
+#[async_trait]
+impl TableProvider for ExampleTableProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ // Implementers can choose to mark these filters as exact or inexact.
+ // If marked as exact they cannot have false positives and must always
be applied.
+ // If marked as Inexact they can have false positives and at runtime
the rewriter
+ // can decide to not rewrite / ignore some filters since they will be
re-evaluated upstream.
+ Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ filters: &[Expr],
+ limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = self.schema.clone();
+ let df_schema = DFSchema::try_from(schema.clone())?;
+ let filter = state.create_physical_expr(
+ conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)),
+ &df_schema,
+ )?;
+
+ let parquet_source = ParquetSource::default()
+ .with_predicate(self.schema.clone(), filter)
+ .with_pushdown_filters(true)
+ .with_filter_expression_rewriter(Arc::new(StructFieldRewriter) as
_);
+
+ let object_store_url = ObjectStoreUrl::parse("memory://")?;
+
+ let store = state.runtime_env().object_store(object_store_url)?;
+
+ let mut files = vec![];
+ let mut listing = store.list(None);
+ while let Some(file) = listing.next().await {
+ if let Ok(file) = file {
+ files.push(file);
+ }
+ }
+
+ let file_group = files
+ .iter()
+ .map(|file| {
+ PartitionedFile::new(
+ file.location.clone(),
+ u64::try_from(file.size.clone()).expect("fits in a u64"),
+ )
+ })
+ .collect();
+
+ let file_scan_config = FileScanConfig::new(
+ ObjectStoreUrl::parse("memory://")?,
+ schema,
+ Arc::new(parquet_source),
+ )
+ .with_projection(projection.cloned())
+ .with_limit(limit)
+ .with_file_group(file_group);
+
+ Ok(file_scan_config.build())
+ }
+}
+
+/// Rewriter that converts struct field access to flattened column references
+#[derive(Debug)]
+struct StructFieldRewriter;
+
+impl FileExpressionRewriter for StructFieldRewriter {
+ fn rewrite(
+ &self,
+ _table_schema: SchemaRef,
+ file_schema: SchemaRef,
+ expr: Arc<dyn PhysicalExpr>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let mut rewrite = StructFieldRewriterImpl { file_schema };
+ expr.rewrite(&mut rewrite).data()
+ }
+}
+
+struct StructFieldRewriterImpl {
+ file_schema: SchemaRef,
+}
+
+impl TreeNodeRewriter for StructFieldRewriterImpl {
+ type Node = Arc<dyn PhysicalExpr>;
+
+ fn f_down(
+ &mut self,
+ expr: Arc<dyn PhysicalExpr>,
+ ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ if let Some(scalar_function) =
expr.as_any().downcast_ref::<ScalarFunctionExpr>()
+ {
+ if scalar_function.name() == "get_field" {
+ if scalar_function.args().len() == 2 {
+ // First argument is the column, second argument is the
field name
+ let column = scalar_function.args()[0].clone();
+ let field_name = scalar_function.args()[1].clone();
+ if let Some(literal) =
+
field_name.as_any().downcast_ref::<expressions::Literal>()
+ {
+ if let Some(field_name) =
literal.value().try_as_str().flatten() {
+ if let Some(column) =
+
column.as_any().downcast_ref::<expressions::Column>()
+ {
+ let column_name = column.name();
+ let source_field =
+
self.file_schema.field_with_name(column_name)?;
+ let expected_flattened_column_name =
+ format!("_{}.{}", column_name, field_name);
+ // Check if the flattened column exists in the
file schema and has the same type
+ if let Ok(shredded_field) = self
+ .file_schema
+
.field_with_name(&expected_flattened_column_name)
+ {
+ if source_field.data_type()
+ == shredded_field.data_type()
+ {
+ // Rewrite the expression to use the
flattened column
+ let rewritten_expr = expressions::col(
+ &expected_flattened_column_name,
+ &self.file_schema,
+ )?;
+ return
Ok(Transformed::yes(rewritten_expr));
Review Comment:
This is completely arbitrary: this could return another expression, it
doesn't _have_ to be just a column.
--
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]
