dongjoon-hyun commented on code in PR #59972:
URL: https://github.com/apache/spark/pull/59972#discussion_r4430398461
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -690,7 +693,13 @@ object PushDownUtils extends Logging {
scanBuilder match {
case r: SupportsPushDownRequiredColumns if
SQLConf.get.nestedSchemaPruningEnabled =>
- val rootFields = SchemaPruning.identifyRootFields(projects, filters)
+ val pruneDynamicSelectors = scanBuilder match {
Review 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
```scala
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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetReadSupport.scala:
##########
@@ -106,20 +106,61 @@ class ParquetReadSupport(
new ParquetRecordMaterializer(
parquetRequestedSchema,
ParquetReadSupport.expandUDT(catalystRequestedSchema),
- new ParquetToSparkSchemaConverter(conf),
+ new ParquetReadSchemaConverter(conf, parquetRequestedSchema, fileSchema),
convertTz,
datetimeRebaseSpec,
int96RebaseSpec)
}
}
+/**
+ * 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(
+ conf: Configuration,
+ requestedSchema: MessageType,
+ fileSchema: MessageType) extends ParquetToSparkSchemaConverter(conf) {
+
+ private lazy val fileRepeatedTypes = {
+ // Equal requested types at different paths can have different layouts in
the file.
+ val types = new util.IdentityHashMap[Type, Type]()
+ def visit(requested: GroupType, file: GroupType): Unit = {
+ if
(file.getLogicalTypeAnnotation.isInstanceOf[ListLogicalTypeAnnotation] &&
+ requested.getFieldCount == 1 && file.getFieldCount == 1) {
+ types.put(requested.getType(0), file.getType(0))
+ }
+ requested.getFields.asScala.foreach {
+ case group: GroupType if file.containsField(group.getName) =>
+ file.getType(group.getName) match {
+ case fileGroup: GroupType => visit(group, fileGroup)
+ case _ =>
+ }
+ case _ =>
+ }
+ }
+ visit(requestedSchema, fileSchema)
+ types
+ }
+
+ override private[parquet] def isElementType(repeatedType: Type, parentName:
String): Boolean = {
+ super.isElementType(fileRepeatedTypes.getOrDefault(repeatedType,
repeatedType), parentName)
+ }
+}
+
object ParquetReadSupport extends Logging {
val SPARK_ROW_REQUESTED_SCHEMA =
"org.apache.spark.sql.parquet.row.requested_schema"
val SPARK_METADATA_KEY = "org.apache.spark.sql.parquet.row.metadata"
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 = {
Review 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`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala:
##########
@@ -35,6 +41,207 @@ abstract class ParquetSchemaPruningSuite extends
SchemaPruningSuite with Adaptiv
override protected val vectorizedReaderNestedEnabledKey: String =
SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key
+ test("dynamic array selection preserves two-level Parquet LIST elements
after pruning") {
+ withTempPath { file =>
+ // After clipping legacy to just wanted, both LISTs have identical
repeated groups.
+ val fileSchema = MessageTypeParser.parseMessageType(
+ """message m {
+ | optional group legacy (LIST) {
+ | repeated group entries { optional int32 sibling = 18; optional
int32 wanted = 17; }
+ | }
+ | optional group standard (LIST) {
+ | repeated group entries { optional int32 wanted = 17; }
+ | }
+ | optional int32 idx;
+ |}""".stripMargin)
+ val writer = ExampleParquetWriter.builder(new Path(file.toURI))
+
.withConf(spark.sessionState.newHadoopConf()).withType(fileSchema).build()
+ try {
+ val row = new SimpleGroupFactory(fileSchema).newGroup().append("idx",
1)
+ val legacy = row.addGroup("legacy")
+ legacy.addGroup("entries").append("sibling", 10).append("wanted", 11)
+ legacy.addGroup("entries").append("sibling", 20).append("wanted", 22)
+ val standard = row.addGroup("standard")
+ standard.addGroup("entries").append("wanted", 30)
+ standard.addGroup("entries").append("wanted", 33)
+ writer.write(row)
+ } finally {
+ writer.close()
+ }
+ val fieldId = new
MetadataBuilder().putLong(ParquetUtils.FIELD_ID_METADATA_KEY, 17).build()
+ val missingId = new
MetadataBuilder().putLong(ParquetUtils.FIELD_ID_METADATA_KEY, 19).build()
+ val element = new StructType()
+ .add("sibling", IntegerType, nullable = true, metadata = missingId)
+ .add("renamed", IntegerType, nullable = true, metadata = fieldId)
+ val readSchema = new StructType().add("legacy", ArrayType(element))
+ .add("standard", ArrayType(IntegerType)).add("idx", IntegerType)
+ Seq("true", "false").foreach { vectorized =>
+ withSQLConf(
+ vectorizedReaderEnabledKey -> vectorized,
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true",
+ SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "true") {
+ val data =
spark.read.schema(readSchema).parquet(file.getCanonicalPath)
+ val query = data.selectExpr("legacy[idx].renamed", "standard[idx]")
+ checkScanSchemata(query,
+
"struct<legacy:array<struct<renamed:int>>,standard:array<int>,idx:int>")
+ checkAnswer(query, Row(22, 33))
+
+ // Only the missing ID is requested. The first physical field is
read for levels
+ // and shares its name, but has a different ID.
+ checkAnswer(data.selectExpr("legacy[idx].sibling"), Row(null))
+ }
+ }
+ }
+ }
+
+ test("a two-level Parquet LIST preserves enclosing struct nullness after
pruning") {
+ withTempPath { file =>
+ val fileSchema = MessageTypeParser.parseMessageType(
+ """message m {
+ | optional group record {
+ | optional group items (LIST) {
+ | repeated group entries { optional int32 x; optional int32 y; }
+ | }
+ | }
+ |}""".stripMargin)
+ val writer = ExampleParquetWriter.builder(new Path(file.toURI))
+
.withConf(spark.sessionState.newHadoopConf()).withType(fileSchema).build()
+ val factory = new SimpleGroupFactory(fileSchema)
+ try {
+ val row = factory.newGroup()
+
row.addGroup("record").addGroup("items").addGroup("entries").append("x",
1).append("y", 2)
+ writer.write(row)
+ writer.write(factory.newGroup())
+ } finally {
+ writer.close()
+ }
+ Seq("true", "false").foreach { vectorized =>
+ withSQLConf(vectorizedReaderEnabledKey -> vectorized) {
+ // The only physical child available for the record's levels is the
clipped LIST.
+ val data = spark.read.schema("record STRUCT<missing:
INT>").parquet(file.getCanonicalPath)
+ checkAnswer(data.selectExpr("record.missing", "record IS NULL"),
+ Seq(Row(null, false), Row(null, true)))
+ }
+ }
+ }
+ }
+
+ test("dynamic array selection preserves root and nested Parquet field IDs") {
+ withSQLConf(
+ SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
+ withTempPath { dir =>
+ val rootId = new MetadataBuilder()
+ .putLong(ParquetUtils.FIELD_ID_METADATA_KEY, 1).build()
+ val leafId = new MetadataBuilder()
+ .putLong(ParquetUtils.FIELD_ID_METADATA_KEY, 17).build()
+ val leafType = new StructType()
+ .add("old_id", IntegerType, nullable = true, metadata = leafId)
+ .add("payload", StringType)
+ def schema(rootName: String, leaf: StructType): StructType = new
StructType()
+ .add(rootName, ArrayType(new StructType().add("children",
ArrayType(leaf))),
+ nullable = true, metadata = rootId)
+ .add("idx", IntegerType)
+ val writeSchema = schema("old_items", leafType)
+ val readSchema = schema("new_items", StructType(Seq(
+ leafType("old_id").copy(name = "new_id"), leafType("payload"))))
+ val rows = Seq(Row(Seq(Row(Seq(Row(42, "unused")))), 0))
+ val path = dir.getCanonicalPath
+ spark.createDataFrame(spark.sparkContext.parallelize(rows),
writeSchema)
+ .write.parquet(path)
+
+ Seq("true", "false").foreach { vectorized =>
+ withSQLConf(vectorizedReaderEnabledKey -> vectorized) {
+ val query = spark.read.schema(readSchema).parquet(path)
+ .selectExpr("new_items[idx].children.new_id")
+ checkScanSchemata(query,
+
"struct<new_items:array<struct<children:array<struct<new_id:int>>>>,idx:int>")
+ checkAnswer(query, Row(Seq(42)))
+ }
+ }
+ }
+ }
+ }
+
+ test("dynamic lookups preserve complex column defaults and constant lookups
still prune") {
+ withTempPath { dir =>
+ withTable("t") {
+ val path = dir.getCanonicalPath
+ sql(
+ s"""CREATE TABLE t (idx INT, k STRING,
+ | items ARRAY<STRUCT<a: INT, b: INT>> DEFAULT array(),
+ | byKey MAP<STRING, STRUCT<a: INT, b: INT>> DEFAULT map()
+ |) USING parquet LOCATION '$path'""".stripMargin)
+ sql("INSERT INTO t SELECT 0, 'k', array(named_struct('a', 5, 'b', 6)),
" +
+ "map('k', named_struct('a', 7, 'b', 8))")
+ sql(
+ """ALTER TABLE t ADD COLUMNS (
+ | arr ARRAY<STRUCT<a: INT, b: INT>> DEFAULT
array(named_struct('a', 11, 'b', 22)),
+ | m MAP<STRING, STRUCT<a: INT, b: INT>>
+ | DEFAULT map('k', named_struct('a', 33, 'b', 44))
+ |)""".stripMargin)
+
+ val data = spark.read.schema(spark.table("t").schema).parquet(path)
+ Seq("true", "false").foreach { vectorized =>
+ withSQLConf(vectorizedReaderEnabledKey -> vectorized) {
+ val query = data.selectExpr("arr[idx].b", "m[k].b", "items[0].b",
"byKey['k'].b")
+ checkAnswer(query, Row(22, 44, 6, 8))
+ checkScanSchemata(query,
"struct<idx:int,k:string,items:array<struct<b:int>>," +
+
"byKey:map<string,struct<b:int>>,arr:array<struct<a:int,b:int>>," +
+ "m:map<string,struct<a:int,b:int>>>")
+ }
+ }
+ }
+ }
+ }
+
+ test("dynamic lookups preserve defaults on nested fields") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ sql(
+ """SELECT 0 AS i,
+ | named_struct('arr', array(named_struct('a', 1, 'b', 22))) AS s,
+ | map(0, named_struct('anchor', 7, 'n', named_struct('a', 3, 'b',
44))) AS m
+ |""".stripMargin).write.parquet(path)
+
+ val pair = StructType.fromDDL("a INT, b INT")
+ val default = "named_struct('a', 5, 'b', 6)"
+ val arr = StructField("arr",
ArrayType(pair)).withExistenceDefaultValue(s"array($default)")
+ val n = StructField("n", pair).withExistenceDefaultValue(default)
+ val schema = new StructType()
+ .add("i", IntegerType)
+ .add("s", new StructType().add(arr))
+ .add("m", MapType(IntegerType, new StructType().add("anchor",
IntegerType).add(n)))
+
+ withSQLConf(vectorizedReaderEnabledKey -> "false") {
Review 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?
##########
sql/core/src/test/scala/org/apache/spark/sql/DatasetOptimizationSuite.scala:
##########
@@ -132,6 +132,29 @@ class DatasetOptimizationSuite extends SharedSparkSession {
}
}
+ test("dynamic map lookup uses the fields emitted by the serializer") {
+ 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"))
Review 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`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala:
##########
@@ -81,16 +81,22 @@ object SelectedField {
* Convert an expression into the parts of the schema (the field) it
accesses.
*/
@scala.annotation.tailrec
- private def selectField(expr: Expression, dataTypeOpt: Option[DataType]):
Option[StructField] = {
+ private def selectField(
+ expr: Expression,
+ dataTypeOpt: Option[DataType],
+ hasDynamicLookup: Boolean = false): Option[StructField] = {
Review 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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -67,7 +67,17 @@ object SchemaPruning extends Rule[LogicalPlan] {
val normalizedProjects = normalizeAttributeRefNames(attrNameMap, projects)
.asInstanceOf[Seq[NamedExpression]]
val normalizedFilters = normalizeAttributeRefNames(attrNameMap, filters)
- val requestedRootFields = identifyRootFields(normalizedProjects,
normalizedFilters)
+ val pruneDynamicSelectors = hadoopFsRelation.fileFormat match {
+ case _: OrcFileFormat => false
+ case _: ParquetFileFormat
+ if (projects.iterator ++ filters.iterator).exists(hasDynamicLookup)
=>
+ val hadoopConf = hadoopFsRelation.sparkSession.sessionState
Review 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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
