adriangb commented on code in PR #25767:
URL: https://github.com/apache/datafusion/pull/25767#discussion_r4299222874
##########
datafusion/expr/src/logical_plan/builder.rs:
##########
@@ -1497,6 +1512,197 @@ impl LogicalPlanBuilder {
)
}
+ /// Build an `INTERSECT ALL` plan, preserving the multiplicity of each row.
+ ///
+ /// Each distinct row is returned `min(m, n)` times, where `m` and `n` are
+ /// its counts on the left and right. `count` must be the `count` aggregate
+ /// function and `range` the `range` scalar function: rows are counted on
+ /// each side and then expanded back with `unnest(range(..))`.
+ ///
+ /// Rows are compared in the same type that `=` compares them in. If the
+ /// type of a column differs between the inputs, both are cast to that type
+ /// before the rows are counted, and the result has that type instead of
the
+ /// type of the left input. Values that are equal after the cast, such as
the
+ /// strings `'1'` and `'01'` and the integer `1`, count as the same row.
+ pub fn intersect_all(
+ left_plan: LogicalPlan,
+ right_plan: LogicalPlan,
+ count: &Arc<AggregateUDF>,
+ range: &Arc<ScalarUDF>,
+ ) -> Result<LogicalPlan> {
+ Self::multiset_set_operation(left_plan, right_plan, count, range,
JoinType::Inner)
+ }
+
+ /// Build an `EXCEPT ALL` plan, subtracting matching row multiplicities.
+ ///
+ /// Each distinct row is returned `m - n` times when `m > n`. `count` and
+ /// `range` are as for [`Self::intersect_all`], and so is the handling of
+ /// inputs whose column types differ.
+ pub fn except_all(
+ left_plan: LogicalPlan,
+ right_plan: LogicalPlan,
+ count: &Arc<AggregateUDF>,
+ range: &Arc<ScalarUDF>,
+ ) -> Result<LogicalPlan> {
+ Self::multiset_set_operation(left_plan, right_plan, count, range,
JoinType::Left)
+ }
+
+ fn multiset_set_operation(
+ left_plan: LogicalPlan,
+ right_plan: LogicalPlan,
+ count: &Arc<AggregateUDF>,
+ range: &Arc<ScalarUDF>,
+ join_type: JoinType,
+ ) -> Result<LogicalPlan> {
+ let left_len = left_plan.schema().fields().len();
+ let right_len = right_plan.schema().fields().len();
+ if left_len != right_len {
+ return plan_err!(
+ "INTERSECT/EXCEPT query must have the same number of columns.
Left is {left_len} and right is {right_len}."
+ );
+ }
+ let from_right = left_plan
+ .schema()
+ .fields()
+ .iter()
+ .zip(right_plan.schema().fields())
+ .map(|(left, right)| {
+ join_type == JoinType::Inner
+ && left.is_nullable()
+ && !right.is_nullable()
+ && left.data_type() == right.data_type()
+ })
+ .collect::<Vec<_>>();
+ // The join below compares the grouped rows with `=`, which casts a
+ // mismatched pair of columns to a common type. Two distinct groups of
+ // one side can be equal in that type, and their counts would then be
+ // matched one by one instead of added. Cast first, so that the rows
+ // are grouped in the type they are compared in.
+ let (left_plan, right_plan) = coerce_set_operation_inputs(left_plan,
right_plan)?;
+ let left_columns = left_plan.schema().columns();
+ let right_columns = right_plan.schema().columns();
+ // The synthetic columns get names that differ from each other and from
+ // every input column, so the sides are only requalified when the
+ // user's own columns conflict.
+ let mut base = "__datafusion_set_operation".to_string();
+ let name_in_use = |name: &str| {
+ [left_plan.schema(), right_plan.schema()]
+ .iter()
+ .any(|schema| schema.fields().iter().any(|field| field.name()
== name))
+ };
+ let suffixes = ["_left_count", "_right_count", "_copies"];
+ while suffixes
+ .iter()
+ .any(|suffix| name_in_use(&format!("{base}{suffix}")))
+ {
+ base.push('_');
+ }
+ let [left_count, right_count, copies] =
+ suffixes.map(|suffix| format!("{base}{suffix}"));
+ // Count each distinct row on both sides. Grouping treats NULLs, and
+ // 0.0 and -0.0, as equal, like the rest of the set operations.
+ let count_rows = |plan: LogicalPlan, columns: &[Column], name: &str| {
+ let count_star =
Expr::AggregateFunction(AggregateFunction::new_udf(
+ Arc::clone(count),
+ vec![lit(COUNT_STAR_EXPANSION)],
+ false,
+ None,
+ vec![],
+ None,
+ ));
+ LogicalPlanBuilder::from(plan)
+ .aggregate(
+ columns.iter().cloned().map(Expr::Column),
+ vec![count_star.alias(name)],
+ )?
+ .build()
+ };
+ let left_plan = count_rows(left_plan, &left_columns, &left_count)?;
+ let right_plan = count_rows(right_plan, &right_columns, &right_count)?;
+ let (left_builder, right_builder, requalified) =
requalify_sides_if_needed(
+ LogicalPlanBuilder::from(left_plan),
+ LogicalPlanBuilder::from(right_plan),
+ )?;
+ let left_plan = left_builder.build()?;
+ let right_plan = right_builder.build()?;
+ // Requalifying can also rename columns (`x` becomes `x:1` when two
+ // relations on one side both have an `x`), so read the join keys from
+ // the schemas the join actually sees.
+ let mut left_join_columns = left_plan.schema().columns();
+ let mut right_join_columns = right_plan.schema().columns();
+ let left_count = Expr::Column(left_join_columns.pop().unwrap());
+ let right_count = Expr::Column(right_join_columns.pop().unwrap());
+ let join_keys = left_join_columns
+ .iter()
+ .cloned()
+ .zip(right_join_columns.iter().cloned())
+ .collect();
+ // The output keeps the left input's qualifiers and names, so a query
+ // can keep referring to them after the set operation.
+ let mut projection = left_columns
+ .iter()
+ .cloned()
+ .zip(left_join_columns)
+ .zip(right_join_columns)
+ .zip(from_right)
+ .zip(left_plan.schema().fields())
+ .map(
+ |((((column, left_column), right_column), from_right), field)|
{
+ let expr = if from_right {
+ let target_field = Arc::new(
+ Field::new(&column.name,
field.data_type().clone(), false)
+ .with_metadata(field.metadata().clone()),
+ );
+ Expr::Cast(Cast::new_from_field(
+ Box::new(Expr::Column(right_column)),
+ target_field,
+ ))
+ } else if requalified {
+ Expr::Column(left_column)
+ } else {
+ return Expr::Column(column);
+ };
+ expr.alias_qualified(column.relation, column.name)
+ },
+ )
+ .collect::<Vec<_>>();
+ let joined = LogicalPlanBuilder::from(left_plan).join_detailed(
+ right_plan,
+ join_type,
+ join_keys,
+ None,
+ NullEquality::NullEqualsNull,
+ )?;
+ let (joined, repeat) = if join_type == JoinType::Inner {
+ let repeat = when(left_count.clone().lt(right_count.clone()),
left_count)
+ .otherwise(right_count)?;
+ (joined, repeat)
+ } else {
+ let unmatched = right_count.clone().is_null();
+ let joined = joined.filter(
+ unmatched
+ .clone()
+ .or(left_count.clone().gt(right_count.clone())),
+ )?;
+ let repeat = when(unmatched, left_count.clone())
+ .otherwise(left_count - right_count)?;
+ (joined, repeat)
+ };
+ // Emit each distinct row `repeat` times.
+ projection.push(
+ Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::clone(range),
+ vec![repeat],
+ ))
+ .alias(&copies),
+ );
+ joined
+ .project(projection)?
+ .unnest_column(Column::from_name(copies))?
Review Comment:
Non-blocking: `unnest(range(n))` builds one list of `n` elements for each
distinct row before it expands it. For a row with a very large count (for
example 100M copies of one value) this is one large allocation. Can we add a
short note here, so a later change can replace it with a streaming expansion if
needed?
##########
datafusion/sql/src/set_expr.rs:
##########
@@ -178,4 +180,21 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}
}
+
+ fn functions_for_set_operation(
+ &self,
+ ) -> Result<(
+ Arc<datafusion_expr::AggregateUDF>,
+ Arc<datafusion_expr::ScalarUDF>,
+ )> {
+ let (Some(count), Some(range)) = (
+ self.context_provider.get_aggregate_meta("count"),
+ self.context_provider.get_function_meta("range"),
Review Comment:
The lookup is by name. If a user registers their own `range` (or `count`)
UDF, `INTERSECT ALL` silently uses it and can return wrong counts. Can we take
the functions directly from the crates (for example `count_udaf()` /
`range_udf()` passed in by `SessionState`) or check the returned function's
type, instead of the name in the registry?
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -1963,7 +1963,16 @@ impl DataFrame {
pub fn intersect(self, dataframe: DataFrame) -> Result<DataFrame> {
let left_plan = self.plan;
let right_plan = dataframe.plan;
- let plan = LogicalPlanBuilder::intersect(left_plan, right_plan, true)?;
+ let (Some(count), Some(range)) = (
+ self.session_state.aggregate_functions().get("count"),
+ self.session_state.scalar_functions().get("range"),
Review Comment:
`range` is in `datafusion-functions-nested`, which is optional
(`nested_expressions` feature). With `default-features = false` this now
returns a planning error, where `main` returns a (wrong) result. Please add
this case to the 56.0.0 upgrade guide note. An alternative that needs no nested
functions: expand with a join against a generated series, or with a small
dedicated `LogicalPlan` node, but I am OK with the documented limit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
