You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@gluten.apache.org by "marin-ma (via GitHub)" <gi...@apache.org> on 2026/09/23 09:05:25 UTC
[PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
marin-ma opened a new pull request, #13109:
URL: https://github.com/apache/gluten/pull/13109
Velox's Hive connector can declare required subfields on a column handle, and its Parquet map reader then materialises only the entries whose key matches, so that reading m['k'] no longer decodes every entry of every row's map. This change lets Gluten derive those declarations for map columns read by a native Parquet scan.
Planner rule ScanMapKeyPruning (Velox backend, post-transform, off by default, spark.gluten.sql.columnar.backend.velox.scanMapKeyPruningEnabled): for each map-typed output of a native Parquet scan it walks up through Filter, Project and Generate nodes only. Every reference to the map on the way must be a constant-key path (GetMapValue, ElementAt on a map, GetStructField steps) or a null check, and some Project or Generate must drop the attribute, and every alias still holding map data, before the chain ends. Since an operator can only reference what its child outputs, nothing above that point can read the map and needs no analysis. Any other shape leaves the map whole: the map still live at an exchange, join, aggregate, union, write or the fragment output, a whole-map use such as size(m) or explode(m), a non-constant key, a key type without a Velox subscript form (date, decimal), or a string key whose bytes are not valid UTF-8. Null checks need no map entries and are declared only
when no value path covers them; a declaration is emitted only when it lets the reader skip something. Spark's NestedColumnAliasing alias of a struct field holding the map (c.m AS _extract_m) is followed with a path prefix.
Transport and native side: paths travel as a structured protobuf (RequiredSubfieldsExtension: column, then field / string key / long key elements) in the ReadRel advanced extension and are rebuilt as common::Subfield on the HiveColumnHandle, with case folded like the schema's names; declared columns that match no scan column are logged.
Spark's GetMapValue is emitted as get_map_value, a Gluten-registered map-only subscript that reports canPushdown(), so a remaining filter such as m['k'].x = v extracts m["k"] instead of the whole map and does not defeat the declaration. Delta scans participate when the table uses no column mapping; BatchScan and FileSourceScan transformers carry the declaration and include it in their equality.
Tests: ScanMapKeyPruningSuite runs each case with adaptive execution off and on and asserts the exact declaration of every native scan, including a direct-declaration test that the reader returns only the declared key; Delta and Hive UDF suites cover column mapping and partial generates.
## Was this patch authored or co-authored using generative AI tooling?
Claude Fable 5.1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "marin-ma (via GitHub)" <gi...@apache.org>.
marin-ma merged PR #13109:
URL: https://github.com/apache/gluten/pull/13109
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "marin-ma (via GitHub)" <gi...@apache.org>.
marin-ma commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4157521956
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,421 @@
+/*
+ * 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;
+ * - 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))
Review Comment:
The Velox schema does not stay `c.Ärger`. The scan's Substrait schema names are
built on the JVM by ConverterUtils.collectAttributeNamesWithoutExprId and
collectStructFieldNames, which normalize every column and nested field name with
normalizeColName, the same function renderSubfieldPaths uses. So both the schema
and the declaration reach the native side as `c.ärger`; the ASCII fold there is
applied to both and is a no-op on an already lowercased name. The umlaut case in
ScanMapKeyPruningSuite compares its result with vanilla Spark, so an unresolved
path would show up as a wrong sum. It passes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4146592072
##########
cpp/velox/substrait/SubstraitToVeloxPlan.cc:
##########
@@ -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";
Review Comment:
This warning can become noisy in production if there are benign mismatches (e.g., schema/name normalization differences across components, or feature interactions where pruning is intentionally ignored). Consider reducing severity (e.g., `VLOG(1)`/`LOG(INFO)`) or adding rate-limiting/context (scan ID / table / query ID) so operators can diagnose without flooding logs.
##########
gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala:
##########
@@ -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)
Review Comment:
Packing required subfields into `ReadRel.advanced_extension.enhancement` makes this scan-level channel single-purpose (protobuf `Any` can hold only one message). If any scan already uses `enhancement` for another read extension (e.g., Iceberg or future features), this will be mutually exclusive and silently drop one side. Consider introducing a wrapper enhancement message that can carry multiple optional sub-messages (or repeated `Any`s), and update both packing/unpacking to merge extensions rather than overwrite.
##########
backends-velox/src-delta/test/scala/org/apache/gluten/execution/VeloxDeltaSuite.scala:
##########
@@ -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(
Review Comment:
The pruning config key is hard-coded here while other tests use `VeloxConfig.SCAN_MAP_KEY_PRUNING_ENABLED.key`. Using the shared constant would avoid drift if the key ever changes and keeps config usage consistent across suites.
##########
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("\"]")
Review Comment:
`toString` is used for logs/tests but currently only escapes backslashes and quotes. Keys containing other control characters (e.g., newline, tab) will produce non-portable / multiline output and can make assertions brittle. Consider a more complete string escaping strategy (e.g., JSON-style escaping for common control chars) to keep output deterministic and debuggable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5792900698
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4084305196
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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;
+ * - 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
Review Comment:
Building `subfields` via `.toMap` can silently drop entries if two scan outputs normalize to the same column name (e.g., duplicate column names or case-fold collisions). That can lead to incomplete required-subfields declarations and incorrect nulls for the losing column. Consider detecting collisions (group by rendered column name and require uniqueness) and skipping pruning for ambiguous names, or merging declarations only when they refer to the same underlying attribute.
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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;
+ * - 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)
+ logInfo(s"Applying scan map-key pruning: $subfields")
+ val newScan = scan.withRequiredMapSubfields(subfields)
+ newScan.copyTagsFrom(scan)
+ newScan
+ }
+ }
+
Review Comment:
Logging the full `subfields` map at INFO can be very verbose for large queries (many scans/paths) and may add noise in production logs. Consider logging only a summary at INFO (e.g., number of scans/columns/paths) and moving full details to DEBUG/TRACE.
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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;
+ * - 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)
+ logInfo(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 pathsToDeclare, 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`.
+ */
+ 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
+ })
+ }
Review Comment:
The emitted `Seq[SubfieldPath]` ordering depends on expression traversal/insertion order. Because scan transformer equality/hashCode and canonicalization can depend on this sequence order, non-canonical ordering can reduce plan reuse (e.g., exchange reuse) even when the declared set is identical. Consider sorting/deduplicating `rendered` into a stable canonical order (e.g., by `toString` or a structural comparator) before storing it on the scan.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4144164786
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,421 @@
+/*
+ * 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;
+ * - 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))
Review Comment:
With case-insensitive analysis this renders fields using Java/Locale.ROOT lowercasing, but the native `foldCase` and `SubstraitParser::parseType` use `folly::toLowerAscii`. For a field such as `Ärger`, the declaration becomes `c.ärger` while the Velox schema remains `c.Ärger`, so the selective reader cannot resolve the required path and the query can return null/missing data; the added umlaut test exercises this mismatch. Align the Scala and native folding (or conservatively skip non-ASCII names).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "JkSelf (via GitHub)" <gi...@apache.org>.
JkSelf commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4195590581
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,421 @@
+/*
+ * 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;
+ * - 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] =>
Review Comment:
@marin-ma This might be an issue, could you help check? `element_at(m, 'k')` filters are declared here, but the pruning may not take effect for them at runtime. Spark `ElementAt` seems to still be emitted as Velox `element_at`, whose `canPushdown()` is false. If the filter becomes the scan's remaining filter, `extractSubfields()` would return the bare column `m`. `FileDataSource` would then append it to the declared subfields, so the whole map might still be read. The PR seems to have switched only GetMapValue to the pushdown-capable `get_map_value`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5833054112
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-6021471135
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4104933268
##########
gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala:
##########
@@ -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)
Review Comment:
`requiredMapSubfields` is dropped by `doCanonicalize`: this `copy` does not pass the new field, so a BatchScan with `m["a"]` and one with `m["b"]` both canonicalize with an empty declaration. Spark uses canonicalized plans for exchange reuse, which can then reuse a scan/exchange materialized with the wrong key set and return nulls for the other branch. Preserve `requiredMapSubfields` in this copy so canonicalization and the custom equality remain consistent.
##########
gluten-delta/src/test/scala/org/apache/gluten/execution/DeltaSuite.scala:
##########
@@ -632,7 +661,62 @@ abstract class DeltaSuite extends WholeStageTransformerSuite {
}
}
- test("deletion vector on partitioned table") {
+ 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")
Review Comment:
This test is in the shared `DeltaSuite`, which is inherited by both `VeloxDeltaSuite` and `BoltDeltaSuite`. The map-key pruning rule and Velox config are only installed for Velox, so Bolt will reach this assertion with an empty `requiredMapSubfields` and fail; move the test to a Velox-only suite or gate the Velox-specific assertions/backend setup.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5910334329
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "marin-ma (via GitHub)" <gi...@apache.org>.
marin-ma commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4157553821
##########
gluten-substrait/src/main/scala/org/apache/gluten/execution/BatchScanExecTransformer.scala:
##########
@@ -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)
Review Comment:
Wrong. Scala copy keeps every field that is not named, so the declaration survives canonicalization. The file-source scan even passes it explicitly. Equality and hashCode include it too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5915258440
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "marin-ma (via GitHub)" <gi...@apache.org>.
marin-ma commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5949468001
@zhouyuan @JkSelf Could you help to review? Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #13109:
URL: https://github.com/apache/gluten/pull/13109#issuecomment-5798546771
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Prune map entries at the Parquet scan to the keys a query reads [gluten]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13109:
URL: https://github.com/apache/gluten/pull/13109#discussion_r4080904928
##########
cpp/velox/substrait/SubstraitToVeloxPlan.cc:
##########
@@ -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;
+ };
Review Comment:
`foldCase` uses `folly::toLowerAscii`, which only lowercases ASCII. On the JVM side, the rule’s comments/tests indicate case folding should follow Spark/ConverterUtils normalization (the suite includes a non-ASCII folding case), so using ASCII-only folding here can cause required-subfield column/field names (or schema names) to stop matching in case-insensitive mode, leading to required subfields being ignored or resolving to nulls. Use the same case-folding strategy as the Substrait schema/type parsing path (and as the Scala normalization), or remove this extra folding if Substrait names are already normalized consistently by the producer.
##########
backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala:
##########
@@ -152,6 +152,32 @@ 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 =>
+ 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())
+ }
Review Comment:
`requiredMapSubfields` redundantly encodes the column name both as the `Map` key (`column`) and inside each `SubfieldPath` (`path.column`). `packRequiredSubfields` currently ignores `path.column` entirely, which makes it easy for callers to accidentally construct inconsistent data that silently serializes under the `Map` key. Add a defensive check (e.g., require `path.column == column`) or remove the redundant `column` field from `SubfieldPath` and treat it purely as “elements under a known column” to prevent hard-to-debug mismatches.
##########
backends-velox/src/test/scala/org/apache/gluten/execution/ScanMapKeyPruningSuite.scala:
##########
@@ -0,0 +1,630 @@
+/*
+ * 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(
Review Comment:
This case asserts the *declared* required subfields for a non-ASCII field name under case-insensitive resolution, but it doesn’t validate that the native reader actually applies the declaration correctly end-to-end (i.e., that the declaration survives transport and matches Velox/Hive field resolution). Consider adding an end-to-end assertion similar to the “native reader applies a declaration” test, specifically for a non-ASCII folded struct field (and/or nested struct field selection like `c.<field>["a"]`), to catch mismatches between JVM normalization and native-side case folding.
##########
backends-velox/src/main/scala/org/apache/gluten/extension/ScanMapKeyPruning.scala:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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;
+ * - 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)
+ logInfo(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 pathsToDeclare, or the fragment's output. Stay whole.
Review Comment:
Fix grammar in comment: `does not pathsToDeclare` is ungrammatical and reads like an accidental symbol reference; rephrase to something like `does not analyze` / `does not handle`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org