peter-toth commented on code in PR #58424: URL: https://github.com/apache/spark/pull/58424#discussion_r3999288633
########## sql/core/src/main/scala/org/apache/spark/sql/execution/RewriteSelfJoinInequalityToAggregate.scala: ########## @@ -0,0 +1,569 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.IN_SUBQUERY +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +/** + * Rewrites supported uncorrelated IN-subquery inequality self-joins into + * `GROUP BY + HAVING MIN(neq) <> MAX(neq)`, avoiding the self-join cross-product. + * + * Supports a direct self-join (Pattern A') and a self-join nested under an outer inner join + * (Pattern A2, where only the self-join child becomes an Aggregate). Unsupported and correlated + * shapes fail closed. + * + * Runs in `extendedOperatorOptimizationRules`, before `RewritePredicateSubquery` turns the + * predicate subquery into a semi/anti/existence join, so it only sees the uncorrelated + * `InSubquery` shape. + * + * Controlled by `spark.sql.optimizer.rewriteSelfJoinInequalityToAggregate.enabled` + * (default false, opt-in). + */ +object RewriteSelfJoinInequalityToAggregate extends Rule[LogicalPlan] with PredicateHelper { + + private val MinNeqAliasName = "_rewrite_selfjoin_inequality_min" + private val MaxNeqAliasName = "_rewrite_selfjoin_inequality_max" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.REWRITE_SELF_JOIN_INEQUALITY_TO_AGGREGATE_ENABLED)) { + return plan + } + + // Fail closed on correlated subqueries: `lq.children` holds the outer references this rule + // does not remap. + plan.transformAllExpressionsWithPruning(_.containsPattern(IN_SUBQUERY)) { + case in @ InSubquery(_, lq: ListQuery) if lq.children.isEmpty => + rewriteSubqueryPlan(lq.plan) match { + case Some(newSub) => in.copy(query = lq.copy(plan = newSub)) + case None => in + } + } + } + + // ============================================================================ + // Shared helpers + // ============================================================================ + + /** + * Build `Filter(min <> max, Aggregate(equiKeys, child))`, taking MIN and MAX over the neq column. + * `MIN(neqCol) <> MAX(neqCol)` is true exactly when the group holds two or more distinct non-null + * values -- the same test as `COUNT(DISTINCT neqCol) > 1`, but avoids the distinct-dedup + * aggregation stages and supports partial aggregation. + * + * The `IsNotNull(equiKeys)` filter preserves the equi-join's NULL semantics: `=` never matches a + * NULL key, but GROUP BY would fold all NULL keys into one group that can leak NULL into a + * `NOT IN`. The neq column needs no filter -- MIN/MAX ignore NULL, and a group with fewer than + * two non-null values has `min = max` (or both NULL, which makes `<>` NULL), so `<>` is never + * true for it and the group is dropped. + */ + private def buildAggregateHavingMultipleDistinct( + equiKeys: Seq[Attribute], + neqCol: Attribute, + child: LogicalPlan): LogicalPlan = { + val minAlias = Alias(Min(neqCol).toAggregateExpression(), MinNeqAliasName)() Review Comment: **Finding 5.** `Min`/`Max` buffer the neq column at its own type, and a var-length buffer takes the aggregate out of hash aggregation. `Min.aggBufferAttributes` is `AttributeReference("min", child.dataType)` (`Min.scala:54,56`), so for a STRING or BINARY neq column the buffer schema is var-length. `Aggregate.supportsHashAggregate` requires `isAggregateBufferMutable` (`basicLogicalOperators.scala:1321`), and `UnsafeRow.isMutable` admits only primitive, decimal, interval and nano-timestamp physical types. `Min`/`Max` are `DeclarativeAggregate`, so `supportsObjectHashAggregate` is false as well, and `AggUtils.scala:79-115` falls through to `SortAggregateExec`. Measured at this head. Same schema, same data, same query, only the neq column's type changed. `v INT`: HashAggregate(keys=[k], functions=[min(v), max(v)]) +- Exchange hashpartitioning(k, 5) +- HashAggregate(keys=[k], functions=[partial_min(v), partial_max(v)]) `v STRING`, and `v BINARY` is identical: SortAggregate(key=[k], functions=[min(v), max(v)]) +- Sort [k ASC NULLS FIRST], false, 0 +- Exchange hashpartitioning(k, 5) +- SortAggregate(key=[k], functions=[partial_min(v), partial_max(v)]) +- Sort [k ASC NULLS FIRST], false, 0 Two full sorts the INT case does not pay, against a rewrite-off plan that is a `BroadcastHashJoin` with no exchange at all. `COUNT(DISTINCT v) > 1` did not have this shape: `RewriteDistinctAggregates` groups by `(k, v)` with long buffers, so it stayed on `HashAggregate` for every type the allowlist admits. Both types are on the allowlist and the suite exercises STRING (`RewriteSelfJoinInequalityToAggregateSuite:1098`), so this is a live path. Two things follow. The doc at `:75-76` reads as an unconditional improvement over `COUNT(DISTINCT)` - "avoids the distinct-dedup aggregation stages and supports partial aggregation". For a var-length neq column the trade goes the other way. Worth a sentence here, and in the PR description's profitability section, since the benchmark matrix was measured on a fixed-width neq column. If you want the rewrite to stay on `HashAggregate`, the neq column needs its own gate - `Aggregate.isAggregateBufferMutable(StructType(Seq(StructField("min", neqType))))`, or simply dropping StringType/BinaryType from the neq side while keeping them for equi keys. That is the same split the ordering point at [r3998997764](https://github.com/apache/spark/pull/58424#discussion_r3998997764) needs: `isSafeComparisonGroupingType` is one predicate serving two roles, and the equi keys need grouping equality while the neq column needs ordering equality and a fixed-width buffer. -- 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]
