diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala index d63928527df..703060201e7 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala @@ -118,6 +118,7 @@ object VeloxRuleApi { injector.injectPostTransform(_ => EnsureLocalSortRequirements) injector.injectPostTransform(_ => EliminateLocalSort) injector.injectPostTransform(_ => CollapseProjectExecTransformer) + injector.injectPostTransform(c => LazyAggregateExpandRule(c.session)) injector.injectPostTransform(c => FlushableHashAggregateRule.apply(c.session)) injector.injectPostTransform(_ => CollectLimitTransformerRule()) injector.injectPostTransform(_ => CollectTailTransformerRule()) diff --git a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala index e411c9d1904..d1aa35b6896 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala @@ -62,6 +62,9 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) { def enableVeloxFlushablePartialAggregation: Boolean = getConf(VELOX_FLUSHABLE_PARTIAL_AGGREGATION_ENABLED) + def enableVeloxLazyAggregateExpand: Boolean = + getConf(VELOX_LAZY_AGGREGATE_EXPAND_ENABLED) + def enableBroadcastBuildRelationInOffheap: Boolean = getConf(VELOX_BROADCAST_BUILD_RELATION_USE_OFFHEAP) @@ -431,6 +434,21 @@ object VeloxConfig extends ConfigRegistry { .booleanConf .createWithDefault(true) + val VELOX_LAZY_AGGREGATE_EXPAND_ENABLED = + buildConf("spark.gluten.sql.columnar.backend.velox.lazyAggregateExpand.enabled") + .doc( + "Experimental. For aggregation over grouping sets (rollup/cube), aggregate at the " + + "finest grain below the Expand operator first, then expand only the intermediate " + + "aggregation states and merge them before shuffle. This avoids aggregating one copy " + + "of every input row per grouping set and is beneficial when the number of distinct " + + "full-grouping-key combinations is much smaller than the input row count. Relies on " + + "flushable partial aggregation to stay adaptive on high-cardinality grouping keys; " + + "ignored when spark.gluten.sql.columnar.backend.velox.flushablePartialAggregation" + + "=false." + ) + .booleanConf + .createWithDefault(false) + val MAX_PARTIAL_AGGREGATION_MEMORY = buildConf("spark.gluten.sql.columnar.backend.velox.maxPartialAggregationMemory") .doc( diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/LazyAggregateExpandRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/LazyAggregateExpandRule.scala new file mode 100644 index 00000000000..123af209dab --- /dev/null +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/LazyAggregateExpandRule.scala @@ -0,0 +1,355 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.config.VeloxConfig +import org.apache.gluten.execution._ + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.types._ + +/** + * For aggregation over grouping sets (rollup/cube), Spark expands every input row once per grouping + * set before the partial aggregation, so the partial aggregate consumes and hashes + * `input rows * number of grouping sets` rows: + * + * partial aggregate <- expand <- child + * + * When the number of distinct full-grouping-key combinations is much smaller than the input row + * count, it is cheaper to aggregate at the finest grain once, expand only the intermediate + * aggregation states, and merge the expanded states before shuffle: + * + * partial-merge aggregate <- expand (over aggregation buffers) <- partial aggregate <- child + * + * The pre-shuffle partial-merge stage collapses duplicated coarse-grained groups locally so the + * rewrite does not increase shuffle volume (see the ClickHouse backend's lazy expand and its + * high-cardinality regression, GLUTEN-7986, for why this stage is required). + * + * Both new aggregates rely on Velox's flushable-aggregation machinery: if the input has too many + * distinct full-grouping-key combinations, the finest-grain aggregate abandons itself and streams + * rows through in intermediate format, and the merge stage over non-raw input abandons to an + * identity pass-through. The rewrite is therefore disabled when flushable partial aggregation is + * disabled. + * + * Rewrite invariants: + * - The rewritten sub-plan's output attributes equal the original partial aggregate's output, so + * no operator above the matched aggregate needs adjustment. + * - The new expand's output is exactly `grouping attributes ++ aggregation buffer attributes` in + * the original order. The partial-merge aggregate binds its buffer inputs by name and position + * against `child.output.drop(groupingExpressions.size)`, so this ordering is load-bearing. + * - Aggregate filters are only evaluated in the finest-grain (raw input) aggregate. The + * partial-merge copies drop them, mirroring Spark's AggUtils.mayRemoveAggFilters. + */ +case class LazyAggregateExpandRule(session: SparkSession) extends Rule[SparkPlan] with Logging { + + override def apply(plan: SparkPlan): SparkPlan = { + if ( + !VeloxConfig.get.enableVeloxLazyAggregateExpand || + !VeloxConfig.get.enableVeloxFlushablePartialAggregation + ) { + return plan + } + plan.transformUp { + case agg: RegularHashAggregateExecTransformer if isEligibleAggregate(agg) => + agg.child match { + case expand: ExpandExecTransformer => + rewrite(agg, expand, preProject = None, preFilter = None).getOrElse(agg) + case project @ ProjectExecTransformer(_, expand: ExpandExecTransformer) => + rewrite(agg, expand, preProject = Some(project), preFilter = None).getOrElse(agg) + case filter @ FilterExecTransformer(_, expand: ExpandExecTransformer) => + rewrite(agg, expand, preProject = None, preFilter = Some(filter)).getOrElse(agg) + case _ => agg + } + } + } + + // Matches only Partial-mode regular aggregates with a zero buffer offset. The offset check + // keeps the rule idempotent for measure-less aggregates (whose mode list is empty): every + // partial-merge aggregate this rule emits carries offset >= 1, while Spark plans partial + // aggregates with offset 0. + private def isEligibleAggregate(agg: RegularHashAggregateExecTransformer): Boolean = { + agg.initialInputBufferOffset == 0 && + agg.groupingExpressions.forall(_.isInstanceOf[Attribute]) && + agg.aggregateExpressions.forall(_.mode == Partial) && + agg.aggregateExpressions.forall(isSupportedAggregateExpression) && + !hasUnsafeFloatingPointAggregate(agg.aggregateExpressions) + } + + private def isSupportedAggregateExpression(aggExpr: AggregateExpression): Boolean = { + if (aggExpr.filter.isDefined || aggExpr.isDistinct) { + return false + } + aggExpr.aggregateFunction match { + case s: Sum => !s.prettyName.equals("try_sum") + case a: Average => !a.prettyName.equals("try_avg") + case _: Count => true + case _: Min => true + case _: Max => true + case _ => false + } + } + + // The rewrite reorders how partial states are merged, which can change the result of + // floating-point sum/avg bitwise. Apply the same policy as FlushableHashAggregateRule. + private def hasUnsafeFloatingPointAggregate(aggExprs: Seq[AggregateExpression]): Boolean = { + if (VeloxConfig.get.floatingPointMode == "loose") { + return false + } + + def isFloatingPointType(dataType: DataType): Boolean = { + dataType == DoubleType || dataType == FloatType + } + + aggExprs.exists { + aggExpr => + aggExpr.aggregateFunction match { + case s: Sum => isFloatingPointType(s.child.dataType) + case a: Average => isFloatingPointType(a.child.dataType) + case _ => false + } + } + } + + private def rewrite( + agg: RegularHashAggregateExecTransformer, + expand: ExpandExecTransformer, + preProject: Option[ProjectExecTransformer], + preFilter: Option[FilterExecTransformer]): Option[SparkPlan] = { + val numKeys = agg.groupingExpressions.length + val expandChildOutput = expand.child.output + + // A Partial aggregate's result expressions are its grouping attributes followed by the + // flattened aggregation buffer attributes. Anything else means an unexpected plan shape. + val numBufferAttributes = + agg.aggregateExpressions.map(_.aggregateFunction.aggBufferAttributes.length).sum + if ( + agg.resultExpressions.length != numKeys + numBufferAttributes || + !agg.resultExpressions.forall(_.isInstanceOf[Attribute]) || + !agg.resultExpressions + .take(numKeys) + .zip(agg.groupingExpressions) + .forall { case (result, key) => result.toAttribute.semanticEquals(key.toAttribute) } + ) { + logDebug(s"Lazy expand: unexpected partial aggregate output shape: ${agg.resultExpressions}") + return None + } + val bufferAttributes = agg.resultExpressions.drop(numKeys).map(_.toAttribute) + + // A pull-out pre-projection between the aggregate and the expand computes aggregate inputs + // (e.g. `_pre_1 = coalesce(a * b, 0)`) from columns that pass through the expand unchanged. + // It can be re-grounded onto the expand's child iff its computed expressions reference only + // pre-expand columns. + val preProjectAliases = preProject.map(_.projectList.collect { case a: Alias => a }) + if ( + !preProject.forall( + _.projectList.forall { + case _: Attribute => true + // Non-deterministic expressions must keep their original per-expanded-row evaluation; + // moving them below the expand would share one draw across all grouping sets. + case a: Alias => a.child.deterministic && resolvableFrom(a.references, expandChildOutput) + case _ => false + }) + ) { + logDebug("Lazy expand: pre-project is not re-groundable onto the expand's child") + return None + } + + // Aggregate inputs must come from columns that pass through the expand unchanged (or from + // the re-grounded pre-projection). This also rejects the look-alike Expand produced by + // RewriteDistinctAggregates, whose aggregate functions reference expand-created attributes. + val aggregateInputCandidates = + expandChildOutput ++ preProjectAliases.getOrElse(Seq.empty).map(_.toAttribute) + if ( + !agg.aggregateExpressions.forall( + ae => resolvableFrom(ae.aggregateFunction.references, aggregateInputCandidates)) + ) { + logDebug("Lazy expand: aggregate inputs are not pass-through columns of the expand") + return None + } + + // A filter between the aggregate and the expand may only reference grouping columns; it then + // filters whole (group, grouping set) rows and can equivalently run above the new expand. + if ( + !preFilter.forall( + f => + f.condition.deterministic && + resolvableFrom(f.condition.references, agg.groupingExpressions.map(_.toAttribute))) + ) { + logDebug("Lazy expand: filter references non-grouping columns of the expand") + return None + } + + // Maps each expand output attribute to the pre-expand attribute that passes through in that + // slot. Slots that are literal-only in every projection (grouping id, grouping position, + // constant grouping keys) have no mapping and are re-attached in the new expand as-is. + val replaceMap = buildReplaceAttributeMap(expand) + val bottomGroupingKeys = agg.groupingExpressions + .map(_.toAttribute) + .flatMap(attr => findReplacement(attr, replaceMap)) + .distinct + + // A keyless finest-grain aggregate would emit one row on empty input, producing spurious + // grand-total rows where Spark returns none. Non-atomic key types are excluded because the + // new expand would need typed null literals for them, which is unaudited. + if ( + bottomGroupingKeys.isEmpty || + !bottomGroupingKeys.forall(key => isSupportedGroupingKeyType(key.dataType)) + ) { + logDebug(s"Lazy expand: unsupported finest-grain grouping keys: $bottomGroupingKeys") + return None + } + + val reGroundedPreProject = preProject.map { + project => + val reGrounded = + ProjectExecTransformer(expandChildOutput ++ preProjectAliases.get, expand.child) + reGrounded.copyTagsFrom(project) + reGrounded + } + val bottomChild = reGroundedPreProject.getOrElse(expand.child) + + // Flushable, so Velox can abandon the aggregation when the finest grain barely reduces the + // row count; a regular aggregate here would hash and spill the whole input on + // high-cardinality keys. + val bottomAggregate = FlushableHashAggregateExecTransformer( + requiredChildDistributionExpressions = None, + groupingExpressions = bottomGroupingKeys, + aggregateExpressions = agg.aggregateExpressions, + aggregateAttributes = agg.aggregateAttributes, + initialInputBufferOffset = 0, + resultExpressions = bottomGroupingKeys ++ bufferAttributes, + child = bottomChild + ) + bottomAggregate.copyTagsFrom(agg) + + val newExpandOutput = agg.resultExpressions.map(_.toAttribute) + val newExpandProjections = + buildPostExpandProjections(expand.projections, expand.output, newExpandOutput) + val newExpand = ExpandExecTransformer(newExpandProjections, newExpandOutput, bottomAggregate) + newExpand.copyTagsFrom(expand) + + val newPreFilter = preFilter.map { + filter => + val newFilter = filter.copy(child = newExpand) + newFilter.copyTagsFrom(filter) + newFilter + } + val mergeChild = newPreFilter.getOrElse(newExpand) + + // Deliberately Regular: FlushableHashAggregateRule runs next and converts this merge stage + // to flushable (it walks down from the shuffle and stops here, never reaching the bottom + // aggregate, which is why the bottom one is emitted flushable directly above). + val mergeAggregate = RegularHashAggregateExecTransformer( + requiredChildDistributionExpressions = agg.requiredChildDistributionExpressions, + groupingExpressions = agg.groupingExpressions, + aggregateExpressions = + agg.aggregateExpressions.map(_.copy(mode = PartialMerge, filter = None)), + aggregateAttributes = agg.aggregateAttributes, + initialInputBufferOffset = numKeys, + resultExpressions = agg.resultExpressions, + child = mergeChild + ) + mergeAggregate.copyTagsFrom(agg) + + val newNodes: Seq[SparkPlan] = + reGroundedPreProject.toSeq ++ Seq(bottomAggregate, newExpand) ++ + newPreFilter.toSeq :+ mergeAggregate + if (!newNodes.forall(passesNativeValidation)) { + logDebug("Lazy expand: native validation failed for the rewritten plan; keeping original") + return None + } + logDebug(s"Lazy expand rewrote aggregate over expand: $mergeAggregate") + Some(mergeAggregate) + } + + // The new expand must emit typed null literals for excluded grouping keys; restrict to types + // whose null literals are known to round-trip through the native ExpandRel. + private def isSupportedGroupingKeyType(dataType: DataType): Boolean = { + dataType match { + // Referenced by type name: TimestampNTZType is private[sql] in Spark 3.3. + case dt if dt.typeName == "timestamp_ntz" => true + case BooleanType | StringType | DateType | TimestampType | BinaryType => + true + case _: NumericType => true + case _ => false + } + } + + private def resolvableFrom(references: AttributeSet, candidates: Seq[Attribute]): Boolean = { + references.forall(ref => candidates.exists(_.semanticEquals(ref))) + } + + private def findReplacement( + attribute: Attribute, + replaceMap: Map[Attribute, Attribute]): Option[Attribute] = { + replaceMap.collectFirst { case (k, v) if k.semanticEquals(attribute) => v } + } + + private def buildReplaceAttributeMap(expand: ExpandExecTransformer): Map[Attribute, Attribute] = { + val passThroughBySlot = expand.output.indices.map { + i => + expand.projections.collectFirst { + case projection if projection(i).isInstanceOf[Attribute] => + projection(i).asInstanceOf[Attribute] + } + } + expand.output + .zip(passThroughBySlot) + .collect { case (out, Some(passThrough)) => out -> passThrough } + .toMap + } + + // Rebuilds the expand projections against the finest-grain aggregate's output: slots that + // existed in the original expand keep their per-projection expression (pass-through attribute + // or literal), and aggregation buffer attributes pass through every projection unchanged. + private def buildPostExpandProjections( + originalProjections: Seq[Seq[Expression]], + originalOutput: Seq[Attribute], + newOutput: Seq[Attribute]): Seq[Seq[Expression]] = { + originalProjections.map { + projection => + newOutput.map { + attr => + val index = originalOutput.indexWhere(_.semanticEquals(attr)) + if (index != -1) { + projection(index) + } else { + attr + } + } + } + } + + private def passesNativeValidation(plan: SparkPlan): Boolean = { + plan match { + case validatable: ValidatablePlan => + try { + validatable.doValidate().ok() + } catch { + case e: Exception => + logDebug(s"Lazy expand: validation threw for ${plan.nodeName}: ${e.getMessage}") + false + } + case _ => true + } + } +} diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/LazyAggregateExpandSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/LazyAggregateExpandSuite.scala new file mode 100644 index 00000000000..355d78d60fc --- /dev/null +++ b/backends-velox/src/test/scala/org/apache/gluten/execution/LazyAggregateExpandSuite.scala @@ -0,0 +1,307 @@ +/* + * 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.gluten.execution + +import org.apache.gluten.config.VeloxConfig + +import org.apache.spark.SparkConf +import org.apache.spark.sql.DataFrame + +class LazyAggregateExpandSuite extends VeloxWholeStageTransformerSuite { + override protected val resourcePath: String = "/tpch-data-parquet" + override protected val fileFormat: String = "parquet" + + override def beforeAll(): Unit = { + super.beforeAll() + createTPCHNotNullTables() + } + + override protected def sparkConf: SparkConf = { + super.sparkConf + .set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.ColumnarShuffleManager") + .set("spark.sql.shuffle.partitions", "2") + .set("spark.memory.offHeap.size", "2g") + .set("spark.unsafe.exceptionOnMemoryLeak", "true") + .set(VeloxConfig.VELOX_LAZY_AGGREGATE_EXPAND_ENABLED.key, "true") + } + + // The rewrite leaves an ExpandExecTransformer whose child is the FLUSHABLE finest-grain + // aggregate, and FlushableHashAggregateRule must have converted the merge stage above it. + // A regular stage on either side would reintroduce the high-cardinality regression that hit + // the ClickHouse backend (GLUTEN-7986), so both are asserted. + private def checkLazyExpand(df: DataFrame, fired: Boolean = true): Unit = { + val plans = getExecutedPlan(df) + if (fired) { + val lazyExpands = plans.collect { + case e: ExpandExecTransformer + if e.child.isInstanceOf[FlushableHashAggregateExecTransformer] => + e + } + assert( + lazyExpands.nonEmpty, + s"expected lazy expand to fire but it did not:\n${df.queryExecution.executedPlan}") + val flushableMergeExists = plans.exists { + case a: FlushableHashAggregateExecTransformer => + lazyExpands.exists(e => a.child.find(_ eq e).isDefined) + case _ => false + } + assert( + flushableMergeExists, + s"expected a flushable merge stage above the expand:\n${df.queryExecution.executedPlan}") + } else { + val lazyExpands = plans.collect { + case e: ExpandExecTransformer if e.child.isInstanceOf[HashAggregateExecBaseTransformer] => + e + } + assert( + lazyExpands.isEmpty, + s"expected lazy expand to not fire but it did:\n${df.queryExecution.executedPlan}") + } + } + + test("rollup with sum/count/min/max") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_suppkey), count(l_suppkey), " + + "min(l_suppkey), max(l_suppkey) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("cube with grouping and grouping_id") { + runQueryAndCompare( + "select l_orderkey, l_partkey, grouping(l_orderkey), grouping(l_partkey), " + + "grouping_id(l_orderkey, l_partkey), sum(l_suppkey) from lineitem " + + "group by cube(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey, 3, 4, 5") { + df => checkLazyExpand(df) + } + } + + test("avg buffers are merged, not averaged") { + // Groups with very different sizes: any average-of-averages shortcut diverges from vanilla. + withTable("t_lazy_avg") { + sql("create table t_lazy_avg (k1 int, k2 int, v int) using parquet") + sql( + "insert into t_lazy_avg values " + + "(1, 1, 1000), (1, 2, 0), (1, 2, 0), (1, 2, 0), (1, 2, 0), (1, 2, 0), " + + "(2, 1, 10), (2, 1, 20), (2, 2, 500)") + runQueryAndCompare( + "select k1, k2, avg(v), count(*) from t_lazy_avg " + + "group by rollup(k1, k2) order by k1, k2") { + df => checkLazyExpand(df) + } + } + } + + test("genuine null key vs rolled-up null") { + withTable("t_lazy_nulls") { + sql("create table t_lazy_nulls (k1 string, k2 string, v int) using parquet") + sql( + "insert into t_lazy_nulls values " + + "('a', null, 1), ('a', 'x', 2), (null, 'x', 3), (null, null, 4), ('b', 'x', 5)") + runQueryAndCompare( + "select k1, k2, grouping(k1), grouping(k2), grouping_id(k1, k2), sum(v), count(*) " + + "from t_lazy_nulls group by rollup(k1, k2) " + + "order by k1, k2, 3, 4, 5") { + df => checkLazyExpand(df) + } + } + } + + test("decimal sum buffer (sum, isEmpty) binds through the expand") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(cast(l_extendedprice as decimal(12, 2))) " + + "from lineitem group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("pre-projected aggregate input (TPC-DS q67 shape)") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_suppkey * l_linenumber + 1) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("aggregate over a grouping key") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_orderkey) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("empty input returns no rows") { + // AQE's empty-relation propagation collapses the whole plan on empty input, which would + // make the plan assertion vacuous; disable it so the rewritten plan is observable. + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withTable("t_lazy_empty") { + sql("create table t_lazy_empty (k1 int, k2 int, v int) using parquet") + runQueryAndCompare( + "select k1, k2, sum(v), count(*) from t_lazy_empty " + + "group by grouping sets ((k1, k2), ())") { + df => checkLazyExpand(df) + } + } + } + } + + test("degenerate grouping sets without attribute keys are not rewritten") { + // A keyless finest-grain aggregate would emit a spurious row on empty input. AQE disabled + // for the same reason as the empty-input test. + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withTable("t_lazy_degenerate") { + sql("create table t_lazy_degenerate (v int) using parquet") + runQueryAndCompare( + "select sum(v), count(*) from t_lazy_degenerate group by grouping sets ((), ())") { + df => checkLazyExpand(df, fired = false) + } + } + } + } + + test("duplicate grouping sets keep duplicate result rows") { + runQueryAndCompare( + "select l_orderkey, sum(l_suppkey) from lineitem " + + "where l_orderkey < 100 " + + "group by grouping sets ((l_orderkey), (l_orderkey)) " + + "order by l_orderkey") { + df => checkLazyExpand(df) + } + } + + test("single count distinct rewrites the dedup aggregate") { + runQueryAndCompare( + "select l_orderkey, l_partkey, count(distinct l_suppkey), sum(l_linenumber) " + + "from lineitem group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("multiple count distinct stays correct") { + // RewriteDistinctAggregates produces a second, look-alike Expand that must not be rewritten; + // correctness is what matters here, the fire state depends on the surviving plan shape. + runQueryAndCompare( + "select l_orderkey, count(distinct l_suppkey), count(distinct l_partkey) " + + "from lineitem group by rollup(l_orderkey) " + + "order by l_orderkey") { _ => } + } + + test("aggregate filter clause is not rewritten") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_suppkey) filter (where l_linenumber > 2) " + + "from lineitem group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df, fired = false) + } + } + + test("floating point sum is not rewritten in strict mode") { + withSQLConf(VeloxConfig.FLOATING_POINT_MODE.key -> "strict") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(cast(l_quantity as double)) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df, fired = false) + } + } + } + + test("filter on grouping keys between aggregate and expand") { + // The predicate references a nulled key copy, so it cannot be pushed below the Expand and + // sits between the partial aggregate and the Expand, the rule's preFilter branch. + runQueryAndCompare( + "select * from (" + + "select l_orderkey, l_partkey, sum(l_suppkey) as s from lineitem " + + "group by cube(l_orderkey, l_partkey)) " + + "where l_orderkey is not null order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("high cardinality keys with early abandon") { + withSQLConf( + "spark.gluten.sql.columnar.backend.velox.abandonPartialAggregationMinRows" -> "10", + "spark.gluten.sql.columnar.backend.velox.abandonPartialAggregationMinPct" -> "1" + ) { + runQueryAndCompare( + "select l_orderkey, l_partkey, l_suppkey, sum(l_linenumber) from lineitem " + + "group by rollup(l_orderkey, l_partkey, l_suppkey) " + + "order by l_orderkey, l_partkey, l_suppkey") { + df => checkLazyExpand(df) + } + } + } + + test("adaptive execution disabled") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_suppkey) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + } + + test("disabled by default when flushable aggregation is off") { + withSQLConf(VeloxConfig.VELOX_FLUSHABLE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(l_suppkey) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df, fired = false) + } + } + } + + test("non-whitelisted aggregate function is not rewritten") { + runQueryAndCompare( + "select l_orderkey, l_partkey, stddev_samp(l_suppkey) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df, fired = false) + } + } + + test("floating point sum is rewritten in loose mode") { + // floatingPointMode defaults to loose, the same policy that already allows flushing float + // aggregates. Small integers cast to double keep the sums exact regardless of merge order. + runQueryAndCompare( + "select l_orderkey, l_partkey, sum(cast(l_linenumber as double)) from lineitem " + + "group by rollup(l_orderkey, l_partkey) " + + "order by l_orderkey, l_partkey") { + df => checkLazyExpand(df) + } + } + + test("non-atomic grouping key is not rewritten") { + runQueryAndCompare( + "select array(l_orderkey), l_partkey, sum(l_suppkey) from lineitem " + + "group by rollup(array(l_orderkey), l_partkey) " + + "order by 2, 1") { + df => checkLazyExpand(df, fired = false) + } + } +} diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md index 952a6c7f9c8..77b3c4df098 100644 --- a/docs/velox-configuration.md +++ b/docs/velox-configuration.md @@ -35,6 +35,7 @@ nav_order: 16 | spark.gluten.sql.columnar.backend.velox.hashProbe.bloomFilterPushdown.maxSize | 🔄 Dynamic | 0b | The maximum byte size of Bloom filter that can be generated from hash probe. When set to 0, no Bloom filter will be generated. To achieve optimal performance, this should not be too larger than the CPU cache size on the host. | | spark.gluten.sql.columnar.backend.velox.hashProbe.dynamicFilterPushdown.enabled | 🔄 Dynamic | true | Whether hash probe can generate any dynamic filter (including Bloom filter) and push down to upstream operators. | | spark.gluten.sql.columnar.backend.velox.hashShuffle.reader.streamMerge.enabled | 🔄 Dynamic | false | Enables a reader-side raw payload merge fast path for plain hash shuffle payloads within each shuffle input stream. This path merges payload buffers before Velox vectors are materialized, so it has lower per-batch overhead than generic VeloxResizeBatchesExec resizing, but it only covers plain payloads. Complex types and dictionary-encoded payloads are not merged by this path. VeloxResizeBatchesExec can still be enabled separately as a generic complement for types and encodings not covered by this fast path. If false, each hash shuffle payload is returned as its own columnar batch. | +| spark.gluten.sql.columnar.backend.velox.lazyAggregateExpand.enabled | 🔄 Dynamic | false | Experimental. For aggregation over grouping sets (rollup/cube), aggregate at the finest grain below the Expand operator first, then expand only the intermediate aggregation states and merge them before shuffle. This avoids aggregating one copy of every input row per grouping set and is beneficial when the number of distinct full-grouping-key combinations is much smaller than the input row count. Relies on flushable partial aggregation to stay adaptive on high-cardinality grouping keys; ignored when spark.gluten.sql.columnar.backend.velox.flushablePartialAggregation=false. | | spark.gluten.sql.columnar.backend.velox.loadQuantum | ⚓ Static | 256MB | Set the load quantum for velox file scan, recommend to use the default value (256MB) for performance consideration. If Velox cache is enabled, it can be 8MB at most. | | spark.gluten.sql.columnar.backend.velox.maxCoalescedBytes | ⚓ Static | 64MB | Set the max coalesced bytes for velox file scan | | spark.gluten.sql.columnar.backend.velox.maxCoalescedDistance | ⚓ Static | 512KB | Set the max coalesced distance bytes for velox file scan |