Skip to content

[SPARK-60766][SQL] Prune nested fields through dynamic array indices and map keys - #59972

Draft
sunchao wants to merge 12 commits into
apache:masterfrom
sunchao:codex/dynamic-selector-pruning-oss-20261030
Draft

sunchao wants to merge 12 commits into
apache:masterfrom
sunchao:codex/dynamic-selector-pruning-oss-20261030

Conversation

@sunchao

@sunchao sunchao commented Oct 30, 2026 •

Copy link
Copy Markdown
Member

What changes were proposed in this pull request?

Allow Spark to prune nested fields through a row-dependent array index or Map key. For example, a Parquet scan for events[event_index].id can read just id from each element and still read the column that supplies event_index. Spark retains fields the selector itself needs, even when it reads from the same collection. Map lookups keep the complete keys they compare.

Review order. Two existing correctness bugs exposed by the review have independent fixes against master. They can be reviewed and backported separately, in either order.

  1. #59996 — bind Dataset projections to the serializer's actual output
  2. #59997 — preserve Parquet LIST layout and field IDs after pruning
  3. Dynamic-selector changes only

This branch includes both prerequisites so the feature can be tested end to end. GitHub's Files changed tab includes them until they land on master. This PR should merge after both.

Why are the changes needed?

SPARK-42163 limited nested pruning to constant selectors because pruning could discard fields used by the index or key. Tracking those dependencies lets eligible scans skip unused siblings. SPARK-60766

Does this PR introduce any user-facing change?

Eligible scans request fewer nested fields. Results and ANSI array bounds errors are preserved. Dynamic lookups keep the previous read schema on ORC, in Parquet's legacy missing-field mode, and where narrowing a root could change how a stored default is read. There is no new configuration or public API.

How was this patch tested?

  • 131 tests passed across seven Catalyst, Dataset, Parquet V1/V2, and ORC V1/V2 suites against Spark master's 5.0.0-SNAPSHOT CI build. Both test files changed by the split were recompiled against that build. All production sources are identical to the previously reviewed head.
  • A further 207 tests passed across five Parquet field-ID, schema, inference, and I/O suites on the same build. Changed-file ScalaStyle and git diff --check passed.
  • The dynamic cases check results, the requested scan schema, and ANSI errors. Coverage includes selector dependencies in the same collection, null predicates on Map keys, typed Dataset outputs, two-level LISTs, renamed and missing field IDs, stored defaults on old files, and ORC positional evolution. The standalone PRs include constant-selector regressions that reproduce on their unmodified parent.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-6)

@sunchao

sunchao commented Oct 31, 2026

Copy link
Copy Markdown
Member Author

cc @viirya @cloud-fan @peter-toth @dongjoon-hyun could you review this PR? thanks!

@dongjoon-hyun dongjoon-hyun left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for working on this, @sunchao. I reviewed the latest commit, 04b5528, and left 6 inline comments. The dynamic-lookup pruning itself looks right to me: the fields that an index or key reads are collected separately, null checks keep the map keys they compare, and the Parquet readers take the LIST layout from the file schema. These points are about where the format checks live, how the default guard is wired, two tests, and two nits. Here is a summary, ordered by priority.

API and design

  1. PushDownUtils L696: this match makes the shared DSv2 push-down code import OrcScanBuilder and ParquetScanBuilder, although FileScanBuilder.supportsNestedSchemaPruning already lets each builder declare such a trait itself.
  2. SelectedField L87: hasDynamicLookup is threaded through all 14 recursive calls of selectField with a default of false, so a new extractor case that forgets it still compiles and silently turns the default guard off; its foldability test also repeats SchemaPruning.DynamicLookup.

Tests

  1. DatasetOptimizationSuite L139: the new test checks only the answer, so it still passes if the serializer is not pruned at all.
  2. ParquetSchemaPruningSuite L217: the nested-defaults test runs only the row reader, and nothing says why.

Cleanup

  1. ParquetReadSupport L159: nit, VectorizedParquetRecordReader still reads this key inline instead of calling the new helper.

Performance

  1. SchemaPruning L74: nit, this builds a full Hadoop Configuration during optimization to read one SQL conf.

The most important is 1. Items 2-4 would make the new paths harder to break, and items 5 and 6 are small.

scanBuilder match {
case r: SupportsPushDownRequiredColumns if SQLConf.get.nestedSchemaPruningEnabled =>
val rootFields = SchemaPruning.identifyRootFields(projects, filters)
val pruneDynamicSelectors = scanBuilder match {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[1] Could the scan builders declare this instead of PushDownUtils matching on them?

PushDownUtils is shared by every DSv2 source, and on master it imports no file-format classes. This match makes it import ParquetReadSupport, OrcScanBuilder and ParquetScanBuilder (L37-39) and know how each reader maps nested fields. The builders already declare traits like this themselves: FileScanBuilder.supportsNestedSchemaPruning (FileScanBuilder.scala:39) is overridden by both. Could FileScanBuilder get a member such as private[sql] def supportsDynamicSelectorPruning: Boolean = true, overridden with false in OrcScanBuilder and with !ParquetReadSupport.returnNullStructIfAllFieldsMissing(hadoopConf) in ParquetScanBuilder? This match would then become

        val pruneDynamicSelectors = scanBuilder match {
          case f: FileScanBuilder => f.supportsDynamicSelectorPruning
          case _ => true
        }

and a new file format, or a change to the Parquet condition, would not need to touch this shared code. The V1 match in SchemaPruning.scala:70-78 can stay as it is, since canPruneDataSchema there (SchemaPruning.scala:122) already checks the file format.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved supportsDynamicSelectorPruning to FileScanBuilder in 6ac1708. ORC and Parquet now declare their own behavior. PushDownUtils no longer imports either format.

private def selectField(
expr: Expression,
dataTypeOpt: Option[DataType],
hasDynamicLookup: Boolean = false): Option[StructField] = {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[2] Could the default guard be applied once, instead of threading hasDynamicLookup through every case?

The flag is passed on in all 14 recursive calls of selectField (L99-190). Because it defaults to false, a new extractor case that forgets to pass it still compiles and silently turns the default guard off for that path. Its foldability test (!key.foldable at L126, !index.foldable at L155, !selector.foldable at L166 and L169) also repeats DynamicLookup (SchemaPruning.scala:155-162). Could selectField return the root attribute along with the field, so that unapply (L56) and withDataType (L70) apply the default check once, with one shared test for dynamic lookups? selectField would then be a pure extractor again. If that shared test is SchemaPruning.hasDynamicLookup (SchemaPruning.scala:140-141), it also sees lookups inside lambda functions, which is slightly more conservative than today.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the flag from the recursive calls. Both entry points now check stored defaults once, using the same dynamic-lookup matcher as scan planning. The check follows the collection input: for filter(items, x -> array(true, false)[idx])[0].b, the index inside the lambda should still let Spark read only items.b. The new scan-schema assertion fails on both V1 and V2 if the check walks the whole expression.

val ds = spark.range(1).map { _ =>
((Map("k" -> ("map-a", "map-b")), ("other-a", "other-b")), "k")
}
checkAnswer(ds.selectExpr("_1._1[_2]._2", "_1._2._1"), Row("map-b", "other-a"))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[3] This test still passes if the serializer is not pruned at all.

It checks only the answer, and the answer is the same whether or not ObjectSerializerPruning prunes _1, so a regression back to no pruning for a dynamic map key would go unnoticed. Could it also check the pruned serializer with testSerializer (L48), as the other tests in this suite do, e.g. that the map values keep only _2 and _1._2 keeps only _1?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a testSerializer assertion for _1._1[_2]._2 using the same Dataset. I checked that upstream returns the same row but fails the new assertion because the map value still contains both fields. The original two-output query can legitimately keep the full serializer when it also requests the sibling, so I kept its value check.

.add("s", new StructType().add(arr))
.add("m", MapType(IntegerType, new StructType().add("anchor", IntegerType).add(n)))

withSQLConf(vectorizedReaderEnabledKey -> "false") {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[4] Why does this test run only with the row reader?

The other new dynamic-lookup tests in this suite run both readers (L78, L118, L154, L186 and L230), and the only other row-only test says so in its name ("row reader uses the scan's field ID setting", L280). The two readers evaluate existence defaults in different places (ParquetRowConverter.scala:238-239 and VectorizedParquetRecordReader.java:301), so the vectorized path for nested defaults, and for a dynamic lookup that shares a map with a constant lookup, is not covered here. Could it loop over both readers like the test above, or, if the vectorized reader is left out on purpose, could a comment say why?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The nested-default regression now runs with both row and vectorized readers in both the Parquet V1 and V2 suites. This includes the case where a constant and a dynamic lookup share a map.

def generateFakeColumnName: String = s"_fake_name_${UUID.randomUUID()}"

// In legacy mode, pruning all physical children can also lose the enclosing array length.
def returnNullStructIfAllFieldsMissing(conf: Configuration): Boolean = {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[5] Nit: VectorizedParquetRecordReader still reads this key inline.

VectorizedParquetRecordReader.java:274-277 reads the same key with the same default by hand, so the plan-time checks that now use this helper (SchemaPruning.scala:76 and PushDownUtils.scala:698) and that reader could drift apart if only one copy changes. Could it call the helper, e.g. ParquetReadSupport$.MODULE$.returnNullStructIfAllFieldsMissing(configuration), the way SpecificParquetRecordReaderBase.java:144 reads SPARK_ROW_REQUESTED_SCHEMA?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the vectorized reader to call ParquetReadSupport.returnNullStructIfAllFieldsMissing. It now uses the same helper as the row reader and planner.

case _: OrcFileFormat => false
case _: ParquetFileFormat
if (projects.iterator ++ filters.iterator).exists(hasDynamicLookup) =>
val hadoopConf = hadoopFsRelation.sparkSession.sessionState

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[6] Nit: this builds a full Hadoop Configuration during optimization to read one SQL conf.

newHadoopConfWithOptions copies the session's Hadoop configuration and sets every SQL conf on it, and this rule does that for every V1 Parquet scan with a dynamic lookup. Could it read the value directly and let the relation's options win, e.g. hadoopFsRelation.options.get(entry.key).map(_.toBoolean).getOrElse(conf.getConf(entry)) with entry = SQLConf.LEGACY_PARQUET_RETURN_NULL_STRUCT_IF_ALL_FIELDS_MISSING? If the point is to resolve it exactly as the reader does, through the same getHadoopConf(relation.sparkSession, relation.options) that DataSourceScanExec.scala:797 uses, a short comment saying so would be enough.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kept the effective Hadoop configuration and added the suggested comment. I reproduced a difference with the legacy flag set through spark.hadoop.* and the SQL setting unset. Reading only relation options and SQLConf made an out-of-bounds access return NULL on both V1 readers. Using the same configuration as the reader preserved INVALID_ARRAY_INDEX.

@viirya viirya left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @sunchao. The dynamic-selector pruning itself looks correct to me. My probes cover same-root indices, indices that read the same collection, filter-only lookups, lambda-variable indices, and element_at/try_element_at. Each returns the parent's rows from a narrower scan on V1 and V2 with both Parquet readers. dongjoon's points are addressed in 6ac1708.

My main concern is scope. Several changes fix bugs that already exist on master, independent of dynamic lookups. Some guards cover only the dynamic path while the same failure remains for constant lookups. Could the master bugs go into separate JIRAs so they can be backported, leaving this PR as the feature? Details inline.

* Interpret LISTs using the file schema. Clipping a repeated group to one field can make a
* two-level LIST indistinguishable from a three-level LIST in the requested schema.
*/
private[parquet] class ParquetReadSchemaConverter(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fixes a master bug that has nothing to do with dynamic lookups. On master, with the two-level LIST schema from the new test, legacy[1].renamed fails on both readers. With field IDs, the row reader returns [10, 20] for a plain legacy.sibling instead of [null, null]. Could this go into its own JIRA so it can be backported?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. I split the reader fix into #59997 so it can be backported independently. I reproduced both cases with constant projections on master. The two-level LIST read failed in V1 and V2, and the row reader returned [10, 20] for a missing field ID where the vectorized reader returned [NULL, NULL]. The standalone change fixes both. Its separate JIRA is pending.

val rootFields = SchemaPruning.identifyRootFields(p.projectList, Seq.empty)
// Resolved references can use the spelling from the query. Both pruning and rebuilding
// field accesses need the names from the serializer, including when a root is used twice.
val namesById = s.output.map(a => a.exprId -> a.name).toMap

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two changes also fix master bugs:

  • Under case-insensitive analysis, items[0]._2 and items._2 throw ArrayIndexOutOfBoundsException.
  • With a constant key, ds.selectExpr("_1._1['k']._2", "_1._2._1") on the new test's data returns map-a instead of map-b.

Same request: a separate JIRA.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. #59996 isolates the Dataset serializer fix and has regression tests with constant selectors. On master, the constant Map query returned (map-a, other-a) rather than the requested (map-b, other-a), and the case-insensitive array reference threw ArrayIndexOutOfBoundsException. Both pass with the standalone fix. Its separate JIRA is pending.

newConverter(parquetField, catalystType, NoopUpdater)
// Only the enclosing struct's levels matter. Its clipped child may no longer carry
// enough information to infer its logical type (for example, a two-level LIST).
ParquetRowConverter.discardConverter(parquetField)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"a two-level Parquet LIST preserves enclosing struct nullness after pruning" and "row reader uses the scan's field ID setting" both pass on master without this PR, so I can't find what this change, or the schemaConverter.useFieldId change at L225, fixes on master. Were they needed only on the 4.1.3 base? If so, could they be dropped? Also, V1 still ignores the per-read field-ID option.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right that both tests pass on unmodified master. They protect interactions introduced by the LIST/field-ID fix, so I moved them to #59997. I rebuilt that patch on Spark master's CI build and reverted each change separately. Removing the discard converter made the enclosing-struct test fail because Parquet requested a group converter and received a primitive converter. Reading the session's field-ID setting instead of the scan's made the V2 row reader return NULL where the query expected 0. V1 still ignores the per-read option. The option regression test is V2-only.

selectField(unaliased, Some(dataType))
// A separate projection can read a nested default under this same root. Preserve its
// declared type for dynamic lookups, including when the default is on a sibling.
selectField(unaliased, dataType).filterNot { case (root, field) =>

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This failure is not specific to dynamic lookups. After ADD COLUMNS (arr ARRAY<STRUCT<a,b>> DEFAULT array(named_struct('a', 11, 'b', 22))), arr[0].b, arr.b and even a struct column's s.b fail with INVALID_DEFAULT_VALUE.DATA_TYPE on both master and this PR, because the existence default is coerced to the pruned type.

arr[idx].b requests the same read schema as arr[0].b, so this guard hides the bug on one path only. The constant-lookup assertions pass only because DEFAULT array()/map() fit any pruned type.

Could the coercion be fixed for all nested pruning in a separate change, and hasDynamicLookupOnPath/hasNestedDefault be dropped here?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I reproduced the underlying bug on master with the nonempty default. arr[0].b, arr.b, and the defaulted struct's s.b all fail with INVALID_DEFAULT_VALUE.DATA_TYPE. I also tried removing the guard here. On older files, arr[idx].b then changed from 22 to that same error, on V1 and V2 with both Parquet readers. I kept the guard for this feature. I agree the general fix belongs in a separate change that resolves the default using the declared type and then projects it. The follow-up JIRA is pending.

val normalizedFilters = normalizeAttributeRefNames(attrNameMap, filters)
val requestedRootFields = identifyRootFields(normalizedProjects, normalizedFilters)
val pruneDynamicSelectors = hadoopFsRelation.fileFormat match {
case _: OrcFileFormat => false

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Constant lookups already fail on ORC on master:

  • V2 returns 11 instead of 22 for arr[0].wanted under positional evolution.
  • V1 fails in the timestamp_ntz case.

Excluding dynamic lookups doesn't fix either. Could they be filed as bugs? (Same for OrcScanBuilder.supportsDynamicSelectorPruning.)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed both on master with row and vectorized readers. V2 positional evolution returned the sibling's 11 for arr[0].wanted when the requested value was 22. V1 raised UNSUPPORTED_FEATURE.ORC_TYPE_CAST for the constant lookup of the TIMESTAMP_NTZ field. With the index supplied by a column, the same queries currently return 22 and the requested timestamp. I kept the V1/V2 exclusions to preserve those results. I prepared separate JIRA reports for the positional and timestamp bugs. The keys are pending.


override protected val supportsNestedSchemaPruning: Boolean = true

override private[sql] def supportsDynamicSelectorPruning: Boolean =

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In legacy missing-field mode under ANSI, the constant items[1].missing already returns null, as the conf's doc warns. Is it worth the V1/V2 plumbing to keep dynamic lookups stricter than constant ones?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would keep the V1/V2 check. I tested a one-element array with an index column containing 1. Under ANSI and legacy missing-field mode, items[1].missing already returns NULL, as you noted. But items[idx].missing currently raises INVALID_ARRAY_INDEX on both V1 and V2, with both Parquet readers. Removing the check silently turned that existing exception into NULL in all four combinations. Keeping the collection shape for this mode preserves the observable behavior of the query we're optimizing.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

3 participants