This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 6fb23c7d3d [VL] Prune map entries at the Parquet scan to the keys a
query reads (#13109)
6fb23c7d3d is described below
commit 6fb23c7d3d64e0b9e84fd53472b3aafa04da6621
Author: Rong Ma <[email protected]>
AuthorDate: Fri Oct 9 17:24:16 2026 +0100
[VL] Prune map entries at the Parquet scan to the keys a query reads
(#13109)
---
.../apache/gluten/execution/VeloxDeltaSuite.scala | 87 ++-
.../gluten/proto/RequiredSubfieldsExtension.proto | 37 ++
.../gluten/backendsapi/velox/VeloxRuleApi.scala | 1 +
.../backendsapi/velox/VeloxSparkPlanExecApi.scala | 29 +-
.../backendsapi/velox/VeloxTransformerApi.scala | 33 +-
.../org/apache/gluten/config/VeloxConfig.scala | 14 +
.../gluten/extension/ScanMapKeyPruning.scala | 424 +++++++++++++
.../gluten/execution/ScanMapKeyPruningSuite.scala | 677 +++++++++++++++++++++
.../spark/sql/execution/GlutenHiveUDFSuite.scala | 36 +-
cpp/velox/CMakeLists.txt | 1 +
.../operators/functions/overlay/GetMapValue.cc | 40 ++
.../operators/functions/overlay/GetMapValue.h | 62 ++
.../functions/overlay/RegisterFunctionOverlay.cc | 6 +
cpp/velox/substrait/SubstraitToVeloxPlan.cc | 86 ++-
cpp/velox/tests/SparkFunctionTest.cc | 33 +
docs/velox-configuration.md | 1 +
.../gluten/execution/DeltaScanTransformer.scala | 18 +-
.../org/apache/gluten/execution/DeltaSuite.scala | 2 +-
.../gluten/backendsapi/SparkPlanExecApi.scala | 11 +
.../apache/gluten/backendsapi/TransformerApi.scala | 10 +-
.../execution/BasicScanExecTransformer.scala | 25 +-
.../execution/BatchScanExecTransformer.scala | 18 +-
.../execution/FileSourceScanExecTransformer.scala | 12 +-
.../org/apache/gluten/execution/SubfieldPath.scala | 51 ++
.../gluten/expression/ExpressionConverter.scala | 7 +
.../org/apache/spark/sql/GlutenQueryTest.scala | 4 +-
26 files changed, 1696 insertions(+), 29 deletions(-)
diff --git
a/backends-velox/src-delta/test/scala/org/apache/gluten/execution/VeloxDeltaSuite.scala
b/backends-velox/src-delta/test/scala/org/apache/gluten/execution/VeloxDeltaSuite.scala
index 851cf1dcbc..3ffee52b4b 100644
---
a/backends-velox/src-delta/test/scala/org/apache/gluten/execution/VeloxDeltaSuite.scala
+++
b/backends-velox/src-delta/test/scala/org/apache/gluten/execution/VeloxDeltaSuite.scala
@@ -18,7 +18,8 @@ package org.apache.gluten.execution
import org.apache.gluten.config.VeloxDeltaConfig
-import org.apache.spark.sql.Row
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.util.SparkVersionUtil
class VeloxDeltaSuite extends DeltaSuite {
test("delta: change data feed scan offload can be disabled") {
@@ -49,4 +50,88 @@ class VeloxDeltaSuite extends DeltaSuite {
}
}
}
+
+ Seq("name", "id").foreach {
+ mode =>
+ test(s"map-key pruning is disabled under column mapping mode = $mode") {
+ withTable("delta_cm_map") {
+ spark.sql(s"""
+ |create table delta_cm_map
+ | (id bigint, m map<string, struct<s string, t
bigint>>)
+ |using delta
+ |tblproperties ("delta.columnMapping.mode" = "$mode")
+ |""".stripMargin)
+ spark.sql(
+ "insert into delta_cm_map select id, " +
+ "map('a', named_struct('s', concat('v', cast(id % 7 as string)),
't', id), " +
+ "'b', named_struct('s', 'x', 't', id * 2)) from range(0, 500)")
+ val pruningFlag =
"spark.gluten.sql.columnar.backend.velox.scanMapKeyPruningEnabled"
+ withSQLConf(pruningFlag -> "true", "spark.sql.adaptive.enabled" ->
"false") {
+ runQueryAndCompare(
+ "select sum(m['a'].t), max(m['a'].s) from delta_cm_map where
m['a'].s = 'v3'") {
+ df =>
+ val subfields = getExecutedPlan(df)
+ .collect { case scan: DeltaScanTransformer =>
scan.requiredMapSubfields }
+ .flatten
+ assert(subfields.isEmpty, s"physical column names cannot be
pruned: $subfields")
+ }
+ }
+ }
+ }
+ }
+
+ testWithMinSparkVersion("map-key pruning on delta scan with and without
deletion vector", "3.4") {
+ withTempPath {
+ p =>
+ val path = p.getCanonicalPath
+ spark
+ .range(0, 2000)
+ .selectExpr(
+ "id",
+ "map('a', named_struct('s', concat('v', cast(id % 7 as string)),
't', id), " +
+ "'b', named_struct('s', 'x', 't', id * 2), " +
+ "'c', named_struct('s', 'y', 't', id + 1)) as m"
+ )
+ .coalesce(1)
+ .write
+ .format("delta")
+ .save(path)
+ val sql = s"SELECT count(*), sum(m['a'].t), max(m['a'].s) " +
+ s"FROM delta.`$path` WHERE m['a'].s = 'v3'"
+ val pruningFlag =
"spark.gluten.sql.columnar.backend.velox.scanMapKeyPruningEnabled"
+
+ def scanSubfields(df: DataFrame): Map[String, Seq[String]] =
+ getExecutedPlan(df)
+ .collect { case scan: DeltaScanTransformer =>
scan.requiredMapSubfields }
+ .flatten
+ .toMap
+ .map { case (col, paths) => col -> paths.map(_.toString) }
+
+ withSQLConf(pruningFlag -> "true", "spark.sql.adaptive.enabled" ->
"false") {
+ runQueryAndCompare(sql) {
+ df =>
+ val subfields = scanSubfields(df)
+ assert(subfields.contains("m"), s"expected pruning on m, got
$subfields")
+ assert(subfields("m").exists(_.contains("m[\"a\"]")),
subfields("m").toString)
+ assert(!subfields("m").exists(_.contains("[\"b\"]")),
subfields("m").toString)
+ }
+ }
+
+ spark.sql(
+ s"ALTER TABLE delta.`$path` SET TBLPROPERTIES
('delta.enableDeletionVectors' = true)")
+ spark.sql(s"DELETE FROM delta.`$path` WHERE id % 10 = 0")
+ if (SparkVersionUtil.gteSpark35) {
+ withSQLConf(pruningFlag -> "true", "spark.sql.adaptive.enabled" ->
"false") {
+ runQueryAndCompare(sql) {
+ df =>
+ assert(
+ df.queryExecution.executedPlan
+ .collect { case _: DeltaScanTransformer => true }
+ .nonEmpty)
+ assert(scanSubfields(df).contains("m"))
+ }
+ }
+ }
+ }
+ }
}
diff --git
a/backends-velox/src/main/resources/org/apache/gluten/proto/RequiredSubfieldsExtension.proto
b/backends-velox/src/main/resources/org/apache/gluten/proto/RequiredSubfieldsExtension.proto
new file mode 100644
index 0000000000..576385c79c
--- /dev/null
+++
b/backends-velox/src/main/resources/org/apache/gluten/proto/RequiredSubfieldsExtension.proto
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+syntax = "proto3";
+
+package gluten;
+
+option java_package = "org.apache.gluten.proto";
+option java_multiple_files = true;
+
+// Subfields of scan output columns that the query accesses, carried in
+// ReadRel.advanced_extension.enhancement. The native reader treats a listed
+// column's set as complete: map entries under keys not listed are not
+// materialized and struct fields not listed are null. Columns without an entry
+// are read whole.
+message RequiredSubfieldsExtension {
+ // One step below a column: a struct field, or a map lookup by key.
+ message PathElement {
+ oneof element {
+ string field = 1;
+ string string_key = 2;
+ int64 long_key = 3;
+ }
+ }
+
+ // A path under a column, e.g. elements [string_key "a", field "t"] for
m["a"].t.
+ message Subfield {
+ repeated PathElement elements = 1;
+ }
+
+ message ColumnSubfields {
+ // Column and field names as the scan names them in its schema
(case-folded by the driver when
+ // case-insensitive).
+ string column = 1;
+ repeated Subfield subfields = 2;
+ }
+
+ repeated ColumnSubfields columns = 1;
+}
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 0937d5d065..63c8880e8a 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
@@ -126,6 +126,7 @@ object VeloxRuleApi {
injector.injectPostTransform(_ => PartialFallbackRules())
injector.injectPostTransform(_ => RemoveNativeWriteFilesSortAndProject())
injector.injectPostTransform(_ => PushDownFilterToScan)
+ injector.injectPostTransform(_ => ScanMapKeyPruning)
injector.injectPostTransform(_ => PushDownInputFileExpression.PostOffload)
injector.injectPostTransform(_ => EnsureLocalSortRequirements)
injector.injectPostTransform(_ => EliminateLocalSort)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
index 90dfc0b260..a7fd67ea81 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
@@ -1123,10 +1123,31 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi
with Logging {
left: ExpressionTransformer,
right: ExpressionTransformer,
original: GetMapValue): ExpressionTransformer = {
- GenericExpressionTransformer(
- ExpressionMappings.expressionsMap(classOf[ElementAt]),
- Seq(left, right),
- original)
+ // Emitted as get_map_value (Gluten function overlay) rather than rewritten
+ // to element_at: the map-only subscript is subfield-pushdown capable, so a
+ // scan remaining filter on m['k'] extracts m["k"] and keeps map-key
+ // pruning effective.
+ GenericExpressionTransformer(substraitExprName, Seq(left, right), original)
+ }
+
+ /**
+ * ElementAt on a map is emitted as get_map_value, like GetMapValue, so that
a scan remaining
+ * filter such as element_at(m, 'k').x = 1 extracts the subfield m["k"] and
map-key pruning stays
+ * effective; Velox's element_at also serves arrays and is not
subfield-pushdown capable, so it
+ * would add the bare column m and keep every entry. The two agree on maps:
the value under the
+ * key, or NULL when it is absent. In ANSI mode a missing key is an error
instead, and the mapped
+ * element_at is kept.
+ */
+ override def genElementAtTransformer(
+ substraitExprName: String,
+ left: ExpressionTransformer,
+ right: ExpressionTransformer,
+ original: ElementAt): ExpressionTransformer = {
+ if (original.left.dataType.isInstanceOf[MapType] && !original.failOnError)
{
+ GenericExpressionTransformer(ExpressionNames.GET_MAP_VALUE, Seq(left,
right), original)
+ } else {
+ GenericExpressionTransformer(substraitExprName, Seq(left, right),
original)
+ }
}
override def genStringToMapTransformer(
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala
index 04df524427..5fabab93a6 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala
@@ -18,10 +18,10 @@ package org.apache.gluten.backendsapi.velox
import org.apache.gluten.backendsapi.{BackendsApiManager, TransformerApi}
import org.apache.gluten.exception.GlutenException
-import org.apache.gluten.execution.WriteFilesExecTransformer
+import org.apache.gluten.execution.{SubfieldElement, SubfieldPath,
WriteFilesExecTransformer}
import org.apache.gluten.execution.datasource.GlutenFormatFactory
import org.apache.gluten.expression.ConverterUtils
-import org.apache.gluten.proto.{ConfigMap, IcebergReadExtension}
+import org.apache.gluten.proto.{ConfigMap, IcebergReadExtension,
RequiredSubfieldsExtension}
import org.apache.gluten.runtime.Runtimes
import org.apache.gluten.substrait.SubstraitContext
import org.apache.gluten.substrait.expression.{ExpressionBuilder,
ExpressionNode}
@@ -152,6 +152,35 @@ class VeloxTransformerApi extends TransformerApi with
Logging {
packPBMessage(extensionBuilder.build())
}
+ override def packRequiredSubfields(subfields: Map[String,
Seq[SubfieldPath]]): Any = {
+ val extension = RequiredSubfieldsExtension.newBuilder()
+ subfields.toSeq.sortBy(_._1).foreach {
+ case (column, paths) =>
+ val columnBuilder =
+
RequiredSubfieldsExtension.ColumnSubfields.newBuilder().setColumn(column)
+ paths.foreach {
+ path =>
+ require(
+ path.column == column,
+ s"Subfield path $path declared under column $column")
+ val subfield = RequiredSubfieldsExtension.Subfield.newBuilder()
+ path.elements.foreach {
+ element =>
+ val pe = RequiredSubfieldsExtension.PathElement.newBuilder()
+ element match {
+ case SubfieldElement.Field(name) => pe.setField(name)
+ case SubfieldElement.StringKey(key) => pe.setStringKey(key)
+ case SubfieldElement.LongKey(key) => pe.setLongKey(key)
+ }
+ subfield.addElements(pe)
+ }
+ columnBuilder.addSubfields(subfield)
+ }
+ extension.addColumns(columnBuilder)
+ }
+ packPBMessage(extension.build())
+ }
+
override def invalidateSQLExecutionResource(executionId: String): Unit = {
GlutenDriverEndpoint.invalidateResourceRelation(executionId)
}
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 ed2217f252..6d66080a23 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
@@ -82,6 +82,9 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) {
def veloxOrcScanEnabled: Boolean =
getConf(VELOX_ORC_SCAN_ENABLED)
+ def scanMapKeyPruningEnabled: Boolean =
+ getConf(SCAN_MAP_KEY_PRUNING_ENABLED)
+
def floatingPointMode: String = getConf(FLOATING_POINT_MODE)
def enableRewriteCastArrayToString: Boolean =
@@ -279,6 +282,17 @@ object VeloxConfig extends ConfigRegistry {
.intConf
.createWithDefault(2)
+ val SCAN_MAP_KEY_PRUNING_ENABLED =
+
buildConf("spark.gluten.sql.columnar.backend.velox.scanMapKeyPruningEnabled")
+ .doc(
+ "When true, a map column read by a native Parquet scan is declared to
Velox with only " +
+ "the constant keys the query looks up (m['k'], element_at(m, 'k'),
null checks), so " +
+ "the reader skips the other entries. Applies when every use of the
map sits in the " +
+ "filters and projections directly above the scan and the map is
dropped before any " +
+ "other operator sees it; otherwise the map is read whole. See
ScanMapKeyPruning.")
+ .booleanConf
+ .createWithDefault(false)
+
val COLUMNAR_VELOX_GLOG_VERBOSE_LEVEL =
buildConf("spark.gluten.sql.columnar.backend.velox.glogVerboseLevel")
.internal()
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala
b/backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala
new file mode 100644
index 0000000000..14f8d0a20b
--- /dev/null
+++
b/backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala
@@ -0,0 +1,424 @@
+/*
+ * 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.{BasicScanExecTransformer,
FilterExecTransformerBase, GenerateExecTransformer, ProjectExecTransformerBase,
SubfieldElement, SubfieldPath}
+import org.apache.gluten.execution.SubfieldElement.{Field, LongKey, StringKey
=> StrKey}
+import org.apache.gluten.expression.ConverterUtils
+import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{FilterExec, GenerateExec, ProjectExec,
SparkPlan}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+
+import java.nio.charset.StandardCharsets
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+/**
+ * Physical plan rule that derives Velox required subfields for map columns
read by a native Parquet
+ * scan, so that the reader prunes map entries the query never accesses.
+ *
+ * Background. Parquet stores a map as a repeated `key_value` group: one key
column chunk and one
+ * value column chunk per row group, covering every entry of every row.
Evaluating `m['k1']` on the
+ * scan output therefore decodes all entries into a MapVector, although only
one key is needed.
+ * Velox's Hive connector lets a `HiveColumnHandle` declare required
subfields, written in Velox
+ * `Subfield` syntax such as `m["k1"].f1`; `makeScanSpec` turns the map keys
they name into a key
+ * filter on the map's key stream, so the SelectiveMapColumnReader
materializes only matching
+ * entries. Struct fields not covered by any required subfield are filled with
nulls.
+ *
+ * Velox treats the declared set as complete and returns null for anything
outside it, so the
+ * declared paths must cover every read of the column. The rule therefore only
handles one plan
+ * shape, in which that is easy to show:
+ * {{{
+ * Scan [id, m] filters: m["a"].s = 'v3'
+ * -> Filter* more predicates over keyed paths
+ * -> Project* keyed paths under aliases; m
may be forwarded
+ * -> Generate* m forwarded only if still
needed above
+ * -> a Project (or Generate) that no longer outputs m end of the chain
+ * }}}
+ * Once a Project drops the attribute (and every alias that still carries map
data of it), no
+ * operator above can reference it: in a resolved plan an attribute can only
be referenced if some
+ * child outputs it. Whatever sits above the chain is therefore irrelevant and
needs no analysis.
+ * Aliases of keyed paths that leave the chain (`m['a'].t AS x`) are fine as
well: their value
+ * depends only on entries the declared paths keep.
+ *
+ * Supported, that is, pruned:
+ * - filters on constant-key lookups, `m['a'].s = 'v3'`, `m[1].t > 0`,
`element_at(m, 'a')`,
+ * including the `isnotnull(m)` Spark infers next to them. Such a filter
reaches Velox as the
+ * scan's remaining filter, whose subfields Velox adds to the declaration;
the Velox backend
+ * emits both lookups as the pushdown-capable `get_map_value`, so they
contribute `m["a"]`
+ * rather than the bare column (see
VeloxSparkPlanExecApi.genElementAtTransformer);
+ * - keyed values projected or aggregated, `SELECT m['a'].t`,
`sum(m['a'].t)`, also under a
+ * `count(*)`, a `LIMIT`, an `ORDER BY ... LIMIT`, a `UNION ALL` of such
projections, a write of
+ * such values, a UDF or script transformation applied to them, or a join
whose sides only read
+ * keyed values below the shuffle: in all of these Spark or Gluten places
the keyed access in a
+ * Project that drops the map before anything else sees it;
+ * - maps nested in structs, `c.m['a']`, including the `c.m AS _extract_m`
alias Spark's
+ * NestedColumnAliasing inserts; sibling struct fields are declared
alongside;
+ * - generators over a keyed value, `explode(m['a'])`: Gluten's
pre-projection computes the lookup
+ * below the Generate, and the Generate forwards the map only while an
operator above still
+ * needs it, so the chain may end at the Generate itself;
+ * - window partition keys, `OVER (PARTITION BY m['a'])`, which Spark's
analyzer extracts into a
+ * Project below the Window's shuffle;
+ * - null checks on keyed paths, `m['a'] IS NULL`.
+ *
+ * Not supported, that is, the map is read whole:
+ * - the map itself, or an alias still holding it, reaching an exchange, a
join, an aggregate, a
+ * union, a write, a UDTF, the fragment's output or any other operator
that is not a chain node;
+ * this includes `ORDER BY m['a']`, whose pre-projection must forward the
map through the
+ * shuffle because the Sort still outputs it;
+ * - whole-map uses anywhere on the chain, `size(m)`, `map_keys(m)`,
`explode(m)`, `m[key_col]`;
+ * - key types without a Velox subscript form, such as date or decimal, and
string keys whose
+ * bytes are not valid UTF-8.
+ *
+ * A chain ends at the first operator that is not a chain node, an exchange
included, so a scan's
+ * declaration depends only on the operators between it and that point.
Structurally identical
+ * exchange subtrees therefore always receive identical declarations, and
exchange reuse is
+ * unaffected.
+ *
+ * Example. For
+ * {{{
+ * SELECT sum(m['a'].t) FROM t WHERE m['a'].s = 'v3'
+ * }}}
+ * Spark places `m['a'].t AS _extract_t` in a Project below the aggregate, so
the chain is Scan ->
+ * Filter -> Project and the declared paths are `m["a"].s` and `m["a"].t`
(shown in Velox Subfield
+ * syntax, which SubfieldPath.toString renders; the wire format is structural,
one element per step,
+ * see RequiredSubfieldsExtension.proto).
+ */
+object ScanMapKeyPruning extends Rule[SparkPlan] {
+
+ // One element of a subfield path below the attribute, mirroring Velox
Subfield::PathElement:
+ // NestedField, StringSubscript, LongSubscript. `c.m["k1"].f1` is attribute
`c` followed by
+ // Field("m"), StrKey("k1"), Field("f1").
+ private type Elem = SubfieldElement
+
+ /** What the chain collectDeclarations collects for one scan attribute. */
+ private class CollectedPaths {
+ // Paths whose values are read; the map entries they name must be
materialized.
+ val valuePaths: mutable.LinkedHashSet[List[Elem]] =
+ mutable.LinkedHashSet[List[Elem]]()
+ // Paths referenced only by IsNull / IsNotNull; a null check reads the
null bitmap, which the
+ // reader produces for anything it materializes, so it needs no map
entries of its own.
+ val nullCheckPaths: mutable.LinkedHashSet[List[Elem]] =
+ mutable.LinkedHashSet[List[Elem]]()
+ // Set when some reference cannot be declared; the attribute is then left
alone.
+ var abandoned: Boolean = false
+ }
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!VeloxConfig.get.scanMapKeyPruningEnabled) {
+ return plan
+ }
+ // Walk down from the root, remembering the chain nodes directly above the
current node,
+ // nearest first. Any other operator restarts the chain. Declarations are
kept per scan object,
+ // since two scans of one table are distinct objects.
+ val declarationsByScan =
+ new java.util.IdentityHashMap[SparkPlan, Map[String,
Seq[SubfieldPath]]]()
+ def collectDeclarations(node: SparkPlan, chain: List[SparkPlan]): Unit =
node match {
+ case scan: BasicScanExecTransformer if isPrunableParquetScan(scan) =>
+ val mapAttrs = scan.output.filter(a => containsMapType(a.dataType) &&
a.name.nonEmpty)
+ if (mapAttrs.nonEmpty) {
+ val filters = scan.filterExprs()
+ val subfields = mapAttrs.flatMap {
+ attr =>
+ pathsToDeclare(filters, chain, attr).map(paths =>
renderSubfieldPaths(attr, paths))
+ }.toMap
+ if (subfields.nonEmpty) {
+ declarationsByScan.put(scan, subfields)
+ }
+ }
+ case _ =>
+ val next = if (isChainOperator(node)) node :: chain else Nil
+ node.children.foreach(collectDeclarations(_, next))
+ }
+ collectDeclarations(plan, Nil)
+ if (declarationsByScan.isEmpty) {
+ return plan
+ }
+
+ plan.transformUp {
+ case scan: BasicScanExecTransformer if
declarationsByScan.containsKey(scan) =>
+ val subfields = declarationsByScan.get(scan)
+ logDebug(s"Applying scan map-key pruning: $subfields")
+ val newScan = scan.withRequiredMapSubfields(subfields)
+ newScan.copyTagsFrom(scan)
+ newScan
+ }
+ }
+
+ /** A native Parquet scan that supports the declaration and has none yet. */
+ private def isPrunableParquetScan(scan: BasicScanExecTransformer): Boolean =
+ scan.supportsMapKeyPruning && scan.requiredMapSubfields.isEmpty &&
+ scan.fileFormat == ReadFileFormat.ParquetReadFormat
+
+ /**
+ * The operators a chain may consist of: Filters and Projects, which use the
input rows through
+ * their predicate or project list, and Generate, which forwards its
requiredChildOutput and uses
+ * the input only through the generator. ColumnarPartialProjectExec and
+ * ColumnarPartialGenerateExec are deliberately not among them: part of
their work is in a field
+ * the rule cannot see.
+ */
+ private def isChainOperator(node: SparkPlan): Boolean = node match {
+ case _: FilterExecTransformerBase | _: FilterExec => true
+ case _: ProjectExecTransformerBase | _: ProjectExec => true
+ case _: GenerateExecTransformer | _: GenerateExec => true
+ case _ => false
+ }
+
+ /**
+ * Follows 'attr' from its scan (whose pushed-down 'filters' are visited
first) up the chain and
+ * returns the paths to declare, or None when the attribute must stay whole.
+ *
+ * 'liveIds' holds the ExprIds that still carry map data of the attribute,
with the path from the
+ * attribute to what the id stands for: the attribute itself with an empty
path, a bare alias (`m
+ * AS mm`) with an empty path, or a struct-field alias (`c.m AS _extract_m`,
which Spark's
+ * NestedColumnAliasing inserts) with path [m]. A reference to a live id
resolves by prepending
+ * that path: `_extract_m['a']` is `c.m["a"]`. The chain is complete once no
id is live.
+ */
+ private def pathsToDeclare(
+ filters: Seq[Expression],
+ chain: List[SparkPlan],
+ attr: Attribute): Option[Seq[List[Elem]]] = {
+ val paths = new CollectedPaths
+ var liveIds: Map[ExprId, List[Elem]] = Map(attr.exprId -> Nil)
+
+ def recordUsesIn(e: Expression, nullCheck: Boolean): Unit =
+ recordUses(liveIds, paths, e, nullCheck)
+
+ // A project list decides what stays live. An entry that resolves to a
path under a live id
+ // stays live under its own id when the path is key-free and still holds a
map: a forwarded
+ // `m#3`, a bare alias `m#3 AS mm#12`, or Spark's `c#3.m AS _extract_m#20`
(recording the path
+ // instead would declare all of it and defeat pruning below). A keyed
path, or a path with no
+ // map under it (`c#3.s.v AS _extract_v#8`), is recorded and needs no
further tracking: the
+ // entry then only carries data the declared path keeps. Anything else is
a computation, a use.
+ def applyProjectList(list: Seq[NamedExpression]): Unit = {
+ val next = mutable.LinkedHashMap[ExprId, List[Elem]]()
+ list.foreach {
+ entry =>
+ val expr = entry match {
+ case Alias(child, _) => child
+ case other => other
+ }
+ resolveSubfieldPath(expr) match {
+ case Some((root, path)) if liveIds.contains(root.exprId) =>
+ val full = liveIds(root.exprId) ++ path
+ if (!hasKeyLookup(full) && containsMapType(entry.dataType)) {
+ next(entry.exprId) = full
+ } else {
+ recordPath(paths, full, nullCheck = false)
+ }
+ case _ => recordUsesIn(expr, nullCheck = false)
+ }
+ }
+ liveIds = next.toMap
+ }
+
+ // A Generate uses its input only through the generator
(requiredChildOutput is a forward),
+ // and outputs requiredChildOutput plus the generated columns. Whatever it
does not output
+ // cannot be referenced above it, so it may also end the chain.
+ def applyGenerate(node: SparkPlan, generator: Expression): Unit = {
+ recordUsesIn(generator, nullCheck = false)
+ val forwarded = node.output.map(_.exprId).toSet
+ liveIds = liveIds.filter { case (id, _) => forwarded.contains(id) }
+ }
+
+ // Filters pushed into the scan (PushDownFilterToScan) are evaluated
natively on the scan
+ // output; they read the map like the Filter above them does.
+ filters.foreach(recordUsesIn(_, nullCheck = false))
+ chain.iterator.takeWhile(_ => liveIds.nonEmpty &&
!paths.abandoned).foreach {
+ case f: FilterExecTransformerBase => recordUsesIn(f.cond, nullCheck =
false)
+ case f: FilterExec => recordUsesIn(f.condition, nullCheck = false)
+ case p: ProjectExecTransformerBase => applyProjectList(p.list)
+ case p: ProjectExec => applyProjectList(p.projectList)
+ case g: GenerateExecTransformer => applyGenerate(g, g.generator)
+ case g: GenerateExec => applyGenerate(g, g.generator)
+ }
+ // Still live when the chain ends: the attribute (or an alias holding map
data of it) reaches
+ // an operator the rule does not follow, or the fragment's output. Stay
whole.
+ if (paths.abandoned || liveIds.nonEmpty) {
+ return None
+ }
+
+ val value = paths.valuePaths.toSeq
+ // A null check needs no map entries, so it is dropped when a value path
has it as a prefix
+ // (`startsWith` compares element lists: `m` is a prefix of `m["a"].t`):
+ // WHERE m IS NOT NULL AND m['a'].t > 0 -> declare m["a"].t only;
declaring `m` would keep
+ // the whole column.
+ // WHERE m['b'] IS NULL, SELECT m['a'].t -> declare both; without
m["b"] the reader would
+ // drop key b and the check
would see null.
+ val nullOnly = paths.nullCheckPaths.toSeq.filterNot(nc =>
value.exists(_.startsWith(nc)))
+ val all = value ++ nullOnly
+ // Declare only if the reader can then skip something. A key-free path
keeps everything under
+ // it (Velox lets a shorter path dominate longer ones), so `size(c.m)`
next to `c.m['a']`, or
+ // `m IS NOT NULL` next to `m['a'] IS NULL`, reads all of the map either
way; declaring would
+ // only rewrite the scan for nothing.
+ val keyFree = all.filterNot(hasKeyLookup)
+ val prunes = all.exists(p => hasKeyLookup(p) && !keyFree.exists(prefix =>
p.startsWith(prefix)))
+ if (prunes) Some(all) else None
+ }
+
+ /**
+ * Classifies one expression: a subfield path under a live id is recorded,
anything else is
+ * decomposed into its children. The largest sub-expression that is a path
is the informative one,
+ * so the whole expression is tried first. `cardinality(m)` is not a path;
the descent reaches
+ * `m`, whose empty path means the map is used whole, and recordPath()
abandons. For
+ * `cardinality(c.m)` the descent records the key-free path `c.m`, which at
emission dominates
+ * every keyed path below it and so suppresses the declaration. 'nullCheck'
is set for the operand
+ * of IsNull / IsNotNull.
+ */
+ private def recordUses(
+ liveIds: Map[ExprId, List[Elem]],
+ paths: CollectedPaths,
+ e: Expression,
+ nullCheck: Boolean): Unit = e match {
+ case n: IsNotNull => recordUses(liveIds, paths, n.child, nullCheck = true)
+ case n: IsNull => recordUses(liveIds, paths, n.child, nullCheck = true)
+ case _ =>
+ resolveSubfieldPath(e) match {
+ case Some((root, path)) =>
+ liveIds.get(root.exprId).foreach(prefix => recordPath(paths, prefix
++ path, nullCheck))
+ case None => e.children.foreach(recordUses(liveIds, paths, _,
nullCheck = false))
+ }
+ }
+
+ /**
+ * Records a reference to 'path', or abandons the attribute when the path
cannot be declared: the
+ * empty path, i.e. the attribute itself used whole (`size(m)`, `first(m)`),
unless it is a null
+ * check; or a path naming an empty field, which the native side cannot
represent. Paths come from
+ * resolved expressions, so they fit the schema; Velox keeps everything
under them.
+ */
+ private def recordPath(paths: CollectedPaths, path: List[Elem], nullCheck:
Boolean): Unit = {
+ val emptyName = path.exists {
+ case Field(name) => name.isEmpty
+ case _ => false
+ }
+ if (emptyName || (path.isEmpty && !nullCheck)) {
+ paths.abandoned = true
+ } else if (nullCheck) {
+ paths.nullCheckPaths += path
+ } else {
+ paths.valuePaths += path
+ }
+ }
+
+ /** True if 'dt' is or contains a MapType, so an attribute of this type is a
candidate. */
+ private def containsMapType(dt: DataType): Boolean = dt match {
+ case _: MapType => true
+ case s: StructType => s.exists(f => containsMapType(f.dataType))
+ case a: ArrayType => containsMapType(a.elementType)
+ case _ => false
+ }
+
+ /**
+ * Resolves 'e' to a subfield path: an AttributeReference at the root,
followed by GetStructField
+ * and constant-key GetMapValue / ElementAt steps. Catalyst nests these
extractors outside-in, so
+ * the match recurses to the attribute and appends one element per level on
the way back, leaving
+ * the elements in root-to-leaf order. None for anything else, for example
`cardinality(m)` or a
+ * lookup with a non-constant key.
+ */
+ private def resolveSubfieldPath(e: Expression): Option[(AttributeReference,
List[Elem])] =
+ e match {
+ case a: AttributeReference => Some((a, Nil))
+ case g: GetStructField =>
+ // The schema's spelling, not the user's: with caseSensitive=false,
`c.M` resolves to field
+ // `m`, and the path must name the field as the file stores it.
+ val fieldName =
g.child.dataType.asInstanceOf[StructType](g.ordinal).name
+ resolveSubfieldPath(g.child).map { case (a, p) => (a, p :+
(Field(fieldName): Elem)) }
+ case g: GetMapValue => resolveMapLookup(g.child, g.key)
+ case ea: ElementAt if ea.left.dataType.isInstanceOf[MapType] =>
+ resolveMapLookup(ea.left, ea.right)
+ case _ => None
+ }
+
+ /** One map lookup step of resolveSubfieldPath: the path to 'map' plus the
constant 'key'. */
+ private def resolveMapLookup(
+ map: Expression,
+ key: Expression): Option[(AttributeReference, List[Elem])] = {
+ val keyType = map.dataType.asInstanceOf[MapType].keyType
+ constantKeyElement(key, keyType).flatMap(
+ k => resolveSubfieldPath(map).map { case (a, p) => (a, p :+ k) })
+ }
+
+ /**
+ * The key of a lookup into a map with keys of 'keyType', as a path element,
if the key expression
+ * is foldable (a constant under Catalyst's definition): `m['a']`,
`m[concat('a', '')]` and
+ * `m[upper('a')]` qualify, a key that depends on the input row does not.
The element kind follows
+ * the map's key type, which is what Velox checks the subfield against:
string keys become StrKey,
+ * integral keys LongKey, any other key type (date, decimal, ...) has no
Velox subscript and
+ * yields None.
+ */
+ private def constantKeyElement(key: Expression, keyType: DataType):
Option[Elem] = {
+ if (!key.foldable) {
+ return None
+ }
+ val value =
+ try key.eval(EmptyRow)
+ catch { case NonFatal(_) => return None }
+ if (value == null) {
+ return None
+ }
+ keyType match {
+ case StringType =>
+ // The key is compared as bytes by Spark and by the reader, but
travels as a string.
+ // Invalid UTF-8 (a key built from a binary cast) does not survive the
round trip, so
+ // such a key cannot be declared.
+ val bytes = value.asInstanceOf[UTF8String].getBytes
+ val rendered = value.toString
+ if (java.util.Arrays.equals(rendered.getBytes(StandardCharsets.UTF_8),
bytes)) {
+ Some(StrKey(rendered))
+ } else {
+ None
+ }
+ case ByteType | ShortType | IntegerType | LongType =>
+ Some(LongKey(value.asInstanceOf[Number].longValue()))
+ case _ => None
+ }
+ }
+
+ private def hasKeyLookup(path: Seq[Elem]): Boolean =
+ path.exists(e => e.isInstanceOf[StrKey] || e.isInstanceOf[LongKey])
+
+ /**
+ * Renders one attribute's paths as [[SubfieldPath]]s. Column and field
names go through
+ * ConverterUtils.normalizeColName, exactly like the names of the scan's
Substrait schema do when
+ * the plan is built, so the declaration and the schema always agree,
including for non-ASCII
+ * names that only Java's full lowercasing changes. Transport is structural,
so no character of a
+ * name or key needs quoting; toString on the result gives the Velox
Subfield syntax, for example
+ * `c.m["k1"].f1`. The paths are sorted by that rendering, so two scans
declaring the same set
+ * compare equal whatever order the walk found the paths in, and exchange
reuse is not lost.
+ */
+ private def renderSubfieldPaths(
+ attr: Attribute,
+ paths: Seq[List[Elem]]): (String, Seq[SubfieldPath]) = {
+ val colName = ConverterUtils.normalizeColName(attr.name)
+ val rendered = paths.map {
+ p =>
+ SubfieldPath(
+ colName,
+ p.map {
+ case Field(n) => Field(ConverterUtils.normalizeColName(n))
+ case other => other
+ })
+ }
+ (colName, rendered.sortBy(_.toString))
+ }
+}
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/ScanMapKeyPruningSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/ScanMapKeyPruningSuite.scala
new file mode 100644
index 0000000000..4c54941591
--- /dev/null
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/ScanMapKeyPruningSuite.scala
@@ -0,0 +1,677 @@
+/*
+ * 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.{GlutenConfig, VeloxConfig}
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.execution.{BaseScriptTransformationExec, SparkPlan}
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.functions.{col, lit, map, struct, sum}
+import org.apache.spark.sql.types.StringType
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Every case runs with adaptive execution off and on. Most cases are rows of
one table: a fixture,
+ * a query and the exact declaration expected on each native scan, where an
empty declaration means
+ * the scan must be left whole. The class comment of ScanMapKeyPruning lists
which shapes are
+ * supported; the cases here follow that list.
+ */
+class ScanMapKeyPruningSuite extends VeloxWholeStageTransformerSuite {
+ override protected val resourcePath: String = "/tpch-data-parquet"
+ override protected val fileFormat: String = "parquet"
+
+ override protected def sparkConf: SparkConf = super.sparkConf
+ .set("spark.unsafe.exceptionOnMemoryLeak", "true")
+
+ private val pruningFlag = VeloxConfig.SCAN_MAP_KEY_PRUNING_ENABLED.key
+ private val noBroadcast = "spark.sql.autoBroadcastJoinThreshold" -> "-1"
+
+ //
---------------------------------------------------------------------------------------------
+ // Fixtures: each writes a parquet table and hands its path to the body.
+ //
---------------------------------------------------------------------------------------------
+
+ private type Fixture = (String => Unit) => Unit
+
+ private def table(rows: Long, columns: String*): Fixture = f =>
+ withTempPath {
+ path =>
+ spark.range(0, rows).selectExpr(("id" +: columns):
_*).write.parquet(path.getCanonicalPath)
+ f(path.getCanonicalPath)
+ }
+
+ /** m: MAP<STRING, STRUCT<s STRING, t BIGINT>> with keys a, b, c. */
+ private val mapTable = table(
+ 10000,
+ "map('a', named_struct('s', concat('v', cast(id % 7 as string)), 't', id),
" +
+ "'b', named_struct('s', 'x', 't', id * 2), " +
+ "'c', named_struct('s', 'y', 't', id + 1)) as m"
+ )
+
+ /** m: MAP<STRING, BIGINT>, a map whose values are plain numbers. */
+ private val scalarMapTable = table(1000, "map('a', id, 'b', id * 2, 'c', id
+ 1) as m")
+
+ /** m: MAP<BIGINT, STRUCT<s, t>>. */
+ private val longKeyTable = table(
+ 10000,
+ "map(1L, named_struct('s', concat('v', cast(id % 7 as string)), 't', id),
" +
+ "2L, named_struct('s', 'x', 't', id * 2)) as m")
+
+ /** c: STRUCT<m: MAP<STRING, BIGINT>, v: BIGINT>. */
+ private val nestedMapTable =
+ table(1000, "named_struct('m', map('a', id, 'b', id * 2), 'v', id) as c")
+
+ /** c: STRUCT<m: MAP<STRING, BIGINT>, s: STRUCT<n: MAP<STRING, BIGINT>, v:
BIGINT>>. */
+ private val siblingMapTable = table(
+ 1000,
+ "named_struct('m', map('a', id, 'b', id * 2), " +
+ "'s', named_struct('n', map('p', id), 'v', id)) as c")
+
+ /** m: MAP<STRING, MAP<STRING, BIGINT>>. */
+ private val mapOfMapTable =
+ table(1000, "map('a', map('x', id, 'y', id * 2), 'b', map('z', id)) as m")
+
+ /** m: MAP<STRING, ARRAY<BIGINT>>. */
+ private val arrayMapTable = table(500, "map('a', array(id, id + 1), 'b',
array(id * 2)) as m")
+
+ /** m as in mapTable plus n: MAP<STRING, BIGINT>. */
+ private val twoMapTable = table(
+ 10000,
+ "map('a', named_struct('s', concat('v', cast(id % 7 as string)), 't', id),
" +
+ "'b', named_struct('s', 'x', 't', id * 2)) as m",
+ "map('x', id, 'y', id * 2) as n"
+ )
+
+ // Built from code points rather than literals or unicode escapes, which
scalastyle rejects.
+ private val umlautUpper = 0xc4.toChar // A with diaeresis
+ private val umlautLower = 0xe4.toChar // a with diaeresis
+
+ /**
+ * c: STRUCT<[umlautUpper]rger: MAP<STRING, BIGINT>, v: BIGINT>: a field
name that only full
+ * Unicode lowercasing changes, since the ASCII fold leaves the umlaut alone.
+ */
+ private val umlautFieldTable =
+ table(1000, s"named_struct('${umlautUpper}rger', map('a', id, 'b', id *
2), 'v', id) as c")
+
+ /** m: MAP<STRING, BIGINT> whose first key is a single byte that is not
valid UTF-8. */
+ private val invalidUtf8KeyTable =
+ table(100, "map(cast(x'FF' as string), id, 'a', id * 2) as m")
+
+ /** d: MAP<DATE, BIGINT>, e: MAP<DECIMAL(3, 1), BIGINT>: key types without a
Velox subscript. */
+ private val oddKeyTable = table(
+ 1000,
+ "map(DATE '2020-01-01', id, DATE '2020-01-02', id * 2) as d",
+ "map(CAST(1.5 AS DECIMAL(3, 1)), id, CAST(2.5 AS DECIMAL(3, 1)), id * 2)
as e")
+
+ //
---------------------------------------------------------------------------------------------
+ // Plan inspection.
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * The declaration of every native scan in the executed plan (query stages
and commands included),
+ * each rendered in Velox Subfield syntax. At least one native scan must
exist: a "nothing
+ * declared" assertion over a plan whose scan fell back to Spark would prove
nothing.
+ */
+ private def declaredPerScan(df: DataFrame): Seq[Map[String, Set[String]]] = {
+ val scans = getExecutedPlan(df).collect { case scan:
BasicScanExecTransformer => scan }
+ assert(scans.nonEmpty, s"expected a native
scan:\n${df.queryExecution.executedPlan}")
+ scans.map(_.requiredMapSubfields.map { case (c, paths) => c ->
paths.map(_.toString).toSet })
+ }
+
+ /** The declaration over all native scans, for single-scan plans. */
+ private def declared(df: DataFrame): Map[String, Set[String]] =
declaredPerScan(df).flatten.toMap
+
+ private def planHas(df: DataFrame, cls: Class[_ <: SparkPlan]): Boolean =
+ getExecutedPlan(df).exists(cls.isInstance)
+
+ /** Per-scan declarations compared regardless of scan order. */
+ private def multiset[T](xs: Seq[T]): Map[T, Int] = xs.groupBy(identity).map {
+ case (k, v) => k -> v.size
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // Table-driven cases.
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * One case: 'sql' over 'fixture', run with the flag on and compared with
vanilla Spark, must
+ * leave exactly 'expected' on the native scans (one entry per scan,
order-free; an empty map is a
+ * scan left whole) and, if given, must contain an operator of type
'requires' in its plan.
+ */
+ private case class Case(
+ name: String,
+ fixture: Fixture,
+ sql: String => String,
+ expected: Seq[Map[String, Set[String]]],
+ conf: Seq[(String, String)] = Nil,
+ requires: Option[Class[_ <: SparkPlan]] = None,
+ noFallBack: Boolean = true,
+ compare: Boolean = true,
+ before: () => Unit = () => ())
+
+ private def m(paths: String*): Map[String, Set[String]] = Map("m" ->
paths.toSet)
+ private def c(paths: String*): Map[String, Set[String]] = Map("c" ->
paths.toSet)
+ private val whole: Map[String, Set[String]] = Map.empty
+
+ private val cases = Seq(
+ // ----- supported shapes: the map is dropped by a Project before anything
else sees it -----
+ Case(
+ "keyed filter with aggregates over keyed values",
+ mapTable,
+ p => s"SELECT count(*), sum(m['a'].t), max(m['a'].s) FROM parquet.`$p`
WHERE m['a'].s = 'v3'",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t"))
+ ),
+ Case(
+ "filter only, map dropped above the filter",
+ mapTable,
+ p => s"SELECT id FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(m("m[\"a\"].s"))),
+ Case(
+ "count only",
+ mapTable,
+ p => s"SELECT count(*) FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(m("m[\"a\"].s"))),
+ Case(
+ "projected keyed value without aggregation",
+ mapTable,
+ p => s"SELECT id, m['a'].t FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t"))
+ ),
+ Case(
+ "several keys",
+ mapTable,
+ p => s"SELECT sum(m['a'].t + m['b'].t) FROM parquet.`$p` WHERE m['c'].s
= 'y'",
+ Seq(m("m[\"a\"].t", "m[\"b\"].t", "m[\"c\"].s"))),
+ Case(
+ "constant key expressions are evaluated",
+ mapTable,
+ p => s"SELECT sum(m[concat('a', '')].t) FROM parquet.`$p` WHERE
m[upper('a')].s = 'v3'",
+ Seq(m("m[\"a\"].t", "m[\"A\"].s"))
+ ),
+ Case(
+ "element_at with a constant key",
+ mapTable,
+ p =>
+ s"SELECT sum(element_at(m, 'a').t) FROM parquet.`$p` " +
+ s"WHERE element_at(m, 'a').s = 'v5'",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t"))
+ ),
+ Case(
+ "inferred null check next to a keyed filter is dropped",
+ mapTable,
+ p =>
+ s"SELECT count(*), sum(m['a'].t) FROM parquet.`$p` " +
+ s"WHERE m IS NOT NULL AND m['a'].s = 'v3'",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t"))
+ ),
+ Case(
+ "null check on a keyed path alone",
+ mapTable,
+ p => s"SELECT count(*) FROM parquet.`$p` WHERE m['a'] IS NULL",
+ Seq(m("m[\"a\"]"))),
+ Case(
+ "long keys",
+ longKeyTable,
+ p => s"SELECT count(*), sum(m[1].t) FROM parquet.`$p` WHERE m[1].s =
'v3'",
+ Seq(m("m[1].s", "m[1].t"))),
+ Case(
+ "scalar-valued map",
+ scalarMapTable,
+ p => s"SELECT sum(m['a']) FROM parquet.`$p` WHERE m['b'] > 10",
+ Seq(m("m[\"a\"]", "m[\"b\"]"))),
+ Case(
+ "nested map value used whole under a key",
+ mapOfMapTable,
+ p => s"SELECT sum(cardinality(m['a'])), sum(m['a']['y']) FROM
parquet.`$p`",
+ Seq(m("m[\"a\"]", "m[\"a\"][\"y\"]"))
+ ),
+ Case(
+ "map nested in a struct, through Spark's nested-column alias",
+ nestedMapTable,
+ p => s"SELECT sum(c.m['a']), sum(c.m['b']) FROM parquet.`$p`",
+ Seq(c("c.m[\"a\"]", "c.m[\"b\"]"))
+ ),
+ Case(
+ "sibling struct fields are declared alongside the keyed map",
+ siblingMapTable,
+ p => s"SELECT sum(c.m['a']), sum(c.s.v), sum(cardinality(c.s.n)) FROM
parquet.`$p`",
+ Seq(c("c.m[\"a\"]", "c.s.v", "c.s.n"))
+ ),
+ Case(
+ "case-insensitive field access renders the schema's spelling",
+ nestedMapTable,
+ p => s"SELECT sum(c.M['a']) FROM parquet.`$p`",
+ Seq(c("c.m[\"a\"]")),
+ conf = Seq("spark.sql.caseSensitive" -> "false")
+ ),
+ Case(
+ "collect limit at the root",
+ mapTable,
+ p => s"SELECT id FROM parquet.`$p` WHERE m['a'].s = 'v3' LIMIT 5",
+ Seq(m("m[\"a\"].s"))),
+ Case(
+ "take ordered and project",
+ mapTable,
+ p => s"SELECT id FROM parquet.`$p` WHERE m['a'].s = 'v3' ORDER BY id
LIMIT 5",
+ Seq(m("m[\"a\"].s"))),
+ Case(
+ "window partitioned by a keyed value",
+ scalarMapTable,
+ p => s"SELECT id, rank() OVER (PARTITION BY m['a'] ORDER BY id) AS r
FROM parquet.`$p`",
+ Seq(m("m[\"a\"]"))
+ ),
+ Case(
+ "generate over a keyed value, map still read above the generate",
+ arrayMapTable,
+ p => s"SELECT m['b'], explode(m['a']) FROM parquet.`$p`",
+ Seq(m("m[\"a\"]", "m[\"b\"]")),
+ requires = Some(classOf[GenerateExecTransformer])
+ ),
+ Case(
+ "generate over a keyed value ends the chain when nothing above needs the
map",
+ arrayMapTable,
+ p => s"SELECT explode(m['a']) FROM parquet.`$p`",
+ Seq(m("m[\"a\"]")),
+ requires = Some(classOf[GenerateExecTransformer])
+ ),
+ Case(
+ "non-ASCII struct field name folded like the schema, case-insensitive",
+ umlautFieldTable,
+ // Non-ASCII identifiers must be quoted in Spark SQL.
+ p =>
+ s"SELECT sum(c.`${umlautLower}rger`['a']) FROM parquet.`$p` " +
+ s"WHERE c.`${umlautUpper}RGER`['b'] > 10",
+ // Plain concatenation: scalastyle's parser does not accept \" inside
interpolated strings.
+ Seq(c("c." + umlautLower + "rger[\"a\"]", "c." + umlautLower +
"rger[\"b\"]")),
+ conf = Seq("spark.sql.caseSensitive" -> "false")
+ ),
+ Case(
+ "two map columns on one scan, one pruned and one read whole",
+ twoMapTable,
+ p => s"SELECT sum(m['a'].t), sum(size(n)) FROM parquet.`$p`",
+ Seq(m("m[\"a\"].t"))
+ ),
+ Case(
+ "shuffled join with keyed filters on both sides",
+ mapTable,
+ p =>
+ s"SELECT count(*) FROM parquet.`$p` t1 JOIN parquet.`$p` t2 ON t1.id =
t2.id " +
+ s"WHERE t1.m['a'].s = 'v3' AND t2.m['b'].s = 'x'",
+ Seq(m("m[\"a\"].s"), m("m[\"b\"].s")),
+ conf = Seq(noBroadcast)
+ ),
+ Case(
+ "inlined CTE read twice with different keys pushed below the join",
+ mapTable,
+ p =>
+ s"WITH c AS (SELECT id, m FROM parquet.`$p`) " +
+ s"SELECT sum(t1.m['a'].t + t2.m['b'].t) FROM c t1 JOIN c t2 ON t1.id
= t2.id",
+ Seq(m("m[\"a\"].t"), m("m[\"b\"].t")),
+ conf = Seq(noBroadcast)
+ ),
+ Case(
+ "identical pruned scans keep exchange reuse",
+ mapTable,
+ p =>
+ s"WITH c AS (SELECT id, m['a'].t AS x FROM parquet.`$p` WHERE m['a'].s
= 'v3') " +
+ s"SELECT sum(t1.x + t2.x) FROM c t1 JOIN c t2 ON t1.id = t2.id",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t")),
+ conf = Seq(noBroadcast),
+ requires = Some(classOf[ReusedExchangeExec])
+ ),
+ Case(
+ // The two branches reference the keys in opposite order, so the walk
finds the same paths
+ // in a different order. Spark canonicalizes the rest of the plan
order-free, so only a
+ // declaration kept in walk order would make the scans differ and lose
the reuse.
+ "same declaration found in another order keeps exchange reuse",
+ mapTable,
+ p =>
+ s"SELECT sum(t1.x + t2.x) FROM " +
+ s"(SELECT id, m['a'].t AS x FROM parquet.`$p` " +
+ s"WHERE m['a'].s = 'v3' AND m['b'].s = 'x') t1 " +
+ s"JOIN (SELECT id, m['a'].t AS x FROM parquet.`$p` " +
+ s"WHERE m['b'].s = 'x' AND m['a'].s = 'v3') t2 ON t1.id = t2.id",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t", "m[\"b\"].s")),
+ conf = Seq(noBroadcast),
+ requires = Some(classOf[ReusedExchangeExec])
+ ),
+ Case(
+ "union of keyed projections prunes each branch",
+ mapTable,
+ p =>
+ s"SELECT m['a'].t AS v FROM parquet.`$p` WHERE m['a'].s = 'v3' " +
+ s"UNION ALL SELECT m['b'].t FROM parquet.`$p` WHERE m['c'].s = 'y'",
+ Seq(m("m[\"a\"].s", "m[\"a\"].t"), m("m[\"b\"].t", "m[\"c\"].s"))
+ ),
+ Case(
+ "UDF in a partially offloaded project above the chain",
+ mapTable,
+ p => s"SELECT sum(mk_len(m['a'].s)), sum(id) FROM parquet.`$p` WHERE
m['b'].s = 'x'",
+ Seq(m("m[\"a\"].s", "m[\"b\"].s")),
+ conf = Seq(GlutenConfig.ENABLE_COLUMNAR_PARTIAL_PROJECT.key -> "true"),
+ requires = Some(classOf[ColumnarPartialProjectExec]),
+ before = () => spark.udf.register("mk_len", (s: String) => if (s ==
null) 0 else s.length)
+ ),
+ Case(
+ "script transformation above the chain",
+ mapTable,
+ p =>
+ s"SELECT TRANSFORM(id, m['a'].s) USING 'cat' AS (a STRING, b STRING) "
+
+ s"FROM parquet.`$p` WHERE m['b'].s = 'x'",
+ Seq(m("m[\"a\"].s", "m[\"b\"].s")),
+ requires = Some(classOf[BaseScriptTransformationExec]),
+ noFallBack = false
+ ),
+ // ----- unsupported shapes: the map stays whole -----
+ Case(
+ "whole-map use",
+ mapTable,
+ p => s"SELECT sum(size(m)) FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(whole)),
+ Case(
+ "non-constant key",
+ mapTable,
+ p => s"SELECT sum(m[if(id % 2 = 0, 'a', 'b')].t) FROM parquet.`$p`",
+ Seq(whole)),
+ Case(
+ "map in the output",
+ mapTable,
+ p => s"SELECT m FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(whole)),
+ Case(
+ "alias still holding the map in the output",
+ nestedMapTable,
+ p => s"SELECT c.m AS mm FROM parquet.`$p` WHERE c.m['a'] > 10",
+ Seq(whole)),
+ Case(
+ "project computing over the whole map",
+ mapTable,
+ p => s"SELECT map_keys(m) FROM parquet.`$p` WHERE m['a'].s = 'v3'",
+ Seq(whole)),
+ Case(
+ "key-free path next to a keyed path prunes nothing",
+ mapTable,
+ p => s"SELECT sum(size(m)), sum(m['a'].t) FROM parquet.`$p`",
+ Seq(whole)),
+ Case(
+ "key-free alias path next to a keyed path prunes nothing",
+ nestedMapTable,
+ p => s"SELECT sum(size(c.m)), sum(c.m['a']) FROM parquet.`$p`",
+ Seq(whole)),
+ Case(
+ "bare null check next to a keyed null check prunes nothing",
+ mapTable,
+ p => s"SELECT count(*) FROM parquet.`$p` WHERE m IS NOT NULL AND m['a']
IS NULL",
+ Seq(whole)
+ ),
+ Case(
+ "generate over the whole map",
+ scalarMapTable,
+ p => s"SELECT id, explode(m) FROM parquet.`$p`",
+ Seq(whole),
+ requires = Some(classOf[GenerateExecTransformer])
+ ),
+ Case(
+ "string key with invalid UTF-8 bytes",
+ invalidUtf8KeyTable,
+ p => s"SELECT sum(m[cast(x'FF' as string)]), sum(m['a']) FROM
parquet.`$p`",
+ Seq(whole),
+ // Vanilla Spark itself returns null for this lookup after the Parquet
round trip while the
+ // native reader finds the entry; only the declaration is checked here.
+ compare = false
+ ),
+ Case(
+ "date keys",
+ oddKeyTable,
+ p => s"SELECT sum(d[DATE '2020-01-01']) FROM parquet.`$p`",
+ Seq(whole),
+ noFallBack = false),
+ Case(
+ "decimal keys",
+ oddKeyTable,
+ p => s"SELECT sum(e[CAST(1.5 AS DECIMAL(3, 1))]) FROM parquet.`$p`",
+ Seq(whole),
+ noFallBack = false),
+ Case(
+ "map forwarded through a shuffle and read above a join",
+ mapTable,
+ p =>
+ s"SELECT sum(element_at(t1.m, 'a').t) FROM parquet.`$p` t1 " +
+ s"JOIN parquet.`$p` t2 ON t1.id = t2.id",
+ Seq(whole, whole),
+ conf = Seq(noBroadcast)
+ ),
+ Case(
+ "maps forwarded through a self join, one read whole",
+ twoMapTable,
+ p =>
+ s"WITH c AS (SELECT id, m, n FROM parquet.`$p`) " +
+ s"SELECT t1.m, t1.n['x'], t2.m['b'], t2.n['y'] FROM c t1 JOIN c t2
ON t1.id = t2.id",
+ Seq(whole),
+ conf = Seq(noBroadcast),
+ requires = Some(classOf[ReusedExchangeExec])
+ ),
+ Case(
+ "sort by a keyed value forwards the map through the shuffle",
+ scalarMapTable,
+ p => s"SELECT id FROM parquet.`$p` WHERE m['b'] > 10 ORDER BY m['a']",
+ Seq(whole)
+ )
+ ) ++ Seq("true", "false").map {
+ nativeUnion =>
+ Case(
+ s"union forwarding the map (native union = $nativeUnion)",
+ mapTable,
+ p =>
+ s"SELECT size(m), m['b'].s FROM (SELECT m FROM parquet.`$p` WHERE
m['a'].s = 'v0' " +
+ s"UNION ALL SELECT m FROM parquet.`$p` WHERE m['c'].s = 'y')",
+ Seq(whole, whole),
+ conf = Seq(GlutenConfig.NATIVE_UNION_ENABLED.key -> nativeUnion)
+ )
+ }
+
+ /** The set of map keys per row of a plan whose single output column is a
string-keyed map. */
+ private def keysPerRow(plan: SparkPlan): Set[Set[String]] =
+ plan
+ .executeCollect()
+
.map(_.getMap(0).keyArray().toArray[UTF8String](StringType).map(_.toString).toSet)
+ .toSet
+
+ /** 'plan' with every native scan declaring only key "a" of column m,
bypassing the rule. */
+ private def declareKeyA(plan: SparkPlan): SparkPlan = plan.transform {
+ case scan: BasicScanExecTransformer =>
+ scan.withRequiredMapSubfields(
+ Map("m" -> Seq(SubfieldPath("m",
Seq(SubfieldElement.StringKey("a"))))))
+ }
+
+ test("the native reader applies a declaration: undeclared keys are not
read") {
+ // Attaches a declaration to the scan directly, bypassing the rule, and
reads the map whole.
+ // Only the declared key may come back; if the transport or the native
side ignored the
+ // declaration, all three keys would.
+ mapTable {
+ path =>
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ val df = spark.read.parquet(path).select("m")
+ val plan = df.queryExecution.executedPlan
+ assert(planHas(df, classOf[BasicScanExecTransformer]), s"expected a
native scan:\n$plan")
+ assert(keysPerRow(plan) == Set(Set("a", "b", "c")))
+ val pruned = declareKeyA(plan)
+ assert(keysPerRow(pruned) == Set(Set("a")), s"declaration not
applied:\n$pruned")
+ }
+ }
+ }
+
+ test("a keyed filter at the scan keeps the declaration effective, element_at
included") {
+ // The filter is pushed to the scan and reaches Velox as the scan's
remaining filter. Velox
+ // adds the subfields that filter reads to the declaration: for a
pushdown-capable lookup the
+ // key path, for anything else the bare column, which keeps every entry.
The Velox backend
+ // emits element_at on a map as get_map_value, so both spellings
contribute m["a"].s and only
+ // the declared key comes back.
+ mapTable {
+ path =>
+ withSQLConf("spark.sql.adaptive.enabled" -> "false") {
+ Seq("element_at(m, 'a').s = 'v5'", "m['a'].s = 'v5'").foreach {
+ predicate =>
+ val df = spark.read.parquet(path).where(predicate).select("m")
+ val plan = df.queryExecution.executedPlan
+ assert(
+ planHas(df, classOf[BasicScanExecTransformer]),
+ s"expected a native scan:\n$plan")
+ assert(keysPerRow(plan) == Set(Set("a", "b", "c")), predicate)
+ val pruned = declareKeyA(plan)
+ assert(
+ keysPerRow(pruned) == Set(Set("a")),
+ s"$predicate read undeclared keys:\n$pruned")
+ }
+ }
+ }
+ }
+
+ Seq(false, true).foreach {
+ adaptive =>
+ val mode = s"adaptive execution ${if (adaptive) "on" else "off"}"
+ val modeConf = "spark.sql.adaptive.enabled" -> adaptive.toString
+
+ cases.foreach {
+ cs =>
+ test(s"${cs.name} ($mode)") {
+ cs.before()
+ cs.fixture {
+ path =>
+ withSQLConf((Seq(pruningFlag -> "true", modeConf) ++ cs.conf):
_*) {
+ runQueryAndCompare(
+ cs.sql(path),
+ compareResult = cs.compare,
+ noFallBack = cs.noFallBack) {
+ df =>
+ val plan = df.queryExecution.executedPlan
+ cs.requires.foreach {
+ cls => assert(planHas(df, cls), s"expected
${cls.getSimpleName}:\n$plan")
+ }
+ val actual = declaredPerScan(df)
+ assert(
+ multiset(actual) == multiset(cs.expected),
+ s"declared $actual, expected ${cs.expected}\n$plan")
+ }
+ }
+ }
+ }
+ }
+
+ // ----- cases that do not fit the table -----
+
+ test(s"pruning is off by default ($mode)") {
+ mapTable {
+ path =>
+ withSQLConf(modeConf) {
+ runQueryAndCompare(
+ s"SELECT sum(m['a'].t) FROM parquet.`$path` WHERE m['a'].s =
'v3'") {
+ df => assert(declared(df).isEmpty)
+ }
+ }
+ }
+ }
+
+ test(s"keys with quotes or backslashes survive the structural transport
($mode)") {
+ // Built with the DataFrame API so the keys reach the plan verbatim,
without SQL
+ // string-literal escaping in between.
+ val quoteKey = "q\"k"
+ val backslashKey = "b\\s"
+ withTempPath {
+ dir =>
+ val file = dir.getCanonicalPath
+ spark
+ .range(0, 100)
+ .select(
+ col("id"),
+ map(
+ lit(quoteKey),
+ struct(lit("v").as("s"), col("id").as("t")),
+ lit(backslashKey),
+ struct(lit("w").as("s"), (col("id") * 2).as("t"))).as("m"))
+ .write
+ .parquet(file)
+ def run(): (Row, Set[String]) = {
+ val df = spark.read
+ .parquet(file)
+ .select(
+ sum(col("m").getItem(quoteKey).getField("t")).as("q"),
+ sum(col("m").getItem(backslashKey).getField("t")).as("b"))
+ val row = df.collect().head
+ (row, declared(df).getOrElse("m", Set.empty))
+ }
+ var expected: Row = null
+ withSQLConf(pruningFlag -> "false", modeConf) {
+ val (row, paths) = run()
+ assert(paths.isEmpty)
+ expected = row
+ }
+ withSQLConf(pruningFlag -> "true", modeConf) {
+ val (row, paths) = run()
+ assert(paths == Set("m[\"q\\\"k\"].t", "m[\"b\\\\s\"].t"),
s"got: $paths")
+ assert(row == expected, s"pruned result $row differs from
$expected")
+ }
+ }
+ }
+
+ test(s"writing the whole map keeps it whole ($mode)") {
+ // The Project feeding the write forwards the map, so the chain ends
with it live.
+ mapTable {
+ path =>
+ withTable("mk_write_whole") {
+ spark.sql(
+ "CREATE TABLE mk_write_whole " +
+ "(id BIGINT, m MAP<STRING, STRUCT<s: STRING, t: BIGINT>>)
USING parquet")
+ withSQLConf(pruningFlag -> "true", modeConf) {
+ val df = spark.sql(
+ s"INSERT INTO mk_write_whole SELECT id, m FROM
parquet.`$path` " +
+ s"WHERE m['a'].s = 'v3'")
+ assert(declaredPerScan(df) == Seq(whole), s"map must stay
whole: ${declared(df)}")
+ }
+ // Every written map still holds all three keys.
+ checkAnswer(
+ spark.sql("SELECT count(*) - sum(CASE WHEN size(m) = 3 THEN 1
ELSE 0 END) " +
+ "FROM mk_write_whole"),
+ Row(0L) :: Nil)
+ }
+ }
+ }
+
+ test(s"writing keyed values prunes the scan ($mode)") {
+ // The Project feeding the write computes m['a'].t and drops the map.
+ mapTable {
+ path =>
+ withTable("mk_write_keyed") {
+ spark.sql("CREATE TABLE mk_write_keyed (id BIGINT, t BIGINT)
USING parquet")
+ withSQLConf(pruningFlag -> "true", modeConf) {
+ val df = spark.sql(
+ s"INSERT INTO mk_write_keyed SELECT id, m['a'].t FROM
parquet.`$path` " +
+ s"WHERE m['a'].s = 'v3'")
+ assert(
+ declared(df) == m("m[\"a\"].s", "m[\"a\"].t"),
+ s"expected pruning below the write: ${declared(df)}")
+ }
+ checkAnswer(
+ spark.sql("SELECT count(*), sum(t) FROM mk_write_keyed"),
+ spark.sql(
+ s"SELECT count(*), sum(m['a'].t) FROM parquet.`$path` WHERE
m['a'].s = 'v3'"))
+ }
+ }
+ }
+ }
+}
diff --git
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/GlutenHiveUDFSuite.scala
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/GlutenHiveUDFSuite.scala
index c2ee2f0dbe..9657d1033f 100644
---
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/GlutenHiveUDFSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/GlutenHiveUDFSuite.scala
@@ -16,8 +16,8 @@
*/
package org.apache.spark.sql.execution
-import org.apache.gluten.config.GlutenConfig
-import org.apache.gluten.execution.{ColumnarPartialGenerateExec,
ColumnarPartialProjectExec, GlutenQueryComparisonTest}
+import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
+import org.apache.gluten.execution.{BasicScanExecTransformer,
ColumnarPartialGenerateExec, ColumnarPartialProjectExec,
GlutenQueryComparisonTest}
import org.apache.gluten.expression.UDFMappings
import org.apache.gluten.udf.{CustomerUDF, DuplicateArray}
import org.apache.gluten.udtf.{ConditionalOutputUDTF, CustomerUDTF,
NoInputUDTF, SimpleUDTF}
@@ -135,6 +135,38 @@ class GlutenHiveUDFSuite extends GlutenQueryComparisonTest
with SQLTestUtils {
}
}
+ test("a partial generate over a hive udtf disables map-key pruning") {
+ // The map is forwarded through the partial generate and read above it, so
ScanMapKeyPruning
+ // finds it still live when its Filter / Project chain ends and leaves the
scan whole.
+ withTempFunction("simpleUDTF") {
+ sql(s"CREATE TEMPORARY FUNCTION simpleUDTF AS
'${classOf[SimpleUDTF].getName}'")
+ withTempPath {
+ path =>
+ spark
+ .range(0, 100)
+ .selectExpr("id", "map('a', id, 'b', id * 2) as m")
+ .write
+ .parquet(path.getCanonicalPath)
+ withTempView("map_udtf_src") {
+
spark.read.parquet(path.getCanonicalPath).createOrReplaceTempView("map_udtf_src")
+ withSQLConf(VeloxConfig.SCAN_MAP_KEY_PRUNING_ENABLED.key ->
"true") {
+ runQueryAndCompare(
+ "select col0, m['a'] from map_udtf_src lateral view
simpleUDTF(id) as col0") {
+ df =>
+ checkOperatorMatch[ColumnarPartialGenerateExec](df)
+ val declared = getExecutedPlan(df).collect {
+ case scan: BasicScanExecTransformer =>
scan.requiredMapSubfields
+ }.filter(_.nonEmpty)
+ assert(
+ declared.isEmpty,
+ s"plan with a partial generate must not be pruned:
$declared")
+ }
+ }
+ }
+ }
+ }
+ }
+
test("simple udtf") {
withTempFunction("simpleUDTF") {
sql(s"CREATE TEMPORARY FUNCTION simpleUDTF AS
'${classOf[SimpleUDTF].getName}'")
diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt
index 71bb0f9a60..fb4b0e4679 100644
--- a/cpp/velox/CMakeLists.txt
+++ b/cpp/velox/CMakeLists.txt
@@ -176,6 +176,7 @@ set(VELOX_SRCS
memory/VeloxMemoryManager.cc
operators/functions/RegistrationAllFunctions.cc
operators/functions/delta/DeltaBitmapAggregator.cc
+ operators/functions/overlay/GetMapValue.cc
operators/functions/overlay/RegisterFunctionOverlay.cc
operators/functions/RowConstructorWithNull.cc
operators/functions/SparkExprToSubfieldFilterParser.cc
diff --git a/cpp/velox/operators/functions/overlay/GetMapValue.cc
b/cpp/velox/operators/functions/overlay/GetMapValue.cc
new file mode 100644
index 0000000000..56570750d5
--- /dev/null
+++ b/cpp/velox/operators/functions/overlay/GetMapValue.cc
@@ -0,0 +1,40 @@
+/*
+ * 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.
+ */
+#include "operators/functions/overlay/GetMapValue.h"
+
+#include "velox/expression/VectorFunction.h"
+
+using namespace facebook;
+
+namespace gluten {
+
+void registerSparkGetMapValueFunction(const std::string& name) {
+ velox::exec::registerStatefulVectorFunction(
+ name,
+ SparkGetMapValueFunction::signatures(),
+ [](const std::string&,
+ const std::vector<velox::exec::VectorFunctionArg>& inputArgs,
+ const velox::core::QueryConfig& config) ->
std::shared_ptr<velox::exec::VectorFunction> {
+ VELOX_CHECK_EQ(inputArgs.size(), 2);
+ VELOX_CHECK(inputArgs[0].type->isMap(), "get_map_value expects a map
input");
+ // A per-call instance so the map subscript cache stays private to the
+ // expression, as in velox::functions::registerElementAtFunction.
+ return
std::make_shared<SparkGetMapValueFunction>(config.isExpressionEvaluationCacheEnabled());
+ });
+}
+
+} // namespace gluten
diff --git a/cpp/velox/operators/functions/overlay/GetMapValue.h
b/cpp/velox/operators/functions/overlay/GetMapValue.h
new file mode 100644
index 0000000000..990aca1ddd
--- /dev/null
+++ b/cpp/velox/operators/functions/overlay/GetMapValue.h
@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+#pragma once
+
+#include "velox/functions/lib/SubscriptUtil.h"
+
+namespace gluten {
+
+/// Spark GetMapValue (m[key]): returns the value stored under 'key', or NULL
+/// when the key is absent. Only the map signature is registered.
+///
+/// Unlike Velox's element_at, which serves arrays as well and therefore keeps
+/// Subscript::canPushdown() == false (a Subfield path cannot express a
+/// negative array index), a map lookup with a constant key is exactly the
+/// path m["k"]: the reader keeps only the entries whose key is "k", and the
+/// lookup yields the same value, or NULL, as on the full map. Reporting the
+/// subscript lets Expr::extractSubfields() produce m["k"] for a remaining
+/// filter such as get_map_value(m, 'k').x = 1 instead of the bare column path
+/// m, which would otherwise force the reader to keep every map entry and
+/// defeat map-key pruning declared through
HiveColumnHandle::requiredSubfields.
+class SparkGetMapValueFunction : public
facebook::velox::functions::SubscriptImpl<
+ /* allowNegativeIndices */ false,
+ /* nullOnNegativeIndices */ false,
+ /* allowOutOfBound */ true,
+ /* indexStartsAtOne */ true> {
+ public:
+ explicit SparkGetMapValueFunction(bool allowCaching) :
SubscriptImpl(allowCaching) {}
+
+ bool canPushdown() const override {
+ return true;
+ }
+
+ static
std::vector<std::shared_ptr<facebook::velox::exec::FunctionSignature>>
signatures() {
+ // map(K,V), K -> V
+ return {facebook::velox::exec::FunctionSignatureBuilder()
+ .typeVariable("K")
+ .typeVariable("V")
+ .returnType("V")
+ .argumentType("map(K,V)")
+ .argumentType("K")
+ .build()};
+ }
+};
+
+/// Registers Spark GetMapValue under 'name'.
+void registerSparkGetMapValueFunction(const std::string& name);
+
+} // namespace gluten
diff --git a/cpp/velox/operators/functions/overlay/RegisterFunctionOverlay.cc
b/cpp/velox/operators/functions/overlay/RegisterFunctionOverlay.cc
index d16e7636bd..d5d4097531 100644
--- a/cpp/velox/operators/functions/overlay/RegisterFunctionOverlay.cc
+++ b/cpp/velox/operators/functions/overlay/RegisterFunctionOverlay.cc
@@ -16,6 +16,7 @@
*/
#include "operators/functions/overlay/RegisterFunctionOverlay.h"
+#include "operators/functions/overlay/GetMapValue.h"
#include "operators/functions/overlay/Round.h"
#include "velox/functions/lib/RegistrationHelpers.h"
@@ -40,6 +41,11 @@ void registerRoundFunction() {
void registerFunctionOverlay() {
registerRoundFunction();
+ // Spark GetMapValue (m[key]) is emitted as get_map_value rather than being
+ // rewritten to Velox's element_at: a map-only subscript can report
+ // canPushdown(), so remaining filters like m['k'].x = v extract m["k"] and
+ // do not defeat map-key pruning.
+ registerSparkGetMapValueFunction("get_map_value");
}
} // namespace gluten
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc
b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
index ce9de452fb..dc9833663a 100644
--- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc
+++ b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
@@ -35,6 +35,7 @@
#include "utils/VeloxArrowUtils.h"
#include "utils/VeloxWriterUtils.h"
+#include "RequiredSubfieldsExtension.pb.h"
#include "config.pb.h"
#include "config/GlutenConfig.h"
#include "config/VeloxConfig.h"
@@ -1566,17 +1567,20 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
std::vector<std::string> colNameList;
std::vector<TypePtr> veloxTypeList;
std::vector<ColumnType> columnTypes;
- // Convert field names into lower case when not case-sensitive.
+ // Convert field names into lower case when not case-sensitive. Every name
that has to match a
+ // scan column (schema names, required-subfield column and field names) goes
through foldCase.
bool asLowerCase = !veloxCfg_->get<bool>(kCaseSensitive, false);
+ auto foldCase = [asLowerCase](std::string name) {
+ if (asLowerCase) {
+ folly::toLowerAscii(name);
+ }
+ return name;
+ };
if (readRel.has_base_schema()) {
const auto& baseSchema = readRel.base_schema();
colNameList.reserve(baseSchema.names().size());
for (const auto& name : baseSchema.names()) {
- std::string fieldName = name;
- if (asLowerCase) {
- folly::toLowerAscii(fieldName);
- }
- colNameList.emplace_back(fieldName);
+ colNameList.emplace_back(foldCase(name));
}
veloxTypeList = SubstraitParser::parseNamedStruct(baseSchema, asLowerCase);
SubstraitParser::parseColumnTypes(baseSchema, columnTypes);
@@ -1632,6 +1636,52 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
tableHandle = std::make_shared<connector::hive::HiveTableHandle>(
connectorId, "hive_table", std::move(subfieldFilters), remainingFilter,
dataColumns);
+ // Required subfields of scan columns (map-key pruning), from the
RequiredSubfieldsExtension the
+ // ScanMapKeyPruning rule packs into the ReadRel advanced extension. Each
path is transported as
+ // structured elements and rebuilt as a common::Subfield here, so no path
syntax is parsed.
+ std::unordered_map<std::string, std::vector<common::Subfield>>
requiredSubfieldsByCol;
+ if (readRel.has_advanced_extension() &&
readRel.advanced_extension().has_enhancement()) {
+ const auto& enhancement = readRel.advanced_extension().enhancement();
+ if (enhancement.Is<gluten::RequiredSubfieldsExtension>()) {
+ gluten::RequiredSubfieldsExtension extension;
+ VELOX_USER_CHECK(enhancement.UnpackTo(&extension), "Failed to unpack
RequiredSubfieldsExtension");
+ for (const auto& column : extension.columns()) {
+ auto columnName = foldCase(column.column());
+ VELOX_USER_CHECK(!columnName.empty(), "Required subfields with an
empty column name");
+ std::vector<common::Subfield> subfields;
+ for (const auto& subfield : column.subfields()) {
+ std::vector<std::unique_ptr<common::Subfield::PathElement>> elements;
+
elements.push_back(std::make_unique<common::Subfield::NestedField>(columnName));
+ for (const auto& element : subfield.elements()) {
+ using PathElement =
gluten::RequiredSubfieldsExtension::PathElement;
+ switch (element.element_case()) {
+ case PathElement::kField: {
+ // Same folding as the schema's field names, or the field
would resolve to nothing
+ // and read as null.
+ auto fieldName = foldCase(element.field());
+ VELOX_USER_CHECK(!fieldName.empty(), "Required subfield with
an empty field name");
+
elements.push_back(std::make_unique<common::Subfield::NestedField>(fieldName));
+ break;
+ }
+ case PathElement::kStringKey:
+
elements.push_back(std::make_unique<common::Subfield::StringSubscript>(element.string_key()));
+ break;
+ case PathElement::kLongKey:
+
elements.push_back(std::make_unique<common::Subfield::LongSubscript>(element.long_key()));
+ break;
+ default:
+ VELOX_USER_FAIL("Required subfield of column {} has an element
without a value", columnName);
+ }
+ }
+ subfields.emplace_back(std::move(elements));
+ }
+ if (!subfields.empty()) {
+ requiredSubfieldsByCol[columnName] = std::move(subfields);
+ }
+ }
+ }
+ }
+
// Get assignments and out names.
std::vector<std::string> outNames;
outNames.reserve(colNameList.size());
@@ -1647,9 +1697,7 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
icebergColumn = &columnIt->second;
} else if (asLowerCase) {
for (const auto& [name, column] : icebergSplitInfo->columns) {
- auto normalizedName = name;
- folly::toLowerAscii(normalizedName);
- if (normalizedName == colNameList[idx]) {
+ if (foldCase(name) == colNameList[idx]) {
icebergColumn = &column;
break;
}
@@ -1660,6 +1708,8 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
// handles only for regular columns so all data columns are mapped by field
// ID together, while partition columns keep the existing Hive conversion.
if (icebergColumn && columnType == ColumnType::kRegular) {
+ // Iceberg scans do not declare required subfields (the Scala side does
not prune them); a
+ // declaration reaching this branch is left unused and reported below.
assignments[outName] =
std::make_shared<connector::hive::iceberg::IcebergColumnHandle>(
colNameList[idx],
columnType,
@@ -1668,11 +1718,27 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
std::vector<common::Subfield>{},
icebergColumn->initialDefault);
} else {
+ std::vector<common::Subfield> requiredSubfields;
+ if (columnType == ColumnType::kRegular &&
!requiredSubfieldsByCol.empty()) {
+ auto it = requiredSubfieldsByCol.find(colNameList[idx]);
+ if (it != requiredSubfieldsByCol.end()) {
+ VLOG(1) << "Map-key pruning: applying " << it->second.size() << "
required subfields to column "
+ << colNameList[idx];
+ requiredSubfields = std::move(it->second);
+ requiredSubfieldsByCol.erase(it);
+ }
+ }
assignments[outName] =
std::make_shared<connector::hive::HiveColumnHandle>(
- colNameList[idx], columnType, veloxTypeList[idx],
veloxTypeList[idx]);
+ colNameList[idx], columnType, veloxTypeList[idx],
veloxTypeList[idx], std::move(requiredSubfields));
}
outNames.emplace_back(outName);
}
+ // A declared column that matched no regular scan column is not applied; the
column is then read
+ // whole, which is correct but not what the plan asked for.
+ for (const auto& [columnName, subfields] : requiredSubfieldsByCol) {
+ LOG(WARNING) << "Map-key pruning: " << subfields.size() << " required
subfields declared for column '" << columnName
+ << "' matched no scan column and were ignored";
+ }
auto outputType = ROW(std::move(outNames), std::move(veloxTypeList));
if (readRel.has_virtual_table()) {
diff --git a/cpp/velox/tests/SparkFunctionTest.cc
b/cpp/velox/tests/SparkFunctionTest.cc
index ceb979a5ed..0384a69bee 100644
--- a/cpp/velox/tests/SparkFunctionTest.cc
+++ b/cpp/velox/tests/SparkFunctionTest.cc
@@ -144,3 +144,36 @@ TEST_F(SparkFunctionTest,
expressionLevelLegacyCastIgnoresSessionAnsiOn) {
facebook::velox::test::assertEqualVectors(makeFlatVector<int8_t>({-121}),
evaluate(legacyCast, input));
}
+
+// Spark GetMapValue is emitted as get_map_value and registered through the
+// function overlay as a map-only, pushdown-capable subscript, so that
+// Expr::extractSubfields() yields m["k"] instead of the bare column path. This
+// keeps map-key pruning declared via HiveColumnHandle::requiredSubfields
+// effective when the remaining filter references the map.
+TEST_F(SparkFunctionTest, getMapValueSubfieldExtraction) {
+ // Spark function names throughout: the test base registers only the Spark
+ // functions, so the parser's `>` (gt) does not resolve.
+ auto mapRow = ROW({"m"}, {MAP(VARCHAR(), BIGINT())});
+ auto exprSet = compileExpression("greaterthan(get_map_value(m, 'k'), 0)",
mapRow);
+ auto subfields = exprSet->expr(0)->extractSubfields();
+ ASSERT_EQ(subfields.size(), 1);
+ EXPECT_EQ(subfields[0].toString(), "m[\"k\"]");
+
+ // Velox's element_at stays non-pushdown: it also serves arrays, and a
+ // Subfield path cannot express a negative array index.
+ exprSet = compileExpression("greaterthan(element_at(m, 'k'), 0)", mapRow);
+ subfields = exprSet->expr(0)->extractSubfields();
+ ASSERT_EQ(subfields.size(), 1);
+ EXPECT_EQ(subfields[0].toString(), "m");
+}
+
+TEST_F(SparkFunctionTest, getMapValueSemantics) {
+ auto maps = makeMapVector<std::string, int64_t>({{{"k", 1}}, {{"x", 2}},
{}});
+ auto result = evaluate("get_map_value(c0, 'k')", makeRowVector({maps}));
+
facebook::velox::test::assertEqualVectors(makeNullableFlatVector<int64_t>({1,
std::nullopt, std::nullopt}), result);
+
+ // Integer literals parse as BIGINT, and the key must match the map's key
type exactly.
+ auto intKeyed = makeMapVector<int64_t, std::string>({{{1, "a"}, {2, "b"}},
{{3, "c"}}});
+ result = evaluate("get_map_value(c0, 2)", makeRowVector({intKeyed}));
+
facebook::velox::test::assertEqualVectors(makeNullableFlatVector<std::string>({"b",
std::nullopt}), result);
+}
diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md
index 86b1b116eb..6cc3995645 100644
--- a/docs/velox-configuration.md
+++ b/docs/velox-configuration.md
@@ -79,6 +79,7 @@ nav_order: 16
|
spark.gluten.sql.columnar.backend.velox.resizeBatches.shuffleInputOutput.minSize
| 🔄 Dynamic | <undefined> | The minimum batch size for shuffle input
and output. If size of an input batch is smaller than the value, it will be
combined with other batches before sending to shuffle. The same applies for
batches output by shuffle read. Only functions when
spark.gluten.sql.columnar.backend.velox.resizeBatches.shuffleInput or
spark.gluten.sql.columnar.backend.velox.resizeBatches.shu [...]
| spark.gluten.sql.columnar.backend.velox.resizeBatches.shuffleOutput
| 🔄 Dynamic | false | If true, combine small columnar
batches together right after shuffle read. The default minimum output batch
size is equal to 0.25 * spark.gluten.sql.columnar.maxBatchSize
[...]
| spark.gluten.sql.columnar.backend.velox.scan.bloomFilterPushdown.enabled
| âš“ Static | false | Whether to push Bloom filters into
Velox scans.
[...]
+| spark.gluten.sql.columnar.backend.velox.scanMapKeyPruningEnabled
| 🔄 Dynamic | false | When true, a map column read by a
native Parquet scan is declared to Velox with only the constant keys the query
looks up (m['k'], element_at(m, 'k'), null checks), so the reader skips the
other entries. Applies when every use of the map sits in the filters and
projections directly above the scan and the map is dropped before any other
operator sees it; otherwise the ma [...]
| spark.gluten.sql.columnar.backend.velox.showTaskMetricsWhenFinished
| 🔄 Dynamic | false | Show velox full task metrics when
finished.
[...]
| spark.gluten.sql.columnar.backend.velox.spillFileSystem
| 🔄 Dynamic | local | The filesystem used to store spill
data. local: The local file system. heap-over-local: Write file to JVM heap if
having extra heap space. Otherwise write to local file system.
[...]
| spark.gluten.sql.columnar.backend.velox.spillNumMaxMergeFiles
| 🔄 Dynamic | 0 | The max number of files to merge at a
time when merging sorted files into a single ordered stream. 0 means unlimited.
[...]
diff --git
a/gluten-delta/src/main/scala/org/apache/gluten/execution/DeltaScanTransformer.scala
b/gluten-delta/src/main/scala/org/apache/gluten/execution/DeltaScanTransformer.scala
index fe26accb71..6165b70439 100644
---
a/gluten-delta/src/main/scala/org/apache/gluten/execution/DeltaScanTransformer.scala
+++
b/gluten-delta/src/main/scala/org/apache/gluten/execution/DeltaScanTransformer.scala
@@ -49,7 +49,8 @@ case class DeltaScanTransformer(
override val dataFilters: Seq[Expression],
override val tableIdentifier: Option[TableIdentifier],
override val disableBucketedScan: Boolean = false,
- override val pushDownFilters: Option[Seq[Expression]] = None)
+ override val pushDownFilters: Option[Seq[Expression]] = None,
+ override val requiredMapSubfields: Map[String, Seq[SubfieldPath]] =
Map.empty)
extends FileSourceScanExecTransformerBase(
relation,
stream,
@@ -257,12 +258,25 @@ case class DeltaScanTransformer(
QueryPlan.normalizePredicates(dataFilters, output),
None,
disableBucketedScan,
- pushDownFilters.map(QueryPlan.normalizePredicates(_, output))
+ pushDownFilters.map(QueryPlan.normalizePredicates(_, output)),
+ requiredMapSubfields
)
}
override def withNewPushdownFilters(filters: Seq[Expression]):
BasicScanExecTransformer =
copy(pushDownFilters = Some(filters))
+
+ // Under name or id column mapping the scan reads physical `col-<uuid>`
names, which the rule
+ // renders from the logical schema; the paths would not match the file. Only
a Delta Parquet
+ // format without column mapping prunes; any other format cannot be checked
and is not pruned.
+ override def supportsMapKeyPruning: Boolean = relation.fileFormat match {
+ case d: DeltaParquetFileFormat => d.columnMappingMode == NoMapping
+ case _ => false
+ }
+
+ override def withRequiredMapSubfields(
+ subfields: Map[String, Seq[SubfieldPath]]): BasicScanExecTransformer =
+ copy(requiredMapSubfields = subfields)
}
object DeltaScanTransformer {
diff --git
a/gluten-delta/src/test/scala/org/apache/gluten/execution/DeltaSuite.scala
b/gluten-delta/src/test/scala/org/apache/gluten/execution/DeltaSuite.scala
index 87915b7b77..26caffcefe 100644
--- a/gluten-delta/src/test/scala/org/apache/gluten/execution/DeltaSuite.scala
+++ b/gluten-delta/src/test/scala/org/apache/gluten/execution/DeltaSuite.scala
@@ -632,7 +632,7 @@ abstract class DeltaSuite extends
WholeStageTransformerSuite {
}
}
- test("deletion vector on partitioned table") {
+ testWithMinSparkVersion("deletion vector on partitioned table", "3.4") {
withTempPath {
p =>
import testImplicits._
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
index 916c31c2ed..0141320030 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
@@ -175,6 +175,17 @@ trait SparkPlanExecApi {
right: ExpressionTransformer,
original: GetMapValue): ExpressionTransformer
+ /**
+ * Generate an expression transformer to transform ElementAt to Substrait.
The default emits the
+ * mapped function as is; a backend may pick another native function for
some inputs.
+ */
+ def genElementAtTransformer(
+ substraitExprName: String,
+ left: ExpressionTransformer,
+ right: ExpressionTransformer,
+ original: ElementAt): ExpressionTransformer =
+ GenericExpressionTransformer(substraitExprName, Seq(left, right), original)
+
def genStringToMapTransformer(
substraitExprName: String,
children: Seq[ExpressionTransformer],
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/TransformerApi.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/TransformerApi.scala
index c35a85398a..98330347a3 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/TransformerApi.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/TransformerApi.scala
@@ -16,7 +16,7 @@
*/
package org.apache.gluten.backendsapi
-import org.apache.gluten.execution.WriteFilesExecTransformer
+import org.apache.gluten.execution.{SubfieldPath, WriteFilesExecTransformer}
import org.apache.gluten.substrait.SubstraitContext
import org.apache.gluten.substrait.expression.ExpressionNode
@@ -81,6 +81,14 @@ trait TransformerApi {
throw new UnsupportedOperationException("Iceberg initial-default reads are
not supported")
}
+ /**
+ * Packs the required subfields of scan columns (map-key pruning) into a
backend-specific read
+ * extension; see BasicScanExecTransformer.requiredMapSubfields.
+ */
+ def packRequiredSubfields(subfields: Map[String, Seq[SubfieldPath]]): Any = {
+ throw new UnsupportedOperationException("Map-key pruning is not supported")
+ }
+
/** This method is only used for CH backend tests */
def invalidateSQLExecutionResource(executionId: String): Unit = {}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
index c765ed333d..faa19881e3 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
@@ -79,6 +79,24 @@ trait BasicScanExecTransformer extends LeafTransformSupport
with BaseDataSource
s"${getClass.toString} does not support push down filters.")
}
+ /**
+ * Required subfields per top-level column, set by the backend map-key
pruning rule. Key is the
+ * normalized column name; values are the paths the query accesses under it,
for example
+ * `c.m["k1"].f1`. Serialized into the ReadRel advanced extension so the
native reader keeps only
+ * the map entries those paths name.
+ */
+ def requiredMapSubfields: Map[String, Seq[SubfieldPath]] = Map.empty
+
+ /** Whether this scan supports map-key pruning via required subfields. */
+ def supportsMapKeyPruning: Boolean = false
+
+ /** Copy the scan with required subfield paths for map-key pruning. */
+ def withRequiredMapSubfields(
+ subfields: Map[String, Seq[SubfieldPath]]): BasicScanExecTransformer = {
+ throw new UnsupportedOperationException(
+ s"${getClass.toString} does not support map-key pruning.")
+ }
+
def getMetadataColumns(): Seq[AttributeReference]
/** This can be used to report FileFormat for a file based scan operator. */
@@ -188,7 +206,12 @@ trait BasicScanExecTransformer extends
LeafTransformSupport with BaseDataSource
val optimization =
BackendsApiManager.getTransformerApiInstance.packPBMessage(
StringValue.newBuilder.setValue(s"isMergeTree=$mergeTreeFlag\n").build)
- val extensionNode = ExtensionBuilder.makeAdvancedExtension(optimization,
null)
+ val enhancement = if (requiredMapSubfields.isEmpty) {
+ null
+ } else {
+
BackendsApiManager.getTransformerApiInstance.packRequiredSubfields(requiredMapSubfields)
+ }
+ val extensionNode = ExtensionBuilder.makeAdvancedExtension(optimization,
enhancement)
val readNode = RelBuilder.makeReadRel(
typeNodes,
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala
index a50b252043..c95237c7ed 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala
@@ -48,7 +48,8 @@ case class BatchScanExecTransformer(
override val commonPartitionValues: Option[Seq[(InternalRow, Int)]] = None,
override val applyPartialClustering: Boolean = false,
override val replicatePartitions: Boolean = false,
- override val pushDownFilters: Option[Seq[Expression]] = None)
+ override val pushDownFilters: Option[Seq[Expression]] = None,
+ override val requiredMapSubfields: Map[String, Seq[SubfieldPath]] =
Map.empty)
extends BatchScanExecTransformerBase(
output,
scan,
@@ -80,6 +81,13 @@ case class BatchScanExecTransformer(
this.copy(pushDownFilters = Some(filters))
}
+ override def supportsMapKeyPruning: Boolean = true
+
+ override def withRequiredMapSubfields(
+ subfields: Map[String, Seq[SubfieldPath]]): BatchScanExecTransformerBase
= {
+ this.copy(requiredMapSubfields = subfields)
+ }
+
override def withOutput(newOutput: Seq[AttributeReference]):
BatchScanExecTransformerBase = {
this.copy(output = newOutput)
}
@@ -216,13 +224,19 @@ abstract class BatchScanExecTransformerBase(
case other: BatchScanExecTransformerBase =>
this.keyGroupedPartitioning == other.keyGroupedPartitioning &&
this.pushDownFilters == other.pushDownFilters &&
+ this.requiredMapSubfields == other.requiredMapSubfields &&
super.equals(other)
case _ =>
false
}
override def hashCode(): Int =
- Objects.hashCode(batch, runtimeFilters, keyGroupedPartitioning,
pushDownFilters)
+ Objects.hashCode(
+ batch,
+ runtimeFilters,
+ keyGroupedPartitioning,
+ pushDownFilters,
+ requiredMapSubfields)
/** Return a copy of this scan with a new output schema. */
def withOutput(newOutput: Seq[AttributeReference]):
BatchScanExecTransformerBase
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
index 088c5f2d6d..ba8e66b076 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
@@ -50,7 +50,8 @@ case class FileSourceScanExecTransformer(
override val tableIdentifier: Option[TableIdentifier],
override val disableBucketedScan: Boolean = false,
override val pushDownFilters: Option[Seq[Expression]] = None,
- inputStats: Option[InputStats] = None)
+ inputStats: Option[InputStats] = None,
+ override val requiredMapSubfields: Map[String, Seq[SubfieldPath]] =
Map.empty)
extends FileSourceScanExecTransformerBase(
relation,
stream,
@@ -85,12 +86,19 @@ case class FileSourceScanExecTransformer(
None,
disableBucketedScan,
pushDownFilters.map(QueryPlan.normalizePredicates(_, output)),
- inputStats
+ inputStats,
+ requiredMapSubfields
)
}
override def withNewPushdownFilters(filters: Seq[Expression]):
FileSourceScanExecTransformer =
copy(pushDownFilters = Some(filters))
+
+ override def supportsMapKeyPruning: Boolean = true
+
+ override def withRequiredMapSubfields(
+ subfields: Map[String, Seq[SubfieldPath]]):
FileSourceScanExecTransformer =
+ copy(requiredMapSubfields = subfields)
}
abstract class FileSourceScanExecTransformerBase(
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/SubfieldPath.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/SubfieldPath.scala
new file mode 100644
index 0000000000..6515bd1ef7
--- /dev/null
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/SubfieldPath.scala
@@ -0,0 +1,51 @@
+/*
+ * 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
+
+/** One step of a [[SubfieldPath]] below its column. */
+sealed trait SubfieldElement
+
+object SubfieldElement {
+
+ /** A struct field, by the schema's name. */
+ case class Field(name: String) extends SubfieldElement
+
+ /** A map lookup by a string key. */
+ case class StringKey(key: String) extends SubfieldElement
+
+ /** A map lookup by an integral key. */
+ case class LongKey(key: Long) extends SubfieldElement
+}
+
+/**
+ * A path from a scan output column into its value, e.g. column `m` with
elements [StringKey("a"),
+ * Field("t")] for `m['a'].t`. Declared on a scan as a required subfield so
the native reader keeps
+ * only the map entries such paths name. Transported structurally (see
+ * RequiredSubfieldsExtension.proto); toString renders the Velox Subfield
syntax for logs and tests.
+ */
+case class SubfieldPath(column: String, elements: Seq[SubfieldElement]) {
+ override def toString: String = {
+ val sb = new StringBuilder(column)
+ elements.foreach {
+ case SubfieldElement.Field(name) => sb.append('.').append(name)
+ case SubfieldElement.StringKey(key) =>
+ sb.append("[\"").append(key.replace("\\", "\\\\").replace("\"",
"\\\"")).append("\"]")
+ case SubfieldElement.LongKey(key) =>
sb.append('[').append(key).append(']')
+ }
+ sb.toString
+ }
+}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
index ff8acfb8a2..97e21a375b 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
@@ -319,6 +319,13 @@ object ExpressionConverter extends SQLConfHelper with
Logging {
val children =
c.children.map(replaceWithExpressionTransformer0(_, attributeSeq,
expressionsMap))
CreateMapTransformer(substraitExprName, children, c)
+ case ea: ElementAt =>
+ BackendsApiManager.getSparkPlanExecApiInstance.genElementAtTransformer(
+ substraitExprName,
+ replaceWithExpressionTransformer0(ea.left, attributeSeq,
expressionsMap),
+ replaceWithExpressionTransformer0(ea.right, attributeSeq,
expressionsMap),
+ ea
+ )
case g: GetMapValue =>
BackendsApiManager.getSparkPlanExecApiInstance.genGetMapValueTransformer(
substraitExprName,
diff --git
a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
index 1b541b3267..8d726359ce 100644
--- a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
+++ b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
@@ -30,7 +30,7 @@ import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.util._
import org.apache.spark.sql.classic.ClassicConversions._
import org.apache.spark.sql.execution.{CommandResultExec, SparkPlan,
SQLExecution, UnaryExecNode, WholeStageCodegenExec}
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
BroadcastQueryStageExec, ShuffleQueryStageExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.InMemoryRelation
import org.apache.spark.storage.StorageLevel
@@ -345,6 +345,8 @@ abstract class GlutenQueryTest extends QueryTest with
AdaptiveSparkPlanHelper {
getExecutedPlan(cmd.commandPhysicalPlan)
case s: ShuffleQueryStageExec =>
getExecutedPlan(s.plan)
+ case b: BroadcastQueryStageExec =>
+ getExecutedPlan(b.plan)
case plan =>
plan.children.flatMap(getExecutedPlan)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
