Repository navigation
Conversation
|
cc @viirya @cloud-fan @peter-toth @dongjoon-hyun could you review this PR? thanks! |
dongjoon-hyun
left a comment
There was a problem hiding this comment.
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
PushDownUtilsL696: this match makes the shared DSv2 push-down code importOrcScanBuilderandParquetScanBuilder, althoughFileScanBuilder.supportsNestedSchemaPruningalready lets each builder declare such a trait itself.SelectedFieldL87:hasDynamicLookupis threaded through all 14 recursive calls ofselectFieldwith a default offalse, so a new extractor case that forgets it still compiles and silently turns the default guard off; its foldability test also repeatsSchemaPruning.DynamicLookup.
Tests
DatasetOptimizationSuiteL139: the new test checks only the answer, so it still passes if the serializer is not pruned at all.ParquetSchemaPruningSuiteL217: the nested-defaults test runs only the row reader, and nothing says why.
Cleanup
ParquetReadSupportL159: nit,VectorizedParquetRecordReaderstill reads this key inline instead of calling the new helper.
Performance
SchemaPruningL74: nit, this builds a full HadoopConfigurationduring 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 { |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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] = { |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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")) |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
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") { |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
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 = { |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
These two changes also fix master bugs:
- Under case-insensitive analysis,
items[0]._2anditems._2throwArrayIndexOutOfBoundsException. - With a constant key,
ds.selectExpr("_1._1['k']._2", "_1._2._1")on the new test's data returnsmap-ainstead ofmap-b.
Same request: a separate JIRA.
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
"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.
There was a problem hiding this comment.
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) => |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
Constant lookups already fail on ORC on master:
- V2 returns 11 instead of 22 for
arr[0].wantedunder positional evolution. - V1 fails in the
timestamp_ntzcase.
Excluding dynamic lookups doesn't fix either. Could they be filed as bugs? (Same for OrcScanBuilder.supportsDynamicSelectorPruning.)
There was a problem hiding this comment.
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 = |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
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].idcan read justidfrom each element and still read the column that suppliesevent_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.
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?
git diff --checkpassed.Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-6)