Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 49 additions & 15 deletions datafusion/expr/src/expr_schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,18 +71,31 @@ pub trait ExprSchemable {
-> Result<(DataType, bool)>;
}

/// Derives the output field for a cast expression from the source field.
/// Derives the output field for a cast expression from the source and target
/// fields. Type-only casts preserve source metadata, while an explicit target
/// field supplies its own metadata.
///
/// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL.
fn cast_output_field(
source_field: &FieldRef,
target_type: &DataType,
target_field: &FieldRef,
force_nullable: bool,
) -> Arc<Field> {
// `Cast::new` and `TryCast::new` use this field when only a target
// type is known. In that case, retain the source field's metadata.
let type_only_target = target_field.name().is_empty()
&& target_field.is_nullable()
&& target_field.metadata().is_empty();
let metadata = if type_only_target {
source_field.metadata().clone()
} else {
target_field.metadata().clone()
};
let mut f = source_field
.as_ref()
.clone()
.with_data_type(target_type.clone())
.with_metadata(source_field.metadata().clone());
.with_data_type(target_field.data_type().clone())
.with_metadata(metadata);
if force_nullable {
f = f.with_nullable(true);
}
Expand Down Expand Up @@ -623,20 +636,16 @@ impl ExprSchemable for Expr {
func.return_field_from_args(args)
}
// _ => Ok((self.get_type(schema)?, self.nullable(schema)?)),
Expr::Cast(Cast { expr, field }) => {
expr.to_field(schema).map(|(_table_ref, src)| {
cast_output_field(&src, field.data_type(), false)
})
}
Expr::Cast(Cast { expr, field }) => expr
.to_field(schema)
.map(|(_table_ref, src)| cast_output_field(&src, field, false)),
Expr::Placeholder(Placeholder {
id: _,
field: Some(field),
}) => Ok(Arc::clone(field).renamed(&schema_name)),
Expr::TryCast(TryCast { expr, field }) => {
expr.to_field(schema).map(|(_table_ref, src)| {
cast_output_field(&src, field.data_type(), true)
})
}
Expr::TryCast(TryCast { expr, field }) => expr
.to_field(schema)
.map(|(_table_ref, src)| cast_output_field(&src, field, true)),
Expr::LambdaVariable(LambdaVariable {
field: Some(field), ..
}) => Ok(Arc::clone(field).renamed(&schema_name)),
Expand Down Expand Up @@ -1149,7 +1158,7 @@ mod tests {
}

#[test]
fn test_expr_metadata() {
fn test_expr_metadata() -> Result<()> {
let mut meta = HashMap::new();
meta.insert("bar".to_string(), "buzz".to_string());
let meta = FieldMetadata::from(meta);
Expand Down Expand Up @@ -1179,13 +1188,38 @@ mod tests {
// verify to_field method populates metadata
assert_eq!(meta, expr.metadata(&schema).unwrap());

// An explicit cast target replaces source metadata. A type-only cast
// continues to preserve it.
let target_metadata = HashMap::from([(
"ARROW:extension:name".to_string(),
"arrow.uuid".to_string(),
)]);
let target_field = Arc::new(
Field::new("", DataType::FixedSizeBinary(16), true)
.with_metadata(target_metadata.clone()),
);
let cast = Expr::Cast(Cast::new_from_field(
Box::new(expr.clone()),
Arc::clone(&target_field),
));
assert_eq!(cast.to_field(&schema)?.1.metadata(), &target_metadata);

let try_cast = Expr::TryCast(TryCast::new_from_field(
Box::new(expr.clone()),
target_field,
));
let try_cast_field = try_cast.to_field(&schema)?.1;
assert_eq!(try_cast_field.metadata(), &target_metadata);
assert!(try_cast_field.is_nullable());

// outer ref constructed by `out_ref_col_with_metadata` should be metadata-preserving
let outer_ref = out_ref_col_with_metadata(
DataType::Int32,
meta.to_hashmap(),
Column::from_name("foo"),
);
assert_eq!(meta, outer_ref.metadata(&schema).unwrap());
Ok(())
}

#[test]
Expand Down
156 changes: 145 additions & 11 deletions datafusion/physical-plan/src/projection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1014,6 +1014,11 @@ pub fn remove_unnecessary_projections(
plan: Arc<dyn ExecutionPlan>,
) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
let maybe_modified = if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
// Metadata is observable by physical expressions. Moving a projection that
// defines metadata can therefore change query results.
if projection_overrides_metadata(projection)? {
return Ok(Transformed::no(plan));
}
// If the projection does not cause any change on the input, we can
// safely remove it:
if is_projection_removable(projection) {
Expand All @@ -1029,10 +1034,12 @@ pub fn remove_unnecessary_projections(
Ok(maybe_modified.map_or_else(|| Transformed::no(plan), Transformed::yes))
}

/// Compare the inputs and outputs of the projection. All expressions must be
/// columns without alias, and projection does not change the order of fields.
/// For example, if the input schema is `a, b`, `SELECT a, b` is removable,
/// but `SELECT b, a` and `SELECT a+1, b` and `SELECT a AS c, b` are not.
/// Compare the input and output of a projection. A removable projection contains
/// only unaliased columns in input order and preserves the exact schema, including
/// field and schema metadata.
///
/// For example, if the input schema is `a, b`, `SELECT a, b` is removable, but
/// `SELECT b, a`, `SELECT a + 1, b`, and a metadata-only projection are not.
fn is_projection_removable(projection: &ProjectionExec) -> bool {
let exprs = projection.expr();
exprs.iter().enumerate().all(|(idx, proj_expr)| {
Expand All @@ -1041,6 +1048,7 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool {
};
col.name() == proj_expr.alias && col.index() == idx
}) && exprs.len() == projection.input().schema().fields().len()
&& projection.schema() == projection.input().schema()
Comment on lines 1050 to +1051

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This might be putting more restrictions than just metadata equality. It might be fine, but if we want to play it safe it could be better to just do && projection.schema().metadata() == projection.input().schema().metadata()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to check full schema because even if the schema metadata is equal things like the field metadata might not be thus we nee to check this as well.

I don't see anything in the schema which would be overestricting this. I may be missing something though

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pretty much the order of columns. I bet that's why the current checks are like they are right now.

I think it's fine though, if this becomes too restrictive it will start popping up in tests

}

/// Given the expression set of a projection, checks if the projection causes
Expand Down Expand Up @@ -1073,14 +1081,20 @@ pub fn new_projections_for_columns(
.collect()
}

/// Creates a new [`ProjectionExec`] instance with the given child plan and
/// projected expressions.
/// Creates an equivalent [`ProjectionExec`] with a new child.
///
/// The original output metadata is preserved because a parent expression may
/// observe it; recomputing metadata from `child` could change query results.
pub fn make_with_child(
projection: &ProjectionExec,
child: &Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
ProjectionExec::try_new(projection.expr().to_vec(), Arc::clone(child))
.map(|e| Arc::new(e) as _)
ProjectionExec::try_new_with_schema_metadata(
projection.expr().to_vec(),
Arc::clone(child),
projection.schema().as_ref(),
)
.map(|e| Arc::new(e) as _)
}

/// Returns `true` if all the expressions in the argument are `Column`s.
Expand Down Expand Up @@ -1326,17 +1340,44 @@ pub fn update_join_filter(
})
}

/// Returns whether a projection defines metadata that its expressions and input
/// schema cannot reproduce.
///
/// Such a projection is an execution boundary: a parent expression such as
/// `arrow_metadata` can observe its output field metadata.
fn projection_overrides_metadata(projection: &ProjectionExec) -> Result<bool> {
let derived_schema = projection
.projector
.projection()
.project_schema(projection.input().schema().as_ref())?;
let output_schema = projection.schema();
Ok(derived_schema.metadata() != output_schema.metadata()
|| derived_schema
.fields()
.iter()
.zip(output_schema.fields())
.any(|(derived, output)| derived.metadata() != output.metadata()))
}

/// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns
/// `None` if nothing could be merged.
fn try_collapse_projection_chain(
outer: &ProjectionExec,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
if projection_overrides_metadata(outer)? {
return Ok(None);
}

let mut current_exprs: Vec<ProjectionExpr> = outer.expr().to_vec();
let mut current_input: Arc<dyn ExecutionPlan> = Arc::clone(outer.input());
let mut column_ref_map: HashMap<Column, usize> = HashMap::new();
let mut collapsed_any = false;

'outer: while let Some(inner_proj) = current_input.downcast_ref::<ProjectionExec>() {
if projection_overrides_metadata(inner_proj)? {
break;
}

// Collect the column references usage in the outer projection.
column_ref_map.clear();
for proj_expr in &current_exprs {
Expand Down Expand Up @@ -1385,9 +1426,14 @@ fn try_collapse_projection_chain(
return Ok(None);
}

// To unify 3 or more sequential projections:
// Expression substitution must not change the outer projection's output
// metadata contract.
let unified: Arc<dyn ExecutionPlan> =
Arc::new(ProjectionExec::try_new(current_exprs, current_input)?);
Arc::new(ProjectionExec::try_new_with_schema_metadata(
current_exprs,
current_input,
outer.schema().as_ref(),
)?);
remove_unnecessary_projections(unified).data().map(Some)
}

Expand Down Expand Up @@ -1517,11 +1563,14 @@ mod tests {
use crate::test;
use crate::test::exec::StatisticsExec;

use arrow::array::StringArray;
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::ScalarValue;
use datafusion_common::stats::{ColumnStatistics, Precision, Statistics};

use datafusion_expr::Operator;
use datafusion_expr::{Operator, ScalarUDF};
use datafusion_functions::core::arrow_metadata::ArrowMetadataFunc;
use datafusion_physical_expr::ScalarFunctionExpr;
use datafusion_physical_expr::expressions::{
BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit,
};
Expand Down Expand Up @@ -1566,6 +1615,91 @@ mod tests {
Ok(())
}

fn identity_projection_with_metadata(
input: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
let metadata_schema =
Schema::new_with_metadata(
vec![Field::new("i", DataType::Int32, true).with_metadata(
HashMap::from([("event_field".to_string(), "true".to_string())]),
)],
HashMap::from([("schema-key".to_string(), "schema-value".to_string())]),
);
Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata(
[ProjectionExpr {
expr: Arc::new(Column::new("i", 0)),
alias: "i".to_string(),
}],
input,
&metadata_schema,
)?))
}

#[test]
fn test_metadata_projection_is_not_removable() -> Result<()> {
let projection = identity_projection_with_metadata(test::scan_partitioned(1))?;
let expected_schema = projection.schema();

let optimized = remove_unnecessary_projections(projection)?.data;

assert!(optimized.downcast_ref::<ProjectionExec>().is_some());
assert_eq!(optimized.schema(), expected_schema);
Ok(())
}

#[test]
fn test_make_with_child_preserves_output_metadata() -> Result<()> {
let projection = identity_projection_with_metadata(test::scan_partitioned(1))?;
let projection = projection
.downcast_ref::<ProjectionExec>()
.expect("test plan should be a ProjectionExec");

let rebuilt = make_with_child(projection, &test::scan_partitioned(1))?;

assert_eq!(rebuilt.schema(), projection.schema());
Ok(())
}

#[tokio::test]
async fn test_metadata_observing_parent_blocks_projection_collapse() -> Result<()> {
let inner = identity_projection_with_metadata(test::scan_partitioned(1))?;
let arrow_metadata = ScalarFunctionExpr::new(
"arrow_metadata",
Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())),
vec![
Arc::new(Column::new("i", 0)),
Arc::new(Literal::new(ScalarValue::Utf8(Some(
"event_field".to_string(),
)))),
],
Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)),
Arc::new(ConfigOptions::default()),
);
let outer: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
[ProjectionExpr {
expr: Arc::new(arrow_metadata),
alias: "metadata".to_string(),
}],
inner,
)?);

let outer_projection = outer
.downcast_ref::<ProjectionExec>()
.expect("test plan should be a ProjectionExec");
assert!(try_collapse_projection_chain(outer_projection)?.is_none());

let optimized = remove_unnecessary_projections(outer)?.data;
let batches =
collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?;
let values = batches[0]
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.expect("metadata expression should return Utf8");
assert_eq!(values.value(0), "true");
Ok(())
}

#[test]
fn test_collect_column_indices() -> Result<()> {
let expr = Arc::new(BinaryExpr::new(
Expand Down