Skip to content
Draft
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
69 changes: 66 additions & 3 deletions datafusion/core/src/physical_planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ use crate::physical_plan::explain::ExplainExec;
use crate::physical_plan::filter::FilterExecBuilder;
use crate::physical_plan::joins::utils as join_utils;
use crate::physical_plan::joins::{
CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec,
PartitionMode, SortMergeJoinExec,
};
use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr};
Expand Down Expand Up @@ -93,8 +94,8 @@ use datafusion_expr::physical_planning_context::{
use datafusion_expr::utils::{expr_to_columns, split_conjunction};
use datafusion_expr::{
Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan,
WindowFrame, WindowFrameBound, WriteOp,
FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType,
StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp,
};
use datafusion_physical_expr::aggregate::{
AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder,
Expand Down Expand Up @@ -1860,6 +1861,67 @@ impl DefaultPhysicalPlanner {
join
}
}
LogicalPlan::AsOfJoin(join) => {
let [physical_left, physical_right] = children.two()?;
let join_on = join
.on
.iter()
.map(|(left, right)| {
Ok((
create_physical_expr(
left,
join.left.schema(),
execution_props,
planning_ctx,
)?,
create_physical_expr(
right,
join.right.schema(),
execution_props,
planning_ctx,
)?,
))
})
.collect::<Result<join_utils::JoinOn>>()?;
let match_condition = AsOfMatchExpr::new(
create_physical_expr(
&join.match_condition.left,
join.left.schema(),
execution_props,
planning_ctx,
)?,
join.match_condition.op,
create_physical_expr(
&join.match_condition.right,
join.right.schema(),
execution_props,
planning_ctx,
)?,
);
let omitted_right = if join.join_constraint == JoinConstraint::Using {
join.on
.iter()
.map(|(_, right)| {
let column = right.get_as_join_column().ok_or_else(|| {
internal_datafusion_err!("ASOF USING key is not a column")
})?;
join.right.schema().index_of_column(column)
})
.collect::<Result<HashSet<_>>>()?
} else {
HashSet::new()
};
let right_output_indices = (0..join.right.schema().fields().len())
.filter(|index| !omitted_right.contains(index))
.collect();
Arc::new(AsOfJoinExec::try_new(
physical_left,
physical_right,
join_on,
match_condition,
right_output_indices,
)?)
}
LogicalPlan::RecursiveQuery(RecursiveQuery {
name,
is_distinct,
Expand Down Expand Up @@ -2359,6 +2421,7 @@ fn extract_dml_filters(
| LogicalPlan::Sort(_)
| LogicalPlan::Union(_)
| LogicalPlan::Join(_)
| LogicalPlan::AsOfJoin(_)
| LogicalPlan::Repartition(_)
| LogicalPlan::Aggregate(_)
| LogicalPlan::Window(_)
Expand Down
42 changes: 42 additions & 0 deletions datafusion/core/tests/physical_optimizer/filter_pushdown.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ use datafusion_physical_plan::{
coalesce_partitions::CoalescePartitionsExec,
collect,
filter::{FilterExec, FilterExecBuilder},
joins::{AsOfJoinExec, AsOfMatchExpr},
projection::ProjectionExec,
repartition::RepartitionExec,
sorts::sort::SortExec,
Expand Down Expand Up @@ -123,6 +124,47 @@ fn test_pushdown_volatile_functions_not_allowed() {
);
}

#[test]
fn test_asof_join_pushes_only_left_filters() {
let left = TestScanBuilder::new(schema()).with_support(true).build();
let right = TestScanBuilder::new(schema()).with_support(true).build();
let join = Arc::new(
AsOfJoinExec::try_new(
left,
right,
vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())],
AsOfMatchExpr::new(
col("b", &schema()).unwrap(),
Operator::GtEq,
col("b", &schema()).unwrap(),
),
vec![0, 1, 2],
)
.unwrap(),
);
let left_filter = Arc::new(BinaryExpr::new(
Arc::new(Column::new("c", 2)),
Operator::Gt,
Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
)) as Arc<dyn PhysicalExpr>;
let right_filter = Arc::new(BinaryExpr::new(
Arc::new(Column::new("c", 5)),
Operator::Gt,
Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
)) as Arc<dyn PhysicalExpr>;
let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter));
let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap());
let mut config = ConfigOptions::default();
config.execution.parquet.pushdown_filters = true;
let optimized = FilterPushdown::new().optimize(plan, &config).unwrap();
let formatted = format_plan_for_test(&optimized);

assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}");
assert!(formatted.contains("AsOfJoinExec:"), "{formatted}");
assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}");
assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}");
}

/// Show that we can use config options to determine how to do pushdown.
#[test]
fn test_pushdown_into_scan_with_config_options() {
Expand Down
107 changes: 103 additions & 4 deletions datafusion/expr/src/logical_plan/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,10 @@ use crate::expr_rewriter::{
rewrite_sort_cols_by_aggs,
};
use crate::logical_plan::{
Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest,
Values, Window,
Aggregate, Analyze, AsOfJoin, AsOfMatch, Distinct, DistinctOn, EmptyRelation,
Explain, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning,
PlanType, Prepare, Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder,
Union, Unnest, Values, Window,
};
use crate::select_expr::SelectExpr;
use crate::utils::{
Expand Down Expand Up @@ -1007,6 +1007,68 @@ impl LogicalPlanBuilder {
)
}

/// Apply a left-preserving ASOF join using equality expressions and one
/// ordered match condition.
pub fn asof_join(
self,
right: LogicalPlan,
on: Vec<(Expr, Expr)>,
match_condition: AsOfMatch,
) -> Result<Self> {
self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On)
}

/// Apply a left-preserving ASOF join using `USING` equality keys.
pub fn asof_join_using(
self,
right: LogicalPlan,
using_keys: Vec<Column>,
match_condition: AsOfMatch,
) -> Result<Self> {
let on = using_keys
.into_iter()
.map(|key| {
let left = Self::normalize(&self.plan, key.clone())?;
let right = Self::normalize(&right, key)?;
Ok((Expr::Column(left), Expr::Column(right)))
})
.collect::<Result<_>>()?;
self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using)
}

fn asof_join_with_constraint(
self,
right: LogicalPlan,
on: Vec<(Expr, Expr)>,
match_condition: AsOfMatch,
join_constraint: JoinConstraint,
) -> Result<Self> {
let normalize = |expr, schema: &DFSchema| {
normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[])
};
let on = on
.into_iter()
.map(|(left, right_expr)| {
Ok((
normalize(left, self.plan.schema())?,
normalize(right_expr, right.schema())?,
))
})
.collect::<Result<_>>()?;
let match_condition = AsOfMatch {
left: normalize(match_condition.left, self.plan.schema())?,
op: match_condition.op,
right: normalize(match_condition.right, right.schema())?,
};
Ok(Self::new(LogicalPlan::AsOfJoin(AsOfJoin::try_new(
self.plan,
Arc::new(right),
on,
match_condition,
join_constraint,
)?)))
}

pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result<Column> {
if column.relation.is_some() {
// column is already normalized
Expand Down Expand Up @@ -1776,6 +1838,43 @@ pub fn build_join_schema(
dfschema.with_functional_dependencies(func_dependencies)
}

/// Creates the schema for a left-preserving ASOF join.
///
/// `ON` emits all left fields followed by nullable right fields. `USING` emits
/// each equality key once by omitting the corresponding right field.
pub fn build_asof_join_schema(
left: &DFSchema,
right: &DFSchema,
on: &[(Expr, Expr)],
join_constraint: JoinConstraint,
) -> Result<DFSchema> {
let omitted_right_indices = if join_constraint == JoinConstraint::Using {
on.iter()
.map(|(_, right_expr)| {
let column = right_expr.get_as_join_column().ok_or_else(|| {
plan_datafusion_err!("ASOF USING keys must be columns")
})?;
right.index_of_column(column)
})
.collect::<Result<HashSet<_>>>()?
} else {
HashSet::new()
};

let full_schema = build_join_schema(left, right, &JoinType::Left)?;
let left_len = left.fields().len();
let fields = full_schema
.iter()
.enumerate()
.filter(|(index, _)| {
*index < left_len || !omitted_right_indices.contains(&(*index - left_len))
})
.map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field)))
.collect();
DFSchema::new_with_metadata(fields, full_schema.metadata().clone())?
.with_functional_dependencies(left.functional_dependencies().clone())
}

/// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise
/// conflict with the columns from the other.
/// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying
Expand Down
23 changes: 19 additions & 4 deletions datafusion/expr/src/logical_plan/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,10 @@ use std::collections::HashMap;
use std::fmt;

use crate::{
Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join,
Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort,
Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values,
Window, expr_vec_fmt,
Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter,
Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition,
Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest,
Values, Window, expr_vec_fmt,
};

use crate::dml::CopyTo;
Expand Down Expand Up @@ -493,6 +493,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
"Filter": format!("{}", filter_expr)
})
}
LogicalPlan::AsOfJoin(AsOfJoin {
on,
match_condition,
join_constraint,
..
}) => {
let join_expr: Vec<String> =
on.iter().map(|(l, r)| format!("{l} = {r}")).collect();
json!({
"Node Type": "AsOf Join",
"Join Constraint": format!("{join_constraint:?}"),
"Join Keys": join_expr.join(", "),
"Match Condition": match_condition.to_string(),
})
}
LogicalPlan::Repartition(Repartition {
partitioning_scheme,
..
Expand Down
16 changes: 8 additions & 8 deletions datafusion/expr/src/logical_plan/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ pub mod tree_node;

pub use builder::{
LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE,
build_join_schema, requalify_sides_if_needed, table_scan, union,
wrap_projection_for_join_if_necessary,
build_asof_join_schema, build_join_schema, requalify_sides_if_needed, table_scan,
union, wrap_projection_for_join_if_necessary,
};
pub use ddl::{
CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction,
Expand All @@ -41,12 +41,12 @@ pub use dml::{
WriteOp,
};
pub use plan::{
Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn,
EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join,
JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection,
RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan,
Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union,
Unnest, Values, Window, projection_schema,
Aggregate, Analyze, AsOfJoin, AsOfMatch, ColumnUnnestList, DescribeTable, Distinct,
DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter,
Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
Projection, RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort,
StringifiedPlan, Subquery, SubqueryAlias, TableScan, TableScanBuilder,
ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema,
};
pub use statement::{
Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement,
Expand Down
Loading
Loading