diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedColumnPruning.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedColumnPruning.scala index 6515ab99fb7..d5c81166eb0 100644 --- a/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedColumnPruning.scala +++ b/backends-clickhouse/src/main/scala/org/apache/gluten/extension/ExtendedColumnPruning.scala @@ -21,7 +21,6 @@ import org.apache.gluten.config.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.expressions.aggregate.AggregateFunction 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._ @@ -30,16 +29,14 @@ import org.apache.spark.sql.catalyst.trees.AlwaysProcess import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import scala.collection.mutable - object ExtendedGeneratorNestedColumnAliasing { def unapply(plan: LogicalPlan): Option[LogicalPlan] = plan match { case pj @ Project(projectList, f @ Filter(condition, g: Generate)) if canPruneGenerator(g.generator) && GlutenConfig.get.enableExtendedColumnPruning && (SQLConf.get.nestedPruningOnExpressions || SQLConf.get.nestedSchemaPruningEnabled) => - val attrToExtractValues = - getAttributeToExtractValues(projectList ++ g.generator.children :+ condition, Seq.empty) + val attrToExtractValues = NestedColumnAliasing + .getAttributeToExtractValues(projectList ++ g.generator.children :+ condition, Seq.empty) if (attrToExtractValues.isEmpty) { return None } @@ -63,10 +60,9 @@ object ExtendedGeneratorNestedColumnAliasing { return Some(pushedThrough) } - // In spark3.2, we could not reuse [[NestedColumnAliasing.getAttributeToExtractValues]] - // which only accepts 2 arguments. Instead we redefine it in current file to avoid moving - // this rule to gluten-shims - attrToExtractValuesOnGenerator = getAttributeToExtractValues( + // Use a Gluten-specific extractor to collect nested `GetStructField`s over generator + // outputs; the vanilla extractor only surfaces top-level [[ExtractValue]]s. + attrToExtractValuesOnGenerator = NestedColumnAliasing.getAttributeToExtractValues( attrToExtractValuesOnGenerator.flatMap(_._2).toSeq, Seq.empty, collectNestedGetStructFields) @@ -140,109 +136,6 @@ object ExtendedGeneratorNestedColumnAliasing { None } - /** - * Returns two types of expressions: - * - Root references that are individually accessed - * - [[GetStructField]] or [[GetArrayStructFields]] on top of other [[ExtractValue]]s or special - * expressions. - */ - private def collectRootReferenceAndExtractValue(e: Expression): Seq[Expression] = e match { - case _: AttributeReference => Seq(e) - case GetStructField(_: ExtractValue | _: AttributeReference, _, _) => Seq(e) - case GetArrayStructFields( - _: MapValues | _: MapKeys | _: ExtractValue | _: AttributeReference, - _, - _, - _, - _) => - Seq(e) - case es if es.children.nonEmpty => es.children.flatMap(collectRootReferenceAndExtractValue) - case _ => Seq.empty - } - - /** - * Creates a map from root [[Attribute]]s to non-redundant nested [[ExtractValue]]s. Nested field - * accessors of `exclusiveAttrs` are not considered in nested fields aliasing. - */ - private def getAttributeToExtractValues( - exprList: Seq[Expression], - exclusiveAttrs: Seq[Attribute], - extractor: (Expression) => Seq[Expression] = collectRootReferenceAndExtractValue) - : Map[Attribute, Seq[ExtractValue]] = { - - val nestedFieldReferences = new mutable.ArrayBuffer[ExtractValue]() - val otherRootReferences = new mutable.ArrayBuffer[AttributeReference]() - exprList.foreach { - e => - extractor(e).foreach { - // we can not alias the attr from lambda variable whose expr id is not available - case ev: ExtractValue if ev.find(_.isInstanceOf[NamedLambdaVariable]).isEmpty => - if (ev.references.size == 1) { - nestedFieldReferences.append(ev) - } - case ar: AttributeReference => otherRootReferences.append(ar) - case _ => // ignore - } - } - val exclusiveAttrSet = AttributeSet(exclusiveAttrs ++ otherRootReferences) - - // Remove cosmetic variations when we group extractors by their references - nestedFieldReferences - .filter(!_.references.subsetOf(exclusiveAttrSet)) - .groupBy(_.references.head.canonicalized.asInstanceOf[Attribute]) - .flatMap { - case (attr: Attribute, nestedFields: collection.Seq[ExtractValue]) => - // Check if `ExtractValue` expressions contain any aggregate functions in their tree. - // Those that do should not have an alias generated as it can lead to pushing the - // aggregate down into a projection. - def containsAggregateFunction(ev: ExtractValue): Boolean = - ev.find(_.isInstanceOf[AggregateFunction]).isDefined - - // Remove redundant [[ExtractValue]]s if they share the same parent nest field. - // For example, when `a.b` and `a.b.c` are in project list, we only need to alias `a.b`. - // Because `a.b` requires all of the inner fields of `b`, we cannot prune `a.b.c`. - val dedupNestedFields = nestedFields - .filter { - // See [[collectExtractValue]]: we only need to deal with [[GetArrayStructFields]] and - // [[GetStructField]] - case e @ (_: GetStructField | _: GetArrayStructFields) => - val child = e.children.head - nestedFields.forall(f => child.find(_.semanticEquals(f)).isEmpty) - case _ => true - } - .distinct - // Discard [[ExtractValue]]s that contain aggregate functions. - .filterNot(containsAggregateFunction) - - // If all nested fields of `attr` are used, we don't need to introduce new aliases. - // By default, the [[ColumnPruning]] rule uses `attr` already. - // Note that we need to remove cosmetic variations first, so we only count a - // nested field once. - val numUsedNestedFields = dedupNestedFields - .map(_.canonicalized) - .distinct - .map(nestedField => totalFieldNum(nestedField.dataType)) - .sum - if (dedupNestedFields.nonEmpty && numUsedNestedFields < totalFieldNum(attr.dataType)) { - Some((attr, dedupNestedFields.toSeq)) - } else { - None - } - } - } - - /** - * Return total number of fields of this type. This is used as a threshold to use nested column - * pruning. It's okay to underestimate. If the number of reference is bigger than this, the parent - * reference is used instead of nested field references. - */ - private def totalFieldNum(dataType: DataType): Int = dataType match { - case StructType(fields) => fields.map(f => totalFieldNum(f.dataType)).sum - case ArrayType(elementType, _) => totalFieldNum(elementType) - case MapType(keyType, valueType, _) => totalFieldNum(keyType) + totalFieldNum(valueType) - case _ => 1 // UDT and others - } - private def replaceGetStructField( g: GetStructField, input: Seq[Attribute], diff --git a/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/CHColumnarWrite.scala b/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/CHColumnarWrite.scala index b2669cadf28..606ded91422 100644 --- a/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/CHColumnarWrite.scala +++ b/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/CHColumnarWrite.scala @@ -143,7 +143,7 @@ case class HadoopMapReduceAdapter(sparkCommitter: HadoopMapReduceCommitProtocol) description: WriteJobDescription): (String, String) = { val stageDir = newTaskAttemptTempPath(description.path) - if (isBucketWrite(description)) { + if (description.bucketSpec.isDefined) { val filePart = getFilename(taskContext, FileNameSpec("", "")) val fileSuffix = CreateFileNameSpec(taskContext, description).suffix (stageDir, s"${filePart}_${FileNamePlaceHolder.BUCKET}$fileSuffix") @@ -152,13 +152,6 @@ case class HadoopMapReduceAdapter(sparkCommitter: HadoopMapReduceCommitProtocol) (stageDir, filename) } } - - private def isBucketWrite(desc: WriteJobDescription): Boolean = { - // In Spark 3.2, bucketSpec is not defined, instead, it uses bucketIdExpression. - val bucketSpecField: Field = desc.getClass.getDeclaredField("bucketSpec") - bucketSpecField.setAccessible(true) - bucketSpecField.get(desc).asInstanceOf[Option[_]].isDefined - } } case class NativeFileWriteResult(filename: String, partition_id: String, record_count: Long) { diff --git a/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/datasources/v1/FakeRowOutputWriter.scala b/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/datasources/v1/FakeRowOutputWriter.scala index 17cb40f147d..12754b51f98 100644 --- a/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/datasources/v1/FakeRowOutputWriter.scala +++ b/backends-clickhouse/src/main/scala/org/apache/spark/sql/execution/datasources/v1/FakeRowOutputWriter.scala @@ -44,8 +44,7 @@ class FakeRowOutputWriter(datasourceJniWrapper: Option[CHDatasourceJniWrapper], datasourceJniWrapper.foreach(_.close()) } - // Do NOT add override keyword for compatibility on spark 3.1. - def path(): String = { + override def path(): String = { outputPath } }