This is an automated email from the ASF dual-hosted git repository.
taiyangli pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 2d03b003c1 [GLUTEN-7868][CH] Nested column pruning for
Project(Filter(Generate)) (#7869)
2d03b003c1 is described below
commit 2d03b003c1b0cf18faae0154e5370a1bd3e829e5
Author: 李扬 <[email protected]>
AuthorDate: Tue Nov 12 16:16:16 2024 +0800
[GLUTEN-7868][CH] Nested column pruning for Project(Filter(Generate))
(#7869)
* save test configs
* wip
* finish dev
* fix import
* fix style
* fix style
* fix failed uts
* fix failec uts
* fix style
---
.../gluten/backendsapi/clickhouse/CHRuleApi.scala | 1 +
.../ExtendedGeneratorNestedColumnAliasing.scala | 126 +++++++++++++++++++++
.../hive/GlutenClickHouseHiveTableSuite.scala | 51 ++++++++-
.../scala/org/apache/gluten/GlutenConfig.scala | 10 ++
4 files changed, 187 insertions(+), 1 deletion(-)
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHRuleApi.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHRuleApi.scala
index 4107844f32..ccb124b613 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHRuleApi.scala
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHRuleApi.scala
@@ -57,6 +57,7 @@ private object CHRuleApi {
injector.injectResolutionRule(spark => new
RewriteToDateExpresstionRule(spark))
injector.injectResolutionRule(spark => new
RewriteDateTimestampComparisonRule(spark))
injector.injectOptimizerRule(spark => new
CommonSubexpressionEliminateRule(spark))
+ injector.injectOptimizerRule(spark => new
ExtendedGeneratorNestedColumnAliasing(spark))
injector.injectOptimizerRule(spark =>
CHAggregateFunctionRewriteRule(spark))
injector.injectOptimizerRule(_ => CountDistinctWithoutExpand)
injector.injectOptimizerRule(_ => EqualToRewrite)
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedGeneratorNestedColumnAliasing.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedGeneratorNestedColumnAliasing.scala
new file mode 100644
index 0000000000..a97e625ae6
--- /dev/null
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedGeneratorNestedColumnAliasing.scala
@@ -0,0 +1,126 @@
+/*
+ * 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.GlutenConfig
+
+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.optimizer.GeneratorNestedColumnAliasing.canPruneGenerator
+import org.apache.spark.sql.catalyst.optimizer.NestedColumnAliasing
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.AlwaysProcess
+import org.apache.spark.sql.internal.SQLConf
+
+// ExtendedGeneratorNestedColumnAliasing process Project(Filter(Generate)),
+// which is ignored by vanilla spark in optimization rule: ColumnPruning
+class ExtendedGeneratorNestedColumnAliasing(spark: SparkSession)
+ extends Rule[LogicalPlan]
+ with Logging {
+
+ override def apply(plan: LogicalPlan): LogicalPlan =
+ plan.transformWithPruning(AlwaysProcess.fn) {
+ case pj @ Project(projectList, f @ Filter(condition, g: Generate))
+ if canPruneGenerator(g.generator) &&
+ GlutenConfig.getConf.enableExtendedGeneratorNestedColumnAliasing &&
+ (SQLConf.get.nestedPruningOnExpressions ||
SQLConf.get.nestedSchemaPruningEnabled) =>
+ val attrToExtractValues =
NestedColumnAliasing.getAttributeToExtractValues(
+ projectList ++ g.generator.children :+ condition,
+ Seq.empty)
+ if (attrToExtractValues.isEmpty) {
+ pj
+ } else {
+ val generatorOutputSet = AttributeSet(g.qualifiedGeneratorOutput)
+ val (_, attrToExtractValuesNotOnGenerator) =
+ attrToExtractValues.partition {
+ case (attr, _) =>
+ attr.references.subsetOf(generatorOutputSet)
+ }
+
+ val pushedThrough = rewritePlanWithAliases(pj,
attrToExtractValuesNotOnGenerator)
+ pushedThrough
+ }
+ case p =>
+ p
+ }
+
+ private def rewritePlanWithAliases(
+ plan: LogicalPlan,
+ attributeToExtractValues: Map[Attribute, Seq[ExtractValue]]):
LogicalPlan = {
+ val attributeToExtractValuesAndAliases =
+ attributeToExtractValues.map {
+ case (attr, evSeq) =>
+ val evAliasSeq = evSeq.map {
+ ev =>
+ val fieldName = ev match {
+ case g: GetStructField => g.extractFieldName
+ case g: GetArrayStructFields => g.field.name
+ }
+ ev -> Alias(ev, s"_extract_$fieldName")()
+ }
+
+ attr -> evAliasSeq
+ }
+
+ val nestedFieldToAlias =
attributeToExtractValuesAndAliases.values.flatten.map {
+ case (field, alias) => field.canonicalized -> alias
+ }.toMap
+
+ // A reference attribute can have multiple aliases for nested fields.
+ val attrToAliases =
+
AttributeMap(attributeToExtractValuesAndAliases.mapValues(_.map(_._2)).toSeq)
+
+ plan match {
+ // Project(Filter(Generate))
+ case p @ Project(projectList, child)
+ if child
+ .isInstanceOf[Filter] &&
child.asInstanceOf[Filter].child.isInstanceOf[Generate] =>
+ val f = child.asInstanceOf[Filter]
+ val g = f.child.asInstanceOf[Generate]
+
+ val newProjectList =
NestedColumnAliasing.getNewProjectList(projectList, nestedFieldToAlias)
+ val newCondition = getNewExpression(f.condition, nestedFieldToAlias)
+ val newGenerator = getNewExpression(g.generator,
nestedFieldToAlias).asInstanceOf[Generator]
+
+ val tmpG = NestedColumnAliasing
+ .replaceWithAliases(g, nestedFieldToAlias, attrToAliases)
+ .asInstanceOf[Generate]
+ val newG = Generate(
+ newGenerator,
+ tmpG.unrequiredChildIndex,
+ tmpG.outer,
+ tmpG.qualifier,
+ tmpG.generatorOutput,
+ tmpG.children.head)
+ val newF = Filter(newCondition, newG)
+ val newP = Project(newProjectList, newF)
+ newP
+ case _ => plan
+ }
+ }
+
+ private def getNewExpression(
+ expr: Expression,
+ nestedFieldToAlias: Map[Expression, Alias]): Expression = {
+ expr.transform {
+ case f: ExtractValue if nestedFieldToAlias.contains(f.canonicalized) =>
+ nestedFieldToAlias(f.canonicalized).toAttribute
+ }
+ }
+}
diff --git
a/backends-clickhouse/src/test/scala/org/apache/gluten/execution/hive/GlutenClickHouseHiveTableSuite.scala
b/backends-clickhouse/src/test/scala/org/apache/gluten/execution/hive/GlutenClickHouseHiveTableSuite.scala
index 8d311614c7..ff2d13996d 100644
---
a/backends-clickhouse/src/test/scala/org/apache/gluten/execution/hive/GlutenClickHouseHiveTableSuite.scala
+++
b/backends-clickhouse/src/test/scala/org/apache/gluten/execution/hive/GlutenClickHouseHiveTableSuite.scala
@@ -17,7 +17,7 @@
package org.apache.gluten.execution.hive
import org.apache.gluten.GlutenConfig
-import
org.apache.gluten.execution.{GlutenClickHouseWholeStageTransformerSuite,
ProjectExecTransformer, TransformSupport}
+import org.apache.gluten.execution.{FileSourceScanExecTransformer,
GlutenClickHouseWholeStageTransformerSuite, ProjectExecTransformer,
TransformSupport}
import org.apache.gluten.test.AllDataTypesWithComplexType
import org.apache.gluten.utils.UTSystemParameters
@@ -28,6 +28,7 @@ import
org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import
org.apache.spark.sql.execution.datasources.v2.clickhouse.ClickHouseConfig
import org.apache.spark.sql.hive.HiveTableScanExecTransformer
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
import org.apache.hadoop.fs.Path
@@ -1450,4 +1451,52 @@ class GlutenClickHouseHiveTableSuite
spark.sql("DROP TABLE test_tbl_7054")
}
+ test("Nested column pruning for Project(Filter(Generate))") {
+ spark.sql("drop table if exists aj")
+ spark.sql(
+ """
+ |CREATE TABLE if not exists aj (
+ | country STRING,
+ | event STRUCT<time:BIGINT, lng:BIGINT, lat:BIGINT, net:STRING,
+ | log_extra:MAP<STRING, STRING>, event_id:STRING,
event_info:MAP<STRING, STRING>>
+ |)
+ |USING orc
+ """.stripMargin)
+
+ spark.sql("""
+ |INSERT INTO aj VALUES
+ | ('USA', named_struct('time', 1622547800, 'lng', -122,
'lat', 37, 'net',
+ | 'wifi', 'log_extra', map('key1', 'value1'), 'event_id',
'event1',
+ | 'event_info', map('tab_type', '5', 'action', '13'))),
+ | ('Canada', named_struct('time', 1622547801, 'lng', -79,
'lat', 43, 'net',
+ | '4g', 'log_extra', map('key2', 'value2'), 'event_id',
'event2',
+ | 'event_info', map('tab_type', '4', 'action', '12')))
+ """.stripMargin)
+
+ val df =
+ spark.sql("""
+ | SELECT * FROM (
+ | SELECT
+ | game_name,
+ | CASE WHEN
+ | event.event_info['tab_type'] IN (5) THEN '1' ELSE
'0' END AS entrance
+ | FROM aj
+ | LATERAL VIEW
explode(split(nvl(event.event_info['game_name'],'0'),','))
+ | as game_name
+ | WHERE event.event_info['action'] IN (13)
+ |) WHERE game_name = 'xxx'
+ """.stripMargin)
+
+ val scan = df.queryExecution.executedPlan.collect {
+ case scan: FileSourceScanExecTransformer => scan
+ }.head
+
+ val schema = scan.schema
+ assert(schema.size == 1)
+ val fieldType = schema.fields.head.dataType.asInstanceOf[StructType]
+ assert(fieldType.size == 1)
+
+ spark.sql("drop table if exists aj")
+ }
+
}
diff --git a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
index 3e491eb275..6579e30dab 100644
--- a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
+++ b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
@@ -107,6 +107,9 @@ class GlutenConfig(conf: SQLConf) extends Logging {
def enableCountDistinctWithoutExpand: Boolean =
conf.getConf(ENABLE_COUNT_DISTINCT_WITHOUT_EXPAND)
+ def enableExtendedGeneratorNestedColumnAliasing: Boolean =
+ conf.getConf(ENABLE_EXTENDED_GENERATOR_NESTED_COLUMN_ALIASING)
+
def veloxOrcScanEnabled: Boolean =
conf.getConf(VELOX_ORC_SCAN_ENABLED)
@@ -1929,6 +1932,13 @@ object GlutenConfig {
.booleanConf
.createWithDefault(false)
+ val ENABLE_EXTENDED_GENERATOR_NESTED_COLUMN_ALIASING =
+ buildConf("spark.gluten.sql.extendedGeneratorNestedColumnAliasing")
+ .internal()
+ .doc("Do nested column aliasing for Project(Filter(Generator))")
+ .booleanConf
+ .createWithDefault(true)
+
val COLUMNAR_VELOX_BLOOM_FILTER_EXPECTED_NUM_ITEMS =
buildConf("spark.gluten.sql.columnar.backend.velox.bloomFilter.expectedNumItems")
.internal()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
