parthchandra commented on code in PR #5369:
URL: https://github.com/apache/datafusion-comet/pull/5369#discussion_r3799521316
##########
native/core/src/execution/operators/parquet_writer.rs:
##########
@@ -533,13 +542,17 @@ impl ExecutionPlan for ParquetWriterExec {
// Rename columns in the batch to match output schema
let renamed_batch = if !column_names.is_empty() {
- RecordBatch::try_new(Arc::clone(&schema_for_write),
batch.columns().to_vec())
- .map_err(|e| {
- DataFusionError::Execution(format!(
- "Failed to rename batch columns: {}",
- e
- ))
- })?
+ // Collection field IDs exist on the target schema, not on
arrays produced by
+ // the placeholder Scan. Ignore nested field
names/metadata during validation
+ // while still checking data types and nested/top-level
nullability.
Review Comment:
Is the nested nullability guaranteed to match with the target schema here?
##########
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:
Maybe also add coverage for nested types. Maybe a test that renames a
struct field and a list element or map key/value?
--
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]
