kosiew commented on code in PR #23691:
URL: https://github.com/apache/datafusion/pull/23691#discussion_r3655775434
##########
datafusion/optimizer/src/extract_leaf_expressions.rs:
##########
@@ -757,6 +757,66 @@ fn try_push_input(
split_and_push_projection(proj, alias_generator)
}
+/// Returns `true` if merging the extracted `pairs` (and the pass-through
+/// `columns_needed`) into `child` would inline a `KeepInPlace` definition into
+/// more than one reference site — e.g. a struct-returning UDF `f(c)` behind
both
+/// `f(c)['a']` and `f(c)['b']`, or behind `f(c)['a']` and a bare `f(c)`.
+///
+/// `CommonSubexprEliminate` hoists such an expression into a shared column so
it
+/// runs once; [`build_extraction_projection_impl`] then resolves that column
+/// back to its definition and re-duplicates it. Mirrors the
+/// `optimize_projections` merge guard (#8296): a compute-once expression used
+/// more than once stays in place.
+fn merge_would_duplicate_kept_expr(
+ pairs: &[(Expr, String)],
+ columns_needed: &IndexSet<Column>,
+ child: &Projection,
+) -> bool {
+ // Columns whose `KeepInPlace` definition must stay put; inlining one would
+ // duplicate a compute-once expression. Plain columns, pushable exprs, and
+ // literals are all cheap to duplicate, so only `KeepInPlace` counts.
+ let kept_columns: std::collections::HashSet<String> = child
+ .schema
+ .iter()
+ .zip(child.expr.iter())
+ .filter_map(|((qualifier, field), expr)| {
+ if matches!(expr.placement(), ExpressionPlacement::KeepInPlace) {
+ Some(Column::from((qualifier, field)).flat_name())
+ } else {
+ None
+ }
+ })
+ .collect();
+ if kept_columns.is_empty() {
+ return false;
+ }
+
+ // Count every site that references a kept column. A pair that references
the
+ // same kept column twice still inlines one copy, so count each pair at
most
+ // once per column.
+ let mut usage: HashMap<String, usize> = HashMap::new();
+ for (expr, _alias) in pairs {
+ let mut seen: std::collections::HashSet<String> =
+ std::collections::HashSet::new();
+ for col in expr.column_refs() {
+ let name = col.flat_name();
+ if kept_columns.contains(&name) && seen.insert(name.clone()) {
+ *usage.entry(name).or_insert(0) += 1;
+ }
+ }
+ }
+ // Pass-through / pre-existing-alias inputs are additional reference sites:
+ // a bare `f(c)` alongside `f(c)['a']` also pins the compute-once column,
so
+ // counting only `pairs` under-counted and re-duplicated it (#23655).
+ for col in columns_needed {
Review Comment:
One thing worth noting here is that the guard counts `columns_needed` as
additional reference sites. However, `columns_needed` also includes
dependencies of extracted expressions because `add_extracted` inserts each
extracted expression's `column_refs`.
This means a single extracted expression such as `get_field(__common_expr_1,
'a')` can be counted once through `pairs` and again through `columns_needed`,
even when there is no bare `__common_expr_1` output.
The conservative behaviour is safe and may be necessary with the current
merge path, since merging can keep the child `__common_expr_1` projection while
also inlining its definition into the extracted field access.
If this is relaxed later, please separate genuine parent pass-through
references from extraction-only dependencies only after confirming that the
merge path cannot retain and inline the same `KeepInPlace` expression twice.
Please also keep the new bare-reference case covered.
--
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]
