sunchao commented on code in PR #59972:
URL: https://github.com/apache/spark/pull/59972#discussion_r4432330733
##########
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:
Moved `supportsDynamicSelectorPruning` to `FileScanBuilder` in 6ac1708fe40.
ORC and Parquet now declare their own behavior. `PushDownUtils` no longer
imports either format.
##########
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:
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.
##########
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:
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.
##########
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:
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.
##########
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:
Updated the vectorized reader to call
`ParquetReadSupport.returnNullStructIfAllFieldsMissing`. It now uses the same
helper as the row reader and planner.
##########
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:
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`.
--
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]
