sunchao commented on code in PR #5369:
URL: https://github.com/apache/datafusion-comet/pull/5369#discussion_r3799830847
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -143,6 +145,168 @@ class CometParquetWriterSuite extends CometTestBase {
}
}
+ test(
+ "native parquet writer preserves Catalyst nullability and honors field ID
write settings") {
+ val requiredMetadata = parquetFieldMetadata(11)
+ val optionalMetadata = parquetFieldMetadata(22)
+ val data = spark
+ .range(0, 3)
+ .select(
+ $"id".as("required_number", requiredMetadata),
+ when($"id" === 1L, $"id".cast(StringType)).as("optional_text",
optionalMetadata),
+ $"id".as("unmapped_number"))
+
+ assert(!data.schema("required_number").nullable)
+ assert(data.schema("optional_text").nullable)
+
+ Seq(None, Some(true), Some(false)).foreach { configuredValue =>
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+
+ withNativeWriter {
+ def writeAndVerify(): Unit = {
+ val plan = captureWritePlan(path => data.write.parquet(path),
outputPath)
+ assertHasCometNativeWriteExec(plan)
+
+ val expectedIds = configuredValue.getOrElse(true)
+ assertParquetSchemas(outputPath) { schema =>
+ val root = schema.asGroupType()
+ val required = root.getType("required_number")
+ val optional = root.getType("optional_text")
+ val unmapped = root.getType("unmapped_number")
+
+ assert(required.getRepetition == Type.Repetition.REQUIRED)
+ assert(optional.getRepetition == Type.Repetition.OPTIONAL)
+ assert(
+ Option(required.getId).map(_.intValue()) ==
+ (if (expectedIds) Some(11) else None))
+ assert(
+ Option(optional.getId).map(_.intValue()) ==
+ (if (expectedIds) Some(22) else None))
+ assert(unmapped.getId == null)
+ }
+ }
+
+ configuredValue match {
+ case Some(enabled) =>
+ withSQLConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key ->
enabled.toString) {
+ writeAndVerify()
+ }
+ case None =>
+
assert(spark.conf.get(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key).toBoolean)
+ writeAndVerify()
+ }
+ }
+ }
+ }
+ }
+
+ test("native parquet writer preserves nested and Delta collection field
IDs") {
+ val detailsMetadata = parquetFieldMetadata(100)
+ val requiredChildMetadata = parquetFieldMetadata(101)
+ val optionalChildMetadata = parquetFieldMetadata(102)
+ val innerMetadata = parquetFieldMetadata(130, "inner.element" -> 131L)
+ val tagsMetadata = parquetFieldMetadata(200, "tags.element" -> 201L)
+ val attrsMetadata =
+ parquetFieldMetadata(300, "attrs.key" -> 301L, "attrs.value" -> 302L)
+
+ val data = spark
+ .range(0, 2)
+ .select(
+ struct(
+ $"id".as("required_child", requiredChildMetadata),
+ when($"id" === 1L, $"id").as("optional_child",
optionalChildMetadata),
+ array($"id").as("inner", innerMetadata)).as("details",
detailsMetadata),
+ array(when($"id" === 1L, $"id")).as("tags", tagsMetadata),
+ map($"id".cast(StringType), when($"id" === 1L, $"id")).as("attrs",
attrsMetadata))
+
+ Seq(true, false).foreach { writeFieldIds =>
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+
+ withNativeWriter {
+ withSQLConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key ->
writeFieldIds.toString) {
+ val plan = captureWritePlan(path => data.write.parquet(path),
outputPath)
+ assertHasCometNativeWriteExec(plan)
+
+ assertParquetSchemas(outputPath) { schema =>
+ val root = schema.asGroupType()
+
+ def assertField(field: Type, id: Int, nullable: Boolean): Unit =
{
+ val expectedRepetition =
+ if (nullable) Type.Repetition.OPTIONAL else
Type.Repetition.REQUIRED
+ assert(field.getRepetition == expectedRepetition)
+ assert(Option(field.getId).map(_.intValue()) ==
+ (if (writeFieldIds) Some(id) else None))
+ }
+
+ val details = root.getType("details")
+ assertField(details, 100, nullable = false)
+ val detailsGroup = details.asGroupType()
+ assertField(detailsGroup.getType("required_child"), 101,
nullable = false)
+ assertField(detailsGroup.getType("optional_child"), 102,
nullable = true)
+
+ val inner = detailsGroup.getType("inner")
+ assertField(inner, 130, nullable = false)
+ val innerList = inner.asGroupType().getType(0)
+ assert(innerList.getRepetition == Type.Repetition.REPEATED)
+ assertField(innerList.asGroupType().getType(0), 131, nullable =
false)
+
+ val tags = root.getType("tags")
+ assertField(tags, 200, nullable = false)
+ val tagsList = tags.asGroupType().getType(0)
+ assert(tagsList.getRepetition == Type.Repetition.REPEATED)
+ assertField(tagsList.asGroupType().getType(0), 201, nullable =
true)
+
+ val attrs = root.getType("attrs")
+ assertField(attrs, 300, nullable = false)
+ val entries = attrs.asGroupType().getType(0)
+ assert(entries.getRepetition == Type.Repetition.REPEATED)
+ val entriesGroup = entries.asGroupType()
+ assertField(entriesGroup.getType("key"), 301, nullable = false)
+ assertField(entriesGroup.getType("value"), 302, nullable = true)
+ }
+ }
+ }
+
+ checkAnswer(spark.read.parquet(outputPath), data)
+ }
+ }
+ }
+
+ test("Spark reads native parquet output by field ID after columns are
renamed and reordered") {
Review Comment:
Added a field-ID read/write round trip covering a renamed nested struct
field, a renamed field inside an `array<struct>`, and a renamed map alongside
reordered top-level columns. The array element and map key/value IDs are
preserved as well. Spark treats `element`, `key`, and `value` as structural
names, so those fields cannot be renamed independently; their IDs are verified
directly in the existing Parquet-footer regression.
--
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]
