From ff0f7a50afafef98a47b7c902af5b536249a433e Mon Sep 17 00:00:00 2001 From: Kadin Rabo Date: Thu, 23 Jul 2026 13:18:04 +0000 Subject: [PATCH 1/2] fix(substrait): handle read filters and projections --- .../consumer/expr/field_reference.rs | 119 +- .../src/logical_plan/consumer/rel/read_rel.rs | 1787 ++++++++++++++++- .../src/logical_plan/consumer/utils.rs | 2 +- 3 files changed, 1795 insertions(+), 113 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs index be084f360358a..16551c3b4622a 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs @@ -16,7 +16,9 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use datafusion::common::{Column, DFSchema, not_impl_err, substrait_err}; +use datafusion::common::{ + Column, DFSchema, not_impl_err, substrait_datafusion_err, substrait_err, +}; use datafusion::logical_expr::Expr; use std::sync::Arc; use substrait::proto::expression::FieldReference; @@ -45,11 +47,20 @@ pub(crate) fn from_substrait_field_reference( "Direct reference StructField with child is not supported" ); } - let field_idx = struct_field.field as usize; + let field_idx = field_reference_index(struct_field.field)?; match &field_ref.root_type { - Some(RootType::RootReference(_)) | None => Ok(Expr::Column( - Column::from(input_schema.qualified_field(field_idx)), - )), + Some(RootType::RootReference(_)) | None => { + if field_idx >= input_schema.fields().len() { + return substrait_err!( + "Field reference index {} is out of bounds for input schema with {} fields", + field_idx, + input_schema.fields().len() + ); + } + Ok(Expr::Column(Column::from( + input_schema.qualified_field(field_idx), + ))) + } Some(RootType::OuterReference(outer_ref)) => { resolve_outer_reference(consumer, outer_ref, field_idx) } @@ -69,6 +80,11 @@ pub(crate) fn from_substrait_field_reference( } } +fn field_reference_index(field: i32) -> datafusion::common::Result { + usize::try_from(field) + .map_err(|_| substrait_datafusion_err!("Invalid field reference index: {field}")) +} + fn resolve_outer_reference( consumer: &impl SubstraitConsumer, outer_ref: &substrait::proto::expression::field_reference::OuterReference, @@ -81,6 +97,13 @@ fn resolve_outer_reference( but no outer schema is available" ); }; + if field_idx >= outer_schema.fields().len() { + return substrait_err!( + "OuterReference field index {} is out of bounds for outer schema with {} fields", + field_idx, + outer_schema.fields().len() + ); + } let (qualifier, field) = outer_schema.qualified_field(field_idx); let col = Column::from((qualifier, field)); Ok(Expr::OuterReferenceColumn(Arc::clone(field), col)) @@ -88,7 +111,10 @@ fn resolve_outer_reference( #[cfg(test)] mod tests { + use std::sync::Arc; + use datafusion::{ + arrow::datatypes::{DataType, Field}, common::{DFSchema, assert_contains}, prelude::SessionContext, }; @@ -152,6 +178,61 @@ mod tests { ); } + #[tokio::test] + async fn test_root_reference_invalid_field_idx() { + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + let input_schema = DFSchema::new_with_metadata( + vec![(None, Arc::new(Field::new("a", DataType::Int64, true)))], + Default::default(), + ) + .unwrap(); + + for (field, expected) in [ + (-1, "Invalid field reference index: -1"), + ( + 1, + "Field reference index 1 is out of bounds for input schema with 1 fields", + ), + ] { + let err = + from_field_reference(&consumer, &root_field_ref(field), &input_schema) + .await + .unwrap_err(); + + assert_contains!(err.to_string(), expected); + } + } + + #[tokio::test] + async fn test_outer_reference_invalid_field_idx() { + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + let outer_schema = Arc::new( + DFSchema::new_with_metadata( + vec![(None, Arc::new(Field::new("a", DataType::Int64, true)))], + Default::default(), + ) + .unwrap(), + ); + consumer.push_outer_schema(outer_schema); + + let err = from_field_reference( + &consumer, + &outer_field_ref(1, 1), + DFSchema::empty_ref(), + ) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "OuterReference field index 1 is out of bounds for outer schema with 1 fields" + ); + } + fn lambda_field_ref(field: i32, steps_out: u32) -> FieldReference { FieldReference { reference_type: Some(field_reference::ReferenceType::DirectReference( @@ -166,4 +247,32 @@ mod tests { )), } } + + fn root_field_ref(field: i32) -> FieldReference { + FieldReference { + reference_type: Some(field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some(ReferenceType::StructField(Box::new( + StructField { field, child: None }, + ))), + }, + )), + root_type: Some(RootType::RootReference(field_reference::RootReference {})), + } + } + + fn outer_field_ref(field: i32, steps_out: u32) -> FieldReference { + FieldReference { + reference_type: Some(field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some(ReferenceType::StructField(Box::new( + StructField { field, child: None }, + ))), + }, + )), + root_type: Some(RootType::OuterReference(field_reference::OuterReference { + steps_out, + })), + } + } } diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 832110e11131c..2a23beb6aeb85 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -18,15 +18,18 @@ use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::from_substrait_literal; use crate::logical_plan::consumer::from_substrait_named_struct; -use crate::logical_plan::consumer::utils::ensure_schema_compatibility; +use crate::logical_plan::consumer::utils::{ + ensure_field_compatibility, ensure_schema_compatibility, rename_expressions, +}; +use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::common::{ - DFSchema, DFSchemaRef, TableReference, not_impl_err, plan_err, - substrait_datafusion_err, substrait_err, + Column, DFSchema, DFSchemaRef, TableReference, not_impl_err, plan_datafusion_err, + plan_err, substrait_datafusion_err, substrait_err, }; use datafusion::datasource::provider_as_source; use datafusion::logical_expr::utils::split_conjunction_owned; use datafusion::logical_expr::{ - EmptyRelation, Expr, LogicalPlan, LogicalPlanBuilder, Values, + Cast, EmptyRelation, Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder, Values, }; use std::sync::Arc; use substrait::proto::expression::MaskExpression; @@ -56,26 +59,61 @@ pub async fn from_read_rel( vec![] }; - let plan = { - let provider = match consumer.resolve_table_ref(&table_ref).await? { - Some(ref provider) => Arc::clone(provider), - _ => return plan_err!("No table named '{table_ref}'"), - }; - - LogicalPlanBuilder::scan_with_filters( - table_ref, - provider_as_source(Arc::clone(&provider)), - None, - filters, - )? - .build()? + let provider = match consumer.resolve_table_ref(&table_ref).await? { + Some(ref provider) => Arc::clone(provider), + _ => return plan_err!("No table named '{table_ref}'"), }; + let unqualified_schema = schema.clone().strip_qualifiers(); + let scan_projection = if arrow_schema_matches_by_position( + provider.schema().as_ref(), + &unqualified_schema, + ) { + ensure_arrow_schema_compatibility_by_position( + provider.schema().as_ref(), + &unqualified_schema, + )?; + projection_indices(&schema, projection)? + } else { + None + }; + let projected_schema = scan_projection + .is_some() + .then(|| apply_masking(schema.clone(), projection)) + .transpose()?; + + let plan = LogicalPlanBuilder::scan_with_filters( + table_ref, + provider_as_source(Arc::clone(&provider)), + scan_projection, + filters, + )? + .build()?; - ensure_schema_compatibility(plan.schema(), schema.clone())?; + if let Some(projected_schema) = projected_schema { + ensure_schema_compatibility_by_position(plan.schema(), &projected_schema)?; + return Ok(plan); + } - let schema = apply_masking(schema, projection)?; + let schema_matches_by_position = + plan.schema().logically_equivalent_names_and_types(&schema); + if schema_matches_by_position { + ensure_schema_compatibility_by_position(plan.schema(), &schema)?; + } else { + if schema_has_duplicate_field_names(&schema) { + return not_impl_err!( + "ReadRel schemas with duplicate field names must match by position" + ); + } + ensure_schema_compatibility(plan.schema(), schema.clone())?; + } - apply_projection(plan, schema) + apply_read_rel_projection( + plan, + schema, + projection, + schema_matches_by_position, + false, + ) } let named_struct = read.base_schema.as_ref().ok_or_else(|| { @@ -114,66 +152,77 @@ pub async fn from_read_rel( .await } Some(ReadType::VirtualTable(vt)) => { - if vt.values.is_empty() && vt.expressions.is_empty() { - return Ok(LogicalPlan::EmptyRelation(EmptyRelation { + let plan = if vt.values.is_empty() && vt.expressions.is_empty() { + LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, - schema: DFSchemaRef::new(substrait_schema), - })); - } - - // Check for produce_one_row pattern in both old (values) and new (expressions) formats. - // A VirtualTable with exactly one row containing only empty/default fields represents - // an EmptyRelation with produce_one_row=true. This pattern is used for queries without - // a FROM clause (e.g., "SELECT 1 AS one") where a single phantom row is needed to - // provide a context for evaluating scalar expressions. This is conceptually similar to - // the SQL "DUAL" table (see: https://en.wikipedia.org/wiki/DUAL_table) which some - // databases provide as a single-row source for selecting constant expressions when no - // real table is present. - let is_produce_one_row = (vt.values.len() == 1 - && vt.expressions.is_empty() - && substrait_schema.fields().is_empty() - && vt.values[0].fields.is_empty()) - || (vt.expressions.len() == 1 - && vt.values.is_empty() + schema: DFSchemaRef::new(substrait_schema.clone()), + }) + } else { + // Check for produce_one_row pattern in both old (values) and new (expressions) formats. + // A VirtualTable with exactly one row containing only empty/default fields represents + // an EmptyRelation with produce_one_row=true. This pattern is used for queries without + // a FROM clause (e.g., "SELECT 1 AS one") where a single phantom row is needed to + // provide a context for evaluating scalar expressions. This is conceptually similar to + // the SQL "DUAL" table (see: https://en.wikipedia.org/wiki/DUAL_table) which some + // databases provide as a single-row source for selecting constant expressions when no + // real table is present. + let is_produce_one_row = (vt.values.len() == 1 + && vt.expressions.is_empty() && substrait_schema.fields().is_empty() - && vt.expressions[0].fields.is_empty()); + && vt.values[0].fields.is_empty()) + || (vt.expressions.len() == 1 + && vt.values.is_empty() + && substrait_schema.fields().is_empty() + && vt.expressions[0].fields.is_empty()); - if is_produce_one_row { - return Ok(LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: true, - schema: DFSchemaRef::new(substrait_schema), - })); - } + if is_produce_one_row { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: true, + schema: DFSchemaRef::new(substrait_schema.clone()), + }) + } else { + let values = if !vt.expressions.is_empty() { + let mut exprs = vec![]; + for row in &vt.expressions { + let mut row_exprs = vec![]; + for expression in &row.fields { + let expr = consumer + .consume_expression(expression, &substrait_schema) + .await?; + row_exprs.push(expr); + } + // For expressions, validate against top-level schema fields, not nested names + if row_exprs.len() != substrait_schema.fields().len() { + return substrait_err!( + "Field count mismatch: expected {} fields but found {} in virtual table row", + substrait_schema.fields().len(), + row_exprs.len() + ); + } + exprs.push(row_exprs); + } + exprs + } else { + convert_literal_rows(consumer, vt, named_struct)? + }; - let values = if !vt.expressions.is_empty() { - let mut exprs = vec![]; - for row in &vt.expressions { - let mut row_exprs = vec![]; - for expression in &row.fields { - let expr = consumer - .consume_expression(expression, &substrait_schema) - .await?; - row_exprs.push(expr); - } - // For expressions, validate against top-level schema fields, not nested names - if row_exprs.len() != substrait_schema.fields().len() { - return substrait_err!( - "Field count mismatch: expected {} fields but found {} in virtual table row", - substrait_schema.fields().len(), - row_exprs.len() - ); - } - exprs.push(row_exprs); + LogicalPlan::Values(Values { + schema: DFSchemaRef::new(substrait_schema.clone()), + values, + }) } - exprs - } else { - convert_literal_rows(consumer, vt, named_struct)? }; - Ok(LogicalPlan::Values(Values { - schema: DFSchemaRef::new(substrait_schema), - values, - })) + let plan = normalize_values_to_schema(plan)?; + + apply_read_rel_filter_and_projection( + consumer, + plan, + substrait_schema, + &read.projection, + &read.filter, + ) + .await } Some(ReadType::LocalFiles(lf)) => { fn extract_filename(name: &str) -> Option { @@ -222,6 +271,184 @@ pub async fn from_read_rel( } } +async fn apply_read_rel_filter_and_projection( + consumer: &impl SubstraitConsumer, + plan: LogicalPlan, + substrait_schema: DFSchema, + projection: &Option, + filter: &Option>, +) -> datafusion::common::Result { + let schema_matches_by_position = plan + .schema() + .logically_equivalent_names_and_types(&substrait_schema); + + if schema_matches_by_position { + ensure_schema_compatibility_by_position(plan.schema(), &substrait_schema)?; + } else { + if schema_has_duplicate_field_names(&substrait_schema) { + return not_impl_err!( + "ReadRel schemas with duplicate field names must match by position" + ); + } + ensure_schema_compatibility(plan.schema(), substrait_schema.clone())?; + } + + let filter_uses_ordinal_aliases = filter.is_some() + && schema_matches_by_position + && schema_has_duplicate_field_names(&substrait_schema); + let (plan, filter_schema) = if filter_uses_ordinal_aliases { + let filter_schema = ordinal_read_filter_schema(&substrait_schema)?; + ( + apply_schema_by_position(plan, &filter_schema)?, + filter_schema, + ) + } else { + (plan, substrait_schema.clone()) + }; + + let plan = if let Some(f) = filter { + let filter_expr = consumer.consume_expression(f, &filter_schema).await?; + LogicalPlanBuilder::from(plan) + .filter(filter_expr)? + .build()? + } else { + plan + }; + + apply_read_rel_projection( + plan, + substrait_schema, + projection, + schema_matches_by_position, + filter_uses_ordinal_aliases, + ) +} + +fn apply_read_rel_projection( + plan: LogicalPlan, + substrait_schema: DFSchema, + projection: &Option, + schema_matches_by_position: bool, + force_aliases: bool, +) -> datafusion::common::Result { + let projection_indices = if schema_matches_by_position { + projection_indices(&substrait_schema, projection)? + } else { + None + }; + let schema = apply_masking(substrait_schema, projection)?; + if force_aliases && schema_has_duplicate_field_names(&schema) { + return not_impl_err!( + "ReadRel filters over duplicate field names require a projection with unique field names" + ); + } + + match projection_indices { + Some(column_indices) if force_aliases => { + apply_projection_with_indices_and_aliases(plan, &schema, column_indices) + } + Some(column_indices) => { + apply_projection_with_indices(plan, &schema, column_indices) + } + None if force_aliases => { + let column_indices = (0..schema.fields().len()).collect(); + apply_projection_with_indices_and_aliases(plan, &schema, column_indices) + } + None => apply_projection(plan, &schema), + } +} + +fn ensure_schema_compatibility_by_position( + table_schema: &DFSchema, + substrait_schema: &DFSchema, +) -> datafusion::common::Result<()> { + table_schema + .fields() + .iter() + .zip(substrait_schema.fields()) + .try_for_each(|(datafusion_field, substrait_field)| { + ensure_field_compatibility(datafusion_field, substrait_field) + }) +} + +fn ensure_arrow_schema_compatibility_by_position( + table_schema: &Schema, + substrait_schema: &DFSchema, +) -> datafusion::common::Result<()> { + if table_schema.fields().len() != substrait_schema.fields().len() { + return substrait_err!( + "ReadRel schema has {} fields but table schema has {} fields", + substrait_schema.fields().len(), + table_schema.fields().len() + ); + } + + table_schema + .fields() + .iter() + .zip(substrait_schema.fields()) + .try_for_each(|(datafusion_field, substrait_field)| { + ensure_field_compatibility(datafusion_field.as_ref(), substrait_field) + }) +} + +fn arrow_schema_matches_by_position( + table_schema: &Schema, + substrait_schema: &DFSchema, +) -> bool { + if table_schema.fields().len() != substrait_schema.fields().len() { + return false; + } + + table_schema + .fields() + .iter() + .zip(substrait_schema.fields()) + .all(|(table_field, substrait_field)| { + table_field.name() == substrait_field.name() + && DFSchema::datatype_is_logically_equal( + table_field.data_type(), + substrait_field.data_type(), + ) + }) +} + +fn ordinal_read_filter_schema(schema: &DFSchema) -> datafusion::common::Result { + let fields = schema + .iter() + .enumerate() + .map(|(index, (_qualifier, field))| { + ( + Some(TableReference::Bare { + table: format!("__datafusion_substrait_read_{index}").into(), + }), + Arc::clone(field), + ) + }) + .collect(); + + DFSchema::new_with_metadata(fields, schema.metadata().clone()) +} + +fn apply_schema_by_position( + plan: LogicalPlan, + schema: &DFSchema, +) -> datafusion::common::Result { + match plan { + LogicalPlan::Values(mut values) => { + values.schema = DFSchemaRef::new(schema.clone()); + Ok(LogicalPlan::Values(values)) + } + LogicalPlan::EmptyRelation(mut empty) => { + empty.schema = DFSchemaRef::new(schema.clone()); + Ok(LogicalPlan::EmptyRelation(empty)) + } + _ => not_impl_err!( + "ReadRel filters over generic plans with duplicate field names are not supported" + ), + } +} + /// Converts Substrait literal rows from a VirtualTable into DataFusion expressions. /// /// This function processes the deprecated `values` field of VirtualTable, converting @@ -262,33 +489,50 @@ fn convert_literal_rows( .collect::>() } +/// Applies a top-level `ReadRel` projection mask to a schema. +/// +/// Nested masks are not supported and return an error. The +/// `maintain_singular_struct` flag does not affect top-level field selection. +/// Projection field indexes must refer to existing top-level fields. pub fn apply_masking( schema: DFSchema, mask_expression: &::core::option::Option, ) -> datafusion::common::Result { + let Some(column_indices) = projection_indices(&schema, mask_expression)? else { + return Ok(schema); + }; + + let fields = column_indices + .iter() + .map(|i| schema.qualified_field(*i)) + .map(|(qualifier, field)| (qualifier.cloned(), Arc::clone(field))) + .collect(); + + DFSchema::new_with_metadata(fields, schema.metadata().clone()) +} + +fn projection_indices( + schema: &DFSchema, + mask_expression: &::core::option::Option, +) -> datafusion::common::Result>> { match mask_expression { Some(MaskExpression { select, .. }) => match &select.as_ref() { - Some(projection) => { - let column_indices: Vec = projection - .struct_items - .iter() - .map(|item| item.field as usize) - .collect(); - - let fields = column_indices - .iter() - .map(|i| schema.qualified_field(*i)) - .map(|(qualifier, field)| (qualifier.cloned(), Arc::clone(field))) - .collect(); - - Ok(DFSchema::new_with_metadata( - fields, - schema.metadata().clone(), - )?) - } - None => Ok(schema), + Some(projection) => projection + .struct_items + .iter() + .map(|item| { + if item.child.is_some() { + return not_impl_err!( + "Nested ReadRel projections are not supported" + ); + } + projection_field_index(schema, item.field) + }) + .collect::>() + .map(Some), + None => Ok(None), }, - None => Ok(schema), + None => Ok(None), } } @@ -296,43 +540,1372 @@ pub fn apply_masking( /// Substrait schema is a subset of the DataFusion schema. fn apply_projection( plan: LogicalPlan, - substrait_schema: DFSchema, + substrait_schema: &DFSchema, ) -> datafusion::common::Result { let df_schema = plan.schema(); - if df_schema.logically_equivalent_names_and_types(&substrait_schema) { + if df_schema.logically_equivalent_names_and_types(substrait_schema) { return Ok(plan); } let df_schema = df_schema.to_owned(); + let column_indices: Vec = substrait_schema + .clone() + .strip_qualifiers() + .fields() + .iter() + .map(|substrait_field| { + unqualified_field_index(&df_schema, substrait_field.name().as_str()) + }) + .collect::>()?; + + apply_projection_with_indices(plan, substrait_schema, column_indices) +} + +fn cast_value_expressions( + exprs: Vec, + input_schema: &DFSchema, + fields: &[Arc], +) -> datafusion::common::Result> { + if exprs.len() != fields.len() { + return plan_err!( + "ReadRel Values projection has {} expressions for {} fields", + exprs.len(), + fields.len() + ); + } + + exprs + .into_iter() + .zip(fields) + .map(|(expr, field)| { + if &expr.get_type(input_schema)? != field.data_type() { + Ok(Expr::Cast(Cast::new( + Box::new(expr), + field.data_type().to_owned(), + ))) + } else { + Ok(expr) + } + }) + .collect() +} + +fn normalize_values_to_schema( + plan: LogicalPlan, +) -> datafusion::common::Result { match plan { - LogicalPlan::TableScan(mut scan) => { - let column_indices: Vec = substrait_schema - .strip_qualifiers() - .fields() - .iter() - .map(|substrait_field| { - Ok(df_schema - .index_of_column_by_name(None, substrait_field.name().as_str()) - .unwrap()) - }) + LogicalPlan::Values(mut values) => { + let schema = Arc::clone(&values.schema); + values.values = values + .values + .into_iter() + .map(|exprs| cast_value_expressions(exprs, &schema, schema.fields())) .collect::>()?; + Ok(LogicalPlan::Values(values)) + } + _ => Ok(plan), + } +} + +fn apply_projection_with_indices( + plan: LogicalPlan, + substrait_schema: &DFSchema, + column_indices: Vec, +) -> datafusion::common::Result { + apply_projection_with_indices_inner(plan, substrait_schema, column_indices, false) +} + +fn apply_projection_with_indices_and_aliases( + plan: LogicalPlan, + substrait_schema: &DFSchema, + column_indices: Vec, +) -> datafusion::common::Result { + apply_projection_with_indices_inner(plan, substrait_schema, column_indices, true) +} + +fn apply_projection_with_indices_inner( + plan: LogicalPlan, + substrait_schema: &DFSchema, + column_indices: Vec, + force_aliases: bool, +) -> datafusion::common::Result { + let df_schema = plan.schema().to_owned(); + match plan { + LogicalPlan::TableScan(mut scan) => { let fields = column_indices .iter() .map(|i| df_schema.qualified_field(*i)) .map(|(qualifier, field)| (qualifier.cloned(), Arc::clone(field))) .collect(); + let scan_projection = match scan.projection.as_ref() { + Some(existing) => Some( + column_indices + .iter() + .map(|i| { + existing.get(*i).copied().ok_or_else(|| { + plan_datafusion_err!( + "ReadRel projection field index {i} is out of bounds for the TableScan projection" + ) + }) + }) + .collect::>()?, + ), + None => Some(column_indices), + }; + scan.projected_schema = DFSchemaRef::new(DFSchema::new_with_metadata( fields, df_schema.metadata().clone(), )?); - scan.projection = Some(column_indices); + scan.projection = scan_projection; Ok(LogicalPlan::TableScan(scan)) } - _ => plan_err!("DataFrame passed to apply_projection must be a TableScan"), + LogicalPlan::Values(mut values) => { + values.values = values + .values + .into_iter() + .map(|row| { + column_indices + .iter() + .map(|i| { + row.get(*i).cloned().ok_or_else(|| { + plan_datafusion_err!( + "ReadRel projection field index {i} is out of bounds for a Values row with {} fields", + row.len() + ) + }) + }) + .collect::>() + }) + .map(|exprs| { + cast_value_expressions(exprs?, &df_schema, substrait_schema.fields()) + }) + .collect::>()?; + values.schema = DFSchemaRef::new(substrait_schema.clone()); + + Ok(LogicalPlan::Values(values)) + } + LogicalPlan::EmptyRelation(mut empty) => { + empty.schema = DFSchemaRef::new(substrait_schema.clone()); + + Ok(LogicalPlan::EmptyRelation(empty)) + } + _ => { + if schema_has_duplicate_field_names(&df_schema) { + return not_impl_err!( + "ReadRel projections over generic plans with duplicate field names are not supported" + ); + } + + let exprs: Vec = column_indices + .iter() + .map(|i| Expr::Column(Column::from(df_schema.qualified_field(*i)))) + .collect(); + let exprs = rename_expressions(exprs, &df_schema, substrait_schema.fields())?; + let exprs = if force_aliases { + exprs + .into_iter() + .zip(substrait_schema.fields()) + .map(|(expr, field)| expr.alias(field.name().to_owned())) + .collect() + } else { + exprs + }; + + LogicalPlanBuilder::from(plan).project(exprs)?.build() + } + } +} + +fn projection_field_index( + schema: &DFSchema, + field: i32, +) -> datafusion::common::Result { + let index = usize::try_from(field).map_err(|_| { + substrait_datafusion_err!("Invalid ReadRel projection field index: {field}") + })?; + + if index >= schema.fields().len() { + return substrait_err!( + "ReadRel projection field index {} is out of bounds for schema with {} fields", + field, + schema.fields().len() + ); + } + + Ok(index) +} + +fn schema_has_duplicate_field_names(schema: &DFSchema) -> bool { + schema + .iter() + .enumerate() + .any(|(left_index, (left_qualifier, left_field))| { + schema + .iter() + .skip(left_index + 1) + .any(|(right_qualifier, right_field)| { + left_qualifier == right_qualifier + && left_field.name() == right_field.name() + }) + }) +} + +fn unqualified_field_index( + schema: &DFSchema, + name: &str, +) -> datafusion::common::Result { + let matches = schema + .iter() + .enumerate() + .filter(|(_, (_, field))| field.name() == name) + .map(|(index, (qualifier, _))| (index, qualifier)) + .collect::>(); + + match matches.len() { + 0 => Err(plan_datafusion_err!( + "No field named '{name}' found in ReadRel input schema" + )), + 1 => Ok(matches[0].0), + _ => { + let unqualified_matches = matches + .iter() + .filter(|(_, qualifier)| qualifier.is_none()) + .collect::>(); + + if unqualified_matches.len() == 1 { + Ok(unqualified_matches[0].0) + } else { + Err(plan_datafusion_err!( + "Ambiguous field name '{name}' found in ReadRel input schema" + )) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::extensions::Extensions; + use crate::logical_plan::consumer::utils::tests::{ + TEST_EXTENSIONS, TEST_SESSION_STATE, test_consumer, + }; + use crate::variation_const::DEFAULT_TYPE_VARIATION_REF; + use async_trait::async_trait; + use datafusion::arrow::array::{Array, ArrayRef, BooleanArray, Int64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::catalog::TableProvider; + use datafusion::common::ScalarValue; + use datafusion::datasource::MemTable; + use datafusion::execution::FunctionRegistry; + use datafusion::logical_expr::Projection; + use datafusion::prelude::SessionContext; + use std::ops::Deref; + use substrait::proto::expression::literal::{LiteralType, Struct as LiteralStruct}; + use substrait::proto::expression::mask_expression::{StructItem, StructSelect}; + use substrait::proto::expression::nested::Struct as ExpressionStruct; + use substrait::proto::expression::{Literal, RexType}; + use substrait::proto::{NamedStruct, Type, read_rel, r#type}; + struct TableConsumer { + table: Arc, + } + + impl TableConsumer { + fn new() -> datafusion::common::Result { + Self::new_with_bool_columns(&["a"], &[vec![true, false]]) + } + + fn new_with_bool_columns( + names: &[&str], + columns: &[Vec], + ) -> datafusion::common::Result { + let schema = Arc::new(Schema::new( + names + .iter() + .map(|name| Field::new(*name, DataType::Boolean, true)) + .collect::>(), + )); + let columns = columns + .iter() + .map(|values| Arc::new(BooleanArray::from(values.clone())) as ArrayRef) + .collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns)?; + let table: Arc = + Arc::new(MemTable::try_new(schema, vec![vec![batch]])?); + + Ok(Self { table }) + } + } + + #[async_trait] + impl SubstraitConsumer for TableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(Some(Arc::clone(&self.table))) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + } + + fn named_table_read_with_bool_fields(table: &str, names: &[&str]) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_bool()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::NamedTable(read_rel::NamedTable { + names: vec![table.to_string()], + advanced_extension: None, + })), + ..Default::default() + } + } + + fn local_files_read_with_bool_fields(file_name: &str, names: &[&str]) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_bool()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::LocalFiles(read_rel::LocalFiles { + items: vec![read_rel::local_files::FileOrFiles { + path_type: Some(UriFile(format!("file:///{file_name}"))), + ..Default::default() + }], + advanced_extension: None, + })), + ..Default::default() + } + } + + fn virtual_table_read(names: &[&str]) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_i64()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + ..Default::default() + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_with_bool_values( + names: &[&str], + rows: &[[bool; 2]], + ) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_bool()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: rows + .iter() + .map(|row| LiteralStruct { + fields: row.iter().map(|value| bool_literal(*value)).collect(), + }) + .collect(), + expressions: vec![], + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_with_i32_values_declared_i64( + name: &str, + rows: &[i32], + ) -> ReadRel { + ReadRel { + base_schema: Some(NamedStruct { + names: vec![name.to_string()], + r#struct: Some(r#type::Struct { + types: vec![nullable_i64()], + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: rows + .iter() + .map(|value| LiteralStruct { + fields: vec![i32_literal(*value)], + }) + .collect(), + expressions: vec![], + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_with_i32_expressions_declared_i64( + name: &str, + rows: &[i32], + ) -> ReadRel { + ReadRel { + base_schema: Some(NamedStruct { + names: vec![name.to_string()], + r#struct: Some(r#type::Struct { + types: vec![nullable_i64()], + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: vec![], + expressions: rows + .iter() + .map(|value| ExpressionStruct { + fields: vec![i32_expression(*value)], + }) + .collect(), + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_with_bool_i32_values_declared_bool_i64( + rows: &[(bool, i32)], + ) -> ReadRel { + ReadRel { + base_schema: Some(NamedStruct { + names: vec!["keep".to_string(), "value".to_string()], + r#struct: Some(r#type::Struct { + types: vec![nullable_bool(), nullable_i64()], + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: rows + .iter() + .map(|(keep, value)| LiteralStruct { + fields: vec![bool_literal(*keep), i32_literal(*value)], + }) + .collect(), + expressions: vec![], + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_with_bool_expressions( + names: &[&str], + rows: &[[bool; 2]], + ) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_bool()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: vec![], + expressions: rows + .iter() + .map(|row| ExpressionStruct { + fields: row.iter().map(|value| bool_expression(*value)).collect(), + }) + .collect(), + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_read_produce_one_row() -> ReadRel { + ReadRel { + base_schema: Some(NamedStruct { + names: vec![], + r#struct: Some(r#type::Struct { + types: vec![], + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: vec![LiteralStruct { fields: vec![] }], + expressions: vec![], + })), + ..Default::default() + } + } + + #[expect(deprecated)] + fn virtual_table_expression_read_produce_one_row() -> ReadRel { + ReadRel { + base_schema: Some(NamedStruct { + names: vec![], + r#struct: Some(r#type::Struct { + types: vec![], + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::VirtualTable(read_rel::VirtualTable { + values: vec![], + expressions: vec![ExpressionStruct { fields: vec![] }], + })), + ..Default::default() + } + } + + fn nullable_i64() -> Type { + Type { + kind: Some(r#type::Kind::I64(r#type::I64 { + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Nullable as i32, + })), + } + } + + fn nullable_bool() -> Type { + Type { + kind: Some(r#type::Kind::Bool(r#type::Boolean { + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Nullable as i32, + })), + } + } + + fn required_bool() -> Type { + Type { + kind: Some(r#type::Kind::Bool(r#type::Boolean { + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + })), + } + } + + fn projection(fields: &[i32]) -> MaskExpression { + MaskExpression { + select: Some(StructSelect { + struct_items: fields + .iter() + .map(|field| StructItem { + field: *field, + child: None, + }) + .collect(), + }), + maintain_singular_struct: false, + } + } + + fn nested_projection(field: i32) -> MaskExpression { + MaskExpression { + select: Some(StructSelect { + struct_items: vec![StructItem { + field, + child: Some(substrait::proto::expression::mask_expression::Select { + r#type: None, + }), + }], + }), + maintain_singular_struct: false, + } + } + + fn maintain_singular_projection(fields: &[i32]) -> MaskExpression { + MaskExpression { + maintain_singular_struct: true, + ..projection(fields) + } + } + + fn bool_expression(value: bool) -> Expression { + Expression { + rex_type: Some(RexType::Literal(bool_literal(value))), + } + } + + fn i32_expression(value: i32) -> Expression { + Expression { + rex_type: Some(RexType::Literal(i32_literal(value))), + } + } + + fn true_filter() -> Expression { + bool_expression(true) + } + + fn false_filter() -> Expression { + bool_expression(false) + } + + fn bool_literal(value: bool) -> Literal { + Literal { + nullable: false, + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + literal_type: Some(LiteralType::Boolean(value)), + } + } + + fn i32_literal(value: i32) -> Literal { + Literal { + nullable: false, + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + literal_type: Some(LiteralType::I32(value)), + } + } + + fn assert_projection_columns(projection: &Projection, expected: &[&str]) { + let actual = projection + .expr + .iter() + .map(|expr| match expr { + Expr::Column(column) => column.name.clone(), + _ => panic!("expected column projection, got {expr:?}"), + }) + .collect::>(); + let expected = expected + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + assert_eq!(actual, expected); + } + + async fn assert_bool_column(plan: LogicalPlan, expected: &[Option]) { + let batches = SessionContext::new() + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + assert_eq!(batch.num_columns(), 1); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|row| { + if array.is_null(row) { + None + } else { + Some(array.value(row)) + } + }) + .collect::>() + }) + .collect::>(); + assert_eq!(actual, expected); + } + + async fn assert_bool_rows(plan: LogicalPlan, expected: &[Vec>]) { + let batches = SessionContext::new() + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()) + .map(|row| { + (0..batch.num_columns()) + .map(|column| { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + if array.is_null(row) { + None + } else { + Some(array.value(row)) + } + }) + .collect::>() + }) + .collect::>() + }) + .collect::>(); + assert_eq!(actual, expected); + } + + async fn assert_i64_column(plan: LogicalPlan, expected: &[Option]) { + let batches = SessionContext::new() + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + assert_eq!(batch.num_columns(), 1); + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|row| { + if array.is_null(row) { + None + } else { + Some(array.value(row)) + } + }) + .collect::>() + }) + .collect::>(); + assert_eq!(actual, expected); + } + + async fn assert_row_count(plan: LogicalPlan, expected: usize) { + let batches = SessionContext::new() + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual = batches.iter().map(|batch| batch.num_rows()).sum::(); + assert_eq!(actual, expected); + } + + fn field_reference_filter(field: i32) -> Expression { + Expression { + rex_type: Some(RexType::Selection(Box::new( + substrait::proto::expression::FieldReference { + reference_type: Some( + substrait::proto::expression::field_reference::ReferenceType::DirectReference( + substrait::proto::expression::ReferenceSegment { + reference_type: Some( + substrait::proto::expression::reference_segment::ReferenceType::StructField( + Box::new( + substrait::proto::expression::reference_segment::StructField { + field, + child: None, + }, + ), + ), + ), + }, + ), + ), + root_type: Some( + substrait::proto::expression::field_reference::RootType::RootReference( + substrait::proto::expression::field_reference::RootReference {}, + ), + ), + }, + ))), + } + } + + fn bool_schema(names: &[&str]) -> DFSchema { + DFSchema::new_with_metadata( + names + .iter() + .map(|name| (None, Arc::new(Field::new(*name, DataType::Boolean, true)))) + .collect(), + Default::default(), + ) + .unwrap() + } + + fn assert_unqualified_schema(schema: &DFSchema, expected: &[&str]) { + let actual = schema + .iter() + .map(|(qualifier, field)| { + assert!(qualifier.is_none(), "unexpected qualifier: {qualifier:?}"); + field.name().as_str() + }) + .collect::>(); + assert_eq!(actual, expected); + } + + #[test] + fn apply_masking_projects_fields() { + let schema = bool_schema(&["a", "b"]); + + let schema = apply_masking(schema, &Some(projection(&[1, 0]))).unwrap(); + + assert_eq!(schema.fields().len(), 2); + assert_eq!(schema.fields()[0].name(), "b"); + assert_eq!(schema.fields()[1].name(), "a"); + } + + #[test] + fn apply_masking_projects_fields_with_maintain_singular_struct() { + let schema = bool_schema(&["a", "b"]); + + let schema = + apply_masking(schema, &Some(maintain_singular_projection(&[1, 0]))).unwrap(); + + assert_eq!(schema.fields().len(), 2); + assert_eq!(schema.fields()[0].name(), "b"); + assert_eq!(schema.fields()[1].name(), "a"); + } + + #[test] + fn apply_masking_rejects_unsupported_masks() { + let schema = bool_schema(&["a"]); + + for (projection, expected) in [ + ( + projection(&[-1]), + "Invalid ReadRel projection field index: -1", + ), + ( + projection(&[1]), + "ReadRel projection field index 1 is out of bounds", + ), + ( + nested_projection(0), + "Nested ReadRel projections are not supported", + ), + ] { + let err = apply_masking(schema.clone(), &Some(projection)).unwrap_err(); + assert!(err.to_string().contains(expected), "got: {err}"); + } + } + + async fn assert_read_rejects_unsupported_projection_masks(read: ReadRel) { + let consumer = TableConsumer::new().unwrap(); + + for (projection, expected) in [ + ( + projection(&[-1]), + "Invalid ReadRel projection field index: -1", + ), + ( + projection(&[1]), + "ReadRel projection field index 1 is out of bounds", + ), + ( + nested_projection(0), + "Nested ReadRel projections are not supported", + ), + ] { + let mut read = read.clone(); + read.projection = Some(projection); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + assert!(err.to_string().contains(expected), "got: {err}"); + } + } + + #[tokio::test] + async fn named_table_read_rejects_unsupported_projection_masks() { + assert_read_rejects_unsupported_projection_masks( + named_table_read_with_bool_fields("source", &["a"]), + ) + .await; + } + + #[tokio::test] + async fn local_files_read_rejects_unsupported_projection_masks() { + assert_read_rejects_unsupported_projection_masks( + local_files_read_with_bool_fields("source", &["a"]), + ) + .await; + } + + async fn assert_read_projects_reordered_table_scan(mut read: ReadRel) { + let consumer = TableConsumer::new_with_bool_columns( + &["a", "b"], + &[vec![true, false], vec![false, true]], + ) + .unwrap(); + read.projection = Some(projection(&[1, 0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_bool_rows( + plan, + &[vec![Some(false), Some(true)], vec![Some(true), Some(false)]], + ) + .await; + } + + #[tokio::test] + async fn named_table_read_projects_reordered_table_scan() { + assert_read_projects_reordered_table_scan(named_table_read_with_bool_fields( + "source", + &["a", "b"], + )) + .await; + } + + #[tokio::test] + async fn local_files_read_projects_reordered_table_scan() { + assert_read_projects_reordered_table_scan(local_files_read_with_bool_fields( + "source", + &["a", "b"], + )) + .await; + } + + async fn assert_read_projects_duplicate_table_scan_column(mut read: ReadRel) { + let consumer = TableConsumer::new_with_bool_columns( + &["a", "a"], + &[vec![true, false], vec![false, true]], + ) + .unwrap(); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + let LogicalPlan::TableScan(scan) = &plan else { + panic!("expected TableScan, got {plan:?}"); + }; + assert_eq!(scan.projection.as_deref(), Some(&[1][..])); + assert_eq!(scan.projected_schema.fields().len(), 1); + assert_eq!(scan.projected_schema.fields()[0].name(), "a"); + } + + #[tokio::test] + async fn named_table_read_projects_duplicate_table_scan_column() { + assert_read_projects_duplicate_table_scan_column( + named_table_read_with_bool_fields("source", &["a", "a"]), + ) + .await; + } + + #[tokio::test] + async fn local_files_read_projects_duplicate_table_scan_column() { + assert_read_projects_duplicate_table_scan_column( + local_files_read_with_bool_fields("source", &["a", "a"]), + ) + .await; + } + + async fn assert_read_rejects_incompatible_omitted_projection_field( + mut read: ReadRel, + ) { + let consumer = TableConsumer::new_with_bool_columns( + &["a", "b"], + &[vec![true, false], vec![false, true]], + ) + .unwrap(); + read.base_schema + .as_mut() + .unwrap() + .r#struct + .as_mut() + .unwrap() + .types[0] = required_bool(); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1])); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "Field 'a' is nullable in the DataFusion schema but not nullable in the Substrait schema" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn named_table_read_rejects_incompatible_omitted_projection_field() { + assert_read_rejects_incompatible_omitted_projection_field( + named_table_read_with_bool_fields("source", &["a", "b"]), + ) + .await; + } + + #[tokio::test] + async fn local_files_read_rejects_incompatible_omitted_projection_field() { + assert_read_rejects_incompatible_omitted_projection_field( + local_files_read_with_bool_fields("source", &["a", "b"]), + ) + .await; + } + + #[tokio::test] + async fn local_files_read_projects_with_maintain_singular_struct() { + let consumer = TableConsumer::new_with_bool_columns( + &["a", "b"], + &[vec![true, false], vec![false, true]], + ) + .unwrap(); + let mut read = local_files_read_with_bool_fields("source", &["a", "b"]); + read.projection = Some(maintain_singular_projection(&[1, 0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_bool_rows( + plan, + &[vec![Some(false), Some(true)], vec![Some(true), Some(false)]], + ) + .await; + } + + #[tokio::test] + async fn virtual_table_read_applies_filter_and_projection() { + let consumer = test_consumer(); + let mut read = virtual_table_read(&["a", "b"]); + read.filter = Some(Box::new(true_filter())); + read.projection = Some(projection(&[0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "a"); + let LogicalPlan::Projection(projection) = plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_projection_columns(&projection, &["a"]); + assert!(matches!(projection.input.as_ref(), LogicalPlan::Filter(_))); + } + + #[tokio::test] + async fn virtual_table_values_read_applies_filter_and_projection() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values( + &["a", "b"], + &[[true, false], [false, true]], + ); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "b"); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_projection_columns(projection, &["b"]); + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected Filter, got {:?}", projection.input); + }; + let Expr::Column(column) = &filter.predicate else { + panic!("expected column filter, got {:?}", filter.predicate); + }; + assert_eq!(column.name, "a"); + let LogicalPlan::Values(values) = filter.input.as_ref() else { + panic!("expected Values, got {:?}", filter.input); + }; + assert_eq!(values.values.len(), 2); + assert_bool_column(plan, &[Some(false)]).await; + } + + #[tokio::test] + async fn virtual_table_projection_casts_values_to_declared_schema() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_i32_values_declared_i64("a", &[25, 100]); + read.projection = Some(projection(&[0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields()[0].data_type(), &DataType::Int64); + assert_i64_column(plan, &[Some(25), Some(100)]).await; + } + + #[tokio::test] + async fn virtual_table_expressions_cast_to_declared_schema() { + let consumer = test_consumer(); + let read = virtual_table_read_with_i32_expressions_declared_i64("a", &[25, 100]); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields()[0].data_type(), &DataType::Int64); + assert_i64_column(plan, &[Some(25), Some(100)]).await; + } + + #[tokio::test] + async fn virtual_table_filter_casts_values_to_declared_schema() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_i32_values_declared_bool_i64(&[ + (true, 25), + (false, 100), + ]); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields()[0].data_type(), &DataType::Int64); + assert_i64_column(plan, &[Some(25)]).await; + } + + #[tokio::test] + async fn virtual_table_read_allows_duplicate_names_without_filter_or_projection() { + let consumer = test_consumer(); + let read = virtual_table_read_with_bool_values(&["a", "a"], &[[true, false]]); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 2); + assert_eq!(plan.schema().fields()[0].name(), "a"); + assert_eq!(plan.schema().fields()[1].name(), "a"); + let LogicalPlan::Values(values) = plan else { + panic!("expected Values, got {plan:?}"); + }; + assert_eq!(values.values.len(), 1); + } + + #[tokio::test] + async fn virtual_table_filter_rejects_duplicate_output_names() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values( + &["a", "a"], + &[[true, false], [true, true]], + ); + read.filter = Some(Box::new(field_reference_filter(1))); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "ReadRel filters over duplicate field names require a projection with unique field names" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn virtual_table_filter_and_projection_use_ordinals_with_duplicate_names() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values( + &["a", "a"], + &[[false, true], [true, false]], + ); + read.filter = Some(Box::new(field_reference_filter(1))); + read.projection = Some(projection(&[0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_unqualified_schema(plan.schema(), &["a"]); + assert_bool_column(plan, &[Some(false)]).await; + } + + #[tokio::test] + async fn virtual_table_empty_filter_and_projection_use_ordinals_with_duplicate_names() + { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values(&["a", "a"], &[]); + read.filter = Some(Box::new(field_reference_filter(1))); + read.projection = Some(projection(&[0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_unqualified_schema(plan.schema(), &["a"]); + assert_row_count(plan, 0).await; + } + + #[tokio::test] + async fn virtual_table_projection_uses_ordinal_with_duplicate_names() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values(&["a", "a"], &[[true, false]]); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "a"); + assert_bool_column(plan, &[Some(false)]).await; + } + + #[tokio::test] + async fn virtual_table_empty_projection_uses_ordinal_with_duplicate_names() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values(&["a", "a"], &[]); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_unqualified_schema(plan.schema(), &["a"]); + assert_row_count(plan, 0).await; + } + + #[tokio::test] + async fn virtual_table_projection_reorders_values_by_ordinal() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values(&["a", "b"], &[[true, false]]); + read.projection = Some(projection(&[1, 0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 2); + assert_eq!(plan.schema().fields()[0].name(), "b"); + assert_eq!(plan.schema().fields()[1].name(), "a"); + assert_bool_rows(plan, &[vec![Some(false), Some(true)]]).await; + } + + #[tokio::test] + async fn virtual_table_empty_projection_preserves_row_count() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_values( + &["a", "b"], + &[[true, false], [false, true]], + ); + read.projection = Some(projection(&[])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 0); + assert_row_count(plan, 2).await; + } + + #[tokio::test] + async fn virtual_table_expressions_read_applies_filter_and_projection() { + let consumer = test_consumer(); + let mut read = virtual_table_read_with_bool_expressions( + &["a", "b"], + &[[true, false], [false, true]], + ); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "b"); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_projection_columns(projection, &["b"]); + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected Filter, got {:?}", projection.input); + }; + let Expr::Column(column) = &filter.predicate else { + panic!("expected column filter, got {:?}", filter.predicate); + }; + assert_eq!(column.name, "a"); + let LogicalPlan::Values(values) = filter.input.as_ref() else { + panic!("expected Values, got {:?}", filter.input); + }; + assert_eq!(values.values.len(), 2); + assert_bool_column(plan, &[Some(false)]).await; + } + + #[tokio::test] + async fn virtual_table_produce_one_row_applies_filter() { + let consumer = test_consumer(); + let mut read = virtual_table_read_produce_one_row(); + read.filter = Some(Box::new(false_filter())); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_row_count(plan.clone(), 0).await; + let LogicalPlan::Filter(filter) = plan else { + panic!("expected Filter, got {plan:?}"); + }; + assert_eq!( + filter.predicate, + Expr::Literal(ScalarValue::Boolean(Some(false)), None) + ); + let LogicalPlan::EmptyRelation(empty) = filter.input.as_ref() else { + panic!("expected EmptyRelation, got {:?}", filter.input); + }; + assert!(empty.produce_one_row); + } + + #[tokio::test] + async fn virtual_table_produce_one_row_keeps_true_filter() { + let consumer = test_consumer(); + let mut read = virtual_table_read_produce_one_row(); + read.filter = Some(Box::new(true_filter())); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_row_count(plan, 1).await; + } + + #[tokio::test] + async fn virtual_table_expression_produce_one_row_applies_filter() { + let consumer = test_consumer(); + let mut read = virtual_table_expression_read_produce_one_row(); + read.filter = Some(Box::new(false_filter())); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_row_count(plan.clone(), 0).await; + let LogicalPlan::Filter(filter) = plan else { + panic!("expected Filter, got {plan:?}"); + }; + assert_eq!( + filter.predicate, + Expr::Literal(ScalarValue::Boolean(Some(false)), None) + ); + let LogicalPlan::EmptyRelation(empty) = filter.input.as_ref() else { + panic!("expected EmptyRelation, got {:?}", filter.input); + }; + assert!(empty.produce_one_row); } } diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index 824c79452d86e..b711538851d6e 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/datafusion/substrait/src/logical_plan/consumer/utils.rs @@ -322,7 +322,7 @@ pub(super) fn ensure_schema_compatibility( /// /// If a Substrait field is not nullable, the Substrait plan may be built around assuming it is not /// nullable. As such if DataFusion has that field as nullable the plan should be rejected. -fn ensure_field_compatibility( +pub(super) fn ensure_field_compatibility( datafusion_field: &Field, substrait_field: &Field, ) -> datafusion::common::Result<()> { From 0821af83bb1de9099a353eba8431bc8231c1b55d Mon Sep 17 00:00:00 2001 From: Kadin Rabo Date: Wed, 22 Jul 2026 21:23:56 +0000 Subject: [PATCH 2/2] feat(substrait): add extension table consumer hook --- .../src/logical_plan/consumer/rel/read_rel.rs | 799 +++++++++++++++++- .../consumer/substrait_consumer.rs | 55 +- 2 files changed, 850 insertions(+), 4 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 2a23beb6aeb85..950937aa56c97 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -15,12 +15,12 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::from_substrait_literal; use crate::logical_plan::consumer::from_substrait_named_struct; use crate::logical_plan::consumer::utils::{ ensure_field_compatibility, ensure_schema_compatibility, rename_expressions, }; +use crate::logical_plan::consumer::{ExtensionTableContext, SubstraitConsumer}; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::common::{ Column, DFSchema, DFSchemaRef, TableReference, not_impl_err, plan_datafusion_err, @@ -265,6 +265,25 @@ pub async fn from_read_rel( ) .await } + Some(ReadType::ExtensionTable(extension_table)) => { + let plan = consumer + .resolve_extension_table(ExtensionTableContext::new( + extension_table, + &substrait_schema, + read.best_effort_filter.as_deref(), + read.advanced_extension.as_ref(), + )) + .await?; + + apply_read_rel_filter_and_projection( + consumer, + plan, + substrait_schema, + &read.projection, + &read.filter, + ) + .await + } _ => { not_impl_err!("Unsupported Readtype: {:?}", read.read_type) } @@ -805,14 +824,33 @@ mod tests { use datafusion::common::ScalarValue; use datafusion::datasource::MemTable; use datafusion::execution::FunctionRegistry; - use datafusion::logical_expr::Projection; + use datafusion::logical_expr::{Projection, TableScan, TableScanBuilder}; use datafusion::prelude::SessionContext; + use pbjson_types::Any as ProtoAny; use std::ops::Deref; + use std::sync::atomic::{AtomicBool, Ordering}; use substrait::proto::expression::literal::{LiteralType, Struct as LiteralStruct}; use substrait::proto::expression::mask_expression::{StructItem, StructSelect}; use substrait::proto::expression::nested::Struct as ExpressionStruct; use substrait::proto::expression::{Literal, RexType}; - use substrait::proto::{NamedStruct, Type, read_rel, r#type}; + use substrait::proto::rel::RelType; + use substrait::proto::rel_common::EmitKind; + use substrait::proto::{ + NamedStruct, Rel, RelCommon, Type, extensions::AdvancedExtension, read_rel, + rel_common, r#type, + }; + + struct ExtensionTableConsumer { + called: AtomicBool, + } + + struct ReorderedExtensionTableConsumer; + struct IncompatibleExtensionTableConsumer; + struct NullableExtensionTableConsumer; + struct OneRowExtensionTableConsumer; + struct ProjectedTableScanExtensionTableConsumer; + struct DuplicateNameTableScanExtensionTableConsumer; + struct DuplicateNameGenericExtensionTableConsumer; struct TableConsumer { table: Arc, } @@ -862,6 +900,376 @@ mod tests { } } + #[async_trait] + impl SubstraitConsumer for ExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + self.called.store(true, Ordering::SeqCst); + + let detail = context.extension_table.detail.as_ref().unwrap(); + assert_eq!(detail.type_url, "type.example/custom"); + assert_eq!(detail.value.as_ref(), &[25, 100]); + + if let Some(filter) = context.best_effort_filter { + assert_eq!(filter, &true_filter()); + } + if let Some(advanced_extension) = context.advanced_extension { + assert_eq!( + advanced_extension.optimization[0].type_url, + "type.example/optimization" + ); + assert_eq!( + advanced_extension.enhancement.as_ref().unwrap().type_url, + "type.example/enhancement" + ); + } + + Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchemaRef::new(context.base_schema.clone()), + })) + } + } + + #[async_trait] + impl SubstraitConsumer for ReorderedExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + _context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let schema = DFSchema::new_with_metadata( + vec![ + (None, Arc::new(Field::new("extra", DataType::Boolean, true))), + (None, Arc::new(Field::new("b", DataType::Boolean, true))), + (None, Arc::new(Field::new("a", DataType::Boolean, true))), + ], + Default::default(), + )?; + Ok(LogicalPlan::Values(Values { + schema: DFSchemaRef::new(schema), + values: vec![ + vec![ + df_bool_expression(false), + df_bool_expression(false), + df_bool_expression(true), + ], + vec![ + df_bool_expression(false), + df_bool_expression(true), + df_bool_expression(false), + ], + ], + })) + } + } + + #[async_trait] + impl SubstraitConsumer for IncompatibleExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + _context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let schema = DFSchema::new_with_metadata( + vec![(None, Arc::new(Field::new("a", DataType::Boolean, true)))], + Default::default(), + )?; + Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchemaRef::new(schema), + })) + } + } + + #[async_trait] + impl SubstraitConsumer for NullableExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + _context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let schema = DFSchema::new_with_metadata( + vec![(None, Arc::new(Field::new("a", DataType::Boolean, true)))], + Default::default(), + )?; + Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchemaRef::new(schema), + })) + } + } + + #[async_trait] + impl SubstraitConsumer for OneRowExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: true, + schema: DFSchemaRef::new(context.base_schema.clone()), + })) + } + } + + #[async_trait] + impl SubstraitConsumer for ProjectedTableScanExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + _context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Boolean, true), + Field::new("extra", DataType::Boolean, true), + Field::new("b", DataType::Boolean, true), + ])); + let columns: Vec = vec![ + Arc::new(BooleanArray::from(vec![true, false])), + Arc::new(BooleanArray::from(vec![false, false])), + Arc::new(BooleanArray::from(vec![false, true])), + ]; + let batch = RecordBatch::try_new(Arc::clone(&schema), columns)?; + let table: Arc = + Arc::new(MemTable::try_new(schema, vec![vec![batch]])?); + let scan = TableScanBuilder::new( + "extension_table_source", + provider_as_source(table), + ) + .with_projection(Some(vec![0, 2])) + .build()?; + + Ok(LogicalPlan::TableScan(scan)) + } + } + + #[async_trait] + impl SubstraitConsumer for DuplicateNameTableScanExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + _context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a_0", DataType::Boolean, true), + Field::new("a_1", DataType::Boolean, true), + ])); + let projected_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Boolean, true), + Field::new("a", DataType::Boolean, true), + ])); + let columns: Vec = vec![ + Arc::new(BooleanArray::from(vec![true, true])), + Arc::new(BooleanArray::from(vec![false, true])), + ]; + let batch = RecordBatch::try_new(Arc::clone(&source_schema), columns)?; + let table: Arc = + Arc::new(MemTable::try_new(source_schema, vec![vec![batch]])?); + let scan = TableScan { + table_name: TableReference::bare("duplicate_name_extension_table"), + source: provider_as_source(table), + projection: None, + projected_schema: DFSchemaRef::new(DFSchema::try_from(projected_schema)?), + filters: vec![], + fetch: None, + statistics_requests: Default::default(), + }; + + Ok(LogicalPlan::TableScan(scan)) + } + } + + #[async_trait] + impl SubstraitConsumer for DuplicateNameGenericExtensionTableConsumer { + async fn resolve_table_ref( + &self, + _table_ref: &TableReference, + ) -> datafusion::common::Result>> { + Ok(None) + } + + fn get_extensions(&self) -> &Extensions { + TEST_EXTENSIONS.deref() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + TEST_SESSION_STATE.deref() + } + + async fn resolve_extension_table( + &self, + context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + let plan = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchemaRef::new(context.base_schema.clone()), + }); + + LogicalPlanBuilder::from(plan) + .filter(df_bool_expression(true))? + .build() + } + } + + fn extension_table_read(detail: Option) -> ReadRel { + extension_table_read_with_names(detail, &["value"]) + } + + fn extension_table_read_with_names( + detail: Option, + names: &[&str], + ) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_i64()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::ExtensionTable(read_rel::ExtensionTable { + detail, + })), + ..Default::default() + } + } + + fn extension_table_read_with_bool_fields( + detail: Option, + names: &[&str], + ) -> ReadRel { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let field_count = names.len(); + + ReadRel { + base_schema: Some(NamedStruct { + names, + r#struct: Some(r#type::Struct { + types: (0..field_count).map(|_| nullable_bool()).collect(), + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + nullability: r#type::Nullability::Required as i32, + }), + }), + read_type: Some(ReadType::ExtensionTable(read_rel::ExtensionTable { + detail, + })), + ..Default::default() + } + } + fn named_table_read_with_bool_fields(table: &str, names: &[&str]) -> ReadRel { let names = names .iter() @@ -1195,6 +1603,10 @@ mod tests { } } + fn df_bool_expression(value: bool) -> Expr { + Expr::Literal(ScalarValue::Boolean(Some(value)), None) + } + fn true_filter() -> Expression { bool_expression(true) } @@ -1373,6 +1785,26 @@ mod tests { } } + fn extension_detail() -> ProtoAny { + ProtoAny { + type_url: "type.example/custom".to_string(), + value: vec![25, 100].into(), + } + } + + fn advanced_extension() -> AdvancedExtension { + AdvancedExtension { + optimization: vec![ProtoAny { + type_url: "type.example/optimization".to_string(), + value: vec![1].into(), + }], + enhancement: Some(ProtoAny { + type_url: "type.example/enhancement".to_string(), + value: vec![2].into(), + }), + } + } + fn bool_schema(names: &[&str]) -> DFSchema { DFSchema::new_with_metadata( names @@ -1614,6 +2046,338 @@ mod tests { .await; } + #[tokio::test] + async fn extension_table_read_uses_consumer_handler() { + let consumer = ExtensionTableConsumer { + called: AtomicBool::new(false), + }; + let mut read = extension_table_read(Some(extension_detail())); + read.best_effort_filter = Some(Box::new(true_filter())); + read.advanced_extension = Some(advanced_extension()); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert!(consumer.called.load(Ordering::SeqCst)); + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "value"); + } + + #[tokio::test] + async fn extension_table_read_applies_filter_and_projection() { + let consumer = ExtensionTableConsumer { + called: AtomicBool::new(false), + }; + let mut read = + extension_table_read_with_names(Some(extension_detail()), &["a", "b"]); + read.filter = Some(Box::new(true_filter())); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert!(consumer.called.load(Ordering::SeqCst)); + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "b"); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_projection_columns(projection, &["b"]); + assert!(matches!(projection.input.as_ref(), LogicalPlan::Filter(_))); + } + + #[tokio::test] + async fn extension_table_read_applies_projection_without_filter() { + let consumer = ExtensionTableConsumer { + called: AtomicBool::new(false), + }; + let mut read = + extension_table_read_with_names(Some(extension_detail()), &["a", "b"]); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert!(consumer.called.load(Ordering::SeqCst)); + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "b"); + assert_row_count(plan, 0).await; + } + + #[tokio::test] + async fn extension_table_read_applies_filter_without_projection() { + let consumer = ExtensionTableConsumer { + called: AtomicBool::new(false), + }; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.filter = Some(Box::new(field_reference_filter(0))); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert!(consumer.called.load(Ordering::SeqCst)); + assert_eq!(plan.schema().fields().len(), 2); + let LogicalPlan::Filter(filter) = plan else { + panic!("expected Filter, got {plan:?}"); + }; + let Expr::Column(column) = &filter.predicate else { + panic!("expected column filter, got {:?}", filter.predicate); + }; + assert_eq!(column.name, "a"); + } + + #[tokio::test] + async fn extension_table_read_rejects_incompatible_handler_schema() { + let consumer = IncompatibleExtensionTableConsumer; + let read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!(err.to_string().contains("No field named b"), "got: {err}"); + } + + #[tokio::test] + async fn extension_table_read_rejects_nullable_handler_field_for_required_schema() { + let consumer = NullableExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a"]); + read.base_schema + .as_mut() + .unwrap() + .r#struct + .as_mut() + .unwrap() + .types[0] = required_bool(); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "Field 'a' is nullable in the DataFusion schema but not nullable in the Substrait schema" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn extension_table_read_rejects_unsupported_projection_masks() { + let consumer = ExtensionTableConsumer { + called: AtomicBool::new(false), + }; + + for (projection, expected) in [ + ( + projection(&[-1]), + "Invalid ReadRel projection field index: -1", + ), + ( + projection(&[1]), + "ReadRel projection field index 1 is out of bounds", + ), + ( + nested_projection(0), + "Nested ReadRel projections are not supported", + ), + ] { + let mut read = extension_table_read(Some(extension_detail())); + read.projection = Some(projection); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + assert!(err.to_string().contains(expected), "got: {err}"); + } + } + + #[tokio::test] + async fn extension_table_projection_composes_table_scan_projection() { + let consumer = ProjectedTableScanExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + let LogicalPlan::TableScan(scan) = &plan else { + panic!("expected TableScan, got {plan:?}"); + }; + assert_eq!(scan.projection.as_deref(), Some(&[2][..])); + assert_eq!(scan.projected_schema.fields().len(), 1); + assert_eq!(scan.projected_schema.fields()[0].name(), "b"); + assert_bool_column(plan, &[Some(false), Some(true)]).await; + } + + #[tokio::test] + async fn extension_table_empty_projection_preserves_table_scan_row_count() { + let consumer = ProjectedTableScanExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.projection = Some(projection(&[])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 0); + assert_row_count(plan, 2).await; + } + + #[tokio::test] + async fn extension_table_filter_uses_read_schema_field_order() { + let consumer = ReorderedExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[0])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_eq!(projection.schema.fields().len(), 1); + assert_eq!(projection.schema.fields()[0].name(), "a"); + assert_projection_columns(projection, &["a"]); + + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected Filter, got {:?}", projection.input); + }; + let Expr::Column(column) = &filter.predicate else { + panic!("expected column filter, got {:?}", filter.predicate); + }; + assert_eq!(column.name, "a"); + assert_bool_column(plan, &[Some(true)]).await; + } + + #[tokio::test] + async fn extension_table_filter_can_reference_projected_away_field() { + let consumer = ReorderedExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected Projection, got {plan:?}"); + }; + assert_eq!(projection.schema.fields().len(), 1); + assert_eq!(projection.schema.fields()[0].name(), "b"); + assert_projection_columns(projection, &["b"]); + + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected Filter, got {:?}", projection.input); + }; + let Expr::Column(column) = &filter.predicate else { + panic!("expected column filter, got {:?}", filter.predicate); + }; + assert_eq!(column.name, "a"); + assert_bool_column(plan, &[Some(false)]).await; + } + + #[tokio::test] + async fn extension_table_filter_rejects_duplicate_table_scan_names() { + let consumer = DuplicateNameTableScanExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "a"]); + read.filter = Some(Box::new(field_reference_filter(1))); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "ReadRel filters over generic plans with duplicate field names are not supported" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn extension_table_read_normalizes_reordered_schema_without_projection() { + let consumer = ReorderedExtensionTableConsumer; + let read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 2); + assert_eq!(plan.schema().fields()[0].name(), "a"); + assert_eq!(plan.schema().fields()[1].name(), "b"); + assert_bool_rows( + plan, + &[vec![Some(true), Some(false)], vec![Some(false), Some(true)]], + ) + .await; + } + + #[tokio::test] + async fn extension_table_read_common_emit_applies_after_filter_and_projection() { + let consumer = ReorderedExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.filter = Some(Box::new(field_reference_filter(0))); + read.projection = Some(projection(&[1, 0])); + read.common = Some(RelCommon { + emit_kind: Some(EmitKind::Emit(rel_common::Emit { + output_mapping: vec![1], + })), + hint: None, + advanced_extension: None, + }); + let rel = Rel { + rel_type: Some(RelType::Read(Box::new(read))), + }; + + let plan = super::super::from_substrait_rel(&consumer, &rel) + .await + .unwrap(); + + assert_eq!(plan.schema().fields().len(), 1); + assert_eq!(plan.schema().fields()[0].name(), "a"); + assert_bool_column(plan, &[Some(true)]).await; + } + + #[tokio::test] + async fn extension_table_projection_rejects_generic_plan_with_duplicate_names() { + let consumer = DuplicateNameGenericExtensionTableConsumer; + let mut read = + extension_table_read_with_names(Some(extension_detail()), &["a", "a"]); + read.projection = Some(projection(&[1])); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "ReadRel projections over generic plans with duplicate field names are not supported" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn extension_table_read_rejects_duplicate_names_on_name_based_path() { + let consumer = IncompatibleExtensionTableConsumer; + let read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "a"]); + + let err = from_read_rel(&consumer, &read).await.unwrap_err(); + + assert!( + err.to_string().contains( + "ReadRel schemas with duplicate field names must match by position" + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn extension_table_empty_projection_preserves_generic_plan_row_count() { + let consumer = OneRowExtensionTableConsumer; + let mut read = + extension_table_read_with_bool_fields(Some(extension_detail()), &["a", "b"]); + read.projection = Some(projection(&[])); + + let plan = from_read_rel(&consumer, &read).await.unwrap(); + + assert_eq!(plan.schema().fields().len(), 0); + assert_row_count(plan, 1).await; + } + #[tokio::test] async fn virtual_table_read_applies_filter_and_projection() { let consumer = test_consumer(); @@ -1908,4 +2672,33 @@ mod tests { }; assert!(empty.produce_one_row); } + + #[tokio::test] + async fn default_consumer_rejects_extension_table_read() { + let consumer = test_consumer(); + let err = + from_read_rel(&consumer, &extension_table_read(Some(extension_detail()))) + .await + .unwrap_err(); + + assert!( + err.to_string() + .contains("Missing handler for ExtensionTable: type.example/custom"), + "got: {err}" + ); + } + + #[tokio::test] + async fn default_consumer_rejects_extension_table_read_without_detail() { + let consumer = test_consumer(); + let err = from_read_rel(&consumer, &extension_table_read(None)) + .await + .unwrap_err(); + + assert!( + err.to_string() + .contains("Missing handler for ExtensionTable"), + "got: {err}" + ); + } } diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index bbd80b4cff001..60ffe0bf5e468 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -50,6 +50,35 @@ use substrait::proto::{ FilterRel, JoinRel, ProjectRel, ReadRel, Rel, SetRel, SortRel, r#type, }; +/// Context passed to [`SubstraitConsumer::resolve_extension_table`]. +#[non_exhaustive] +pub struct ExtensionTableContext<'a> { + /// The typed Substrait extension table read. + pub extension_table: &'a proto::read_rel::ExtensionTable, + /// The decoded `ReadRel` base schema. + pub base_schema: &'a DFSchema, + /// A best-effort filter that the extension table resolver may apply. + pub best_effort_filter: Option<&'a Expression>, + /// Advanced extension metadata from the enclosing `ReadRel`. + pub advanced_extension: Option<&'a proto::extensions::AdvancedExtension>, +} + +impl<'a> ExtensionTableContext<'a> { + pub(crate) fn new( + extension_table: &'a proto::read_rel::ExtensionTable, + base_schema: &'a DFSchema, + best_effort_filter: Option<&'a Expression>, + advanced_extension: Option<&'a proto::extensions::AdvancedExtension>, + ) -> Self { + Self { + extension_table, + base_schema, + best_effort_filter, + advanced_extension, + } + } +} + #[async_trait] /// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans. /// It can be implemented by users to allow for custom handling of relations, expressions, etc. @@ -446,9 +475,33 @@ pub trait SubstraitConsumer: Send + Sync + Sized { // User-Defined Functionality - // The details of extension relations, and how to handle them, are fully up to users to specify. + // The details of extension reads and relations, and how to handle them, + // are fully up to users to specify. // The following methods allow users to customize the consumer behaviour + /// Resolve this extension table as a base plan. + /// + /// Implementations should not apply the `ReadRel` filter, projection, or + /// common emit mapping. Those fields are applied by the generic relation + /// consumer after this hook returns. + /// Implementations are responsible for interpreting the best-effort filter + /// and advanced extension fields if they are relevant. + /// The returned plan may include extra fields whose names do not duplicate + /// `base_schema` fields, but every `base_schema` field must be present with + /// a compatible type. + async fn resolve_extension_table( + &self, + context: ExtensionTableContext<'_>, + ) -> datafusion::common::Result { + if let Some(detail) = context.extension_table.detail.as_ref() { + return substrait_err!( + "Missing handler for ExtensionTable: {}", + detail.type_url + ); + } + substrait_err!("Missing handler for ExtensionTable") + } + async fn consume_extension_leaf( &self, rel: &ExtensionLeafRel,