diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index ee0d25d6d12..d6e81b85e11 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -34,7 +34,7 @@ use crate::parquet::parquet_support::is_hdfs_scheme; #[cfg(feature = "hdfs-opendal")] use crate::parquet::parquet_support::{create_hdfs_operator, prepare_object_store_with_configs}; use arrow::datatypes::{Schema, SchemaRef}; -use arrow::record_batch::RecordBatch; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use async_trait::async_trait; use datafusion::{ error::{DataFusionError, Result}, @@ -232,6 +232,8 @@ pub struct ParquetWriterExec { partition_id: i32, /// Column names to use in the output Parquet file column_names: Vec, + /// Catalyst's target schema, including nullability and Parquet field metadata. + output_schema: Option, /// Object store configuration options object_store_options: HashMap, /// Metrics @@ -252,6 +254,7 @@ impl ParquetWriterExec { compression: ParquetCompression, partition_id: i32, column_names: Vec, + output_schema: Option, object_store_options: HashMap, ) -> Result { // Preserve the input's partitioning so each partition writes its own file @@ -273,6 +276,7 @@ impl ParquetWriterExec { compression, partition_id, column_names, + output_schema, object_store_options, metrics: ExecutionPlanMetricsSet::new(), cache, @@ -453,6 +457,7 @@ impl ExecutionPlan for ParquetWriterExec { self.compression.clone(), self.partition_id, self.column_names.clone(), + self.output_schema.clone(), self.object_store_options.clone(), )?)), _ => Err(DataFusionError::Internal( @@ -483,14 +488,18 @@ impl ExecutionPlan for ParquetWriterExec { assert_eq!(input_schema.fields().len(), column_names.len()); - // Replace the generic column names (col_0, col_1, etc.) with the actual names - let fields: Vec<_> = input_schema - .fields() - .iter() - .enumerate() - .map(|(i, field)| Arc::new(field.as_ref().clone().with_name(&column_names[i]))) - .collect(); - let output_schema = Arc::new(arrow::datatypes::Schema::new(fields)); + // The input schema comes from the placeholder Scan and marks every top-level field + // nullable. Use Catalyst's target schema so Parquet repetition and field IDs match Spark. + // Keep the column-name-only path for plans serialized before output_schema was added. + let output_schema = self.output_schema.clone().unwrap_or_else(|| { + let fields: Vec<_> = input_schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| Arc::new(field.as_ref().clone().with_name(&column_names[i]))) + .collect(); + Arc::new(Schema::new(fields)) + }); // Generate part file name for this partition // If using FileCommitProtocol (work_dir is set), include task_attempt_id in the filename @@ -533,13 +542,18 @@ 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. Both schemas use the same Catalyst data types, and + // disabling field-name matching still recursively validates nested nullability; + // only nested field names and metadata are ignored. + RecordBatch::try_new_with_options( + Arc::clone(&schema_for_write), + batch.columns().to_vec(), + &RecordBatchOptions::new().with_match_field_names(false), + ) + .map_err(|e| { + DataFusionError::Execution(format!("Failed to rename batch columns: {}", e)) + })? } else { batch }; @@ -589,8 +603,14 @@ impl ExecutionPlan for ParquetWriterExec { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Int32Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{Array, Int32Array, ListArray, StringArray}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::prelude::SessionContext; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use parquet::basic::Repetition; + use parquet::file::reader::{FileReader, SerializedFileReader}; use std::sync::Arc; #[test] @@ -618,6 +638,145 @@ mod tests { ); } + #[tokio::test] + async fn test_parquet_writer_preserves_catalyst_schema_in_footer() -> Result<()> { + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(1), None]), + Some(vec![Some(2)]), + ]); + let input_schema = Arc::new(Schema::new(vec![ + Field::new("col_0", DataType::Int32, true), + Field::new("col_1", values.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&input_schema), + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(values)], + )?; + + let DataType::List(input_element) = input_schema.field(1).data_type() else { + panic!("expected list input"); + }; + let list_element = input_element + .as_ref() + .clone() + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "23".to_string(), + )])); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("required_id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "11".to_string(), + )])), + Field::new("values", DataType::List(Arc::new(list_element)), true).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "22".to_string())]), + ), + ])); + + let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let temp_dir = tempfile::tempdir()?; + let work_dir = format!("file://{}", temp_dir.path().display()); + let writer = ParquetWriterExec::try_new( + input, + work_dir.clone(), + work_dir, + None, + None, + ParquetCompression::None, + 0, + vec!["required_id".to_string(), "values".to_string()], + Some(output_schema), + HashMap::new(), + )?; + + let mut stream = writer.execute(0, SessionContext::new().task_ctx())?; + while stream.try_next().await?.is_some() {} + + let file = File::open(temp_dir.path().join("part-00000.parquet"))?; + let reader = SerializedFileReader::new(file)?; + let fields = reader + .metadata() + .file_metadata() + .schema_descr() + .root_schema() + .get_fields(); + let required_id = fields[0].get_basic_info(); + assert_eq!(required_id.repetition(), Repetition::REQUIRED); + assert_eq!(required_id.id(), 11); + + let list = &fields[1]; + assert_eq!(list.get_basic_info().repetition(), Repetition::OPTIONAL); + assert_eq!(list.get_basic_info().id(), 22); + let element = list.get_fields()[0].get_fields()[0].get_basic_info(); + assert_eq!(element.repetition(), Repetition::OPTIONAL); + assert_eq!(element.id(), 23); + + Ok(()) + } + + #[tokio::test] + async fn test_parquet_writer_rejects_mismatched_nested_nullability() -> Result<()> { + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(1)]), + Some(vec![Some(2)]), + ]); + let DataType::List(input_element) = values.data_type() else { + panic!("expected list input"); + }; + assert!(input_element.is_nullable()); + + let input_schema = Arc::new(Schema::new(vec![Field::new( + "col_0", + values.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&input_schema), vec![Arc::new(values)])?; + + let target_element = Field::new("element", DataType::Int32, false).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "23".to_string())]), + ); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "values", + DataType::List(Arc::new(target_element)), + true, + )])); + + let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let temp_dir = tempfile::tempdir()?; + let work_dir = format!("file://{}", temp_dir.path().display()); + let writer = ParquetWriterExec::try_new( + input, + work_dir.clone(), + work_dir, + None, + None, + ParquetCompression::None, + 0, + vec!["values".to_string()], + Some(output_schema), + HashMap::new(), + )?; + + let mut stream = writer.execute(0, SessionContext::new().task_ctx())?; + let error = stream + .try_next() + .await + .expect_err("mismatched nested nullability must be rejected"); + let message = error.to_string(); + assert!( + message.contains("Failed to rename batch columns"), + "unexpected error: {message}" + ); + assert!( + message.contains("column types must match schema types"), + "unexpected error: {message}" + ); + + Ok(()) + } + /// Helper function to create a test RecordBatch with 1000 rows of (int, string) data /// Example batch_id 1 -> 0..1000, 2 -> 1001..2000 #[allow(dead_code)] @@ -828,10 +987,6 @@ mod tests { #[cfg(feature = "hdfs-opendal")] #[ignore = "This test requires a running HDFS cluster"] async fn test_parquet_writer_exec_with_memory_input() -> Result<()> { - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::prelude::SessionContext; - // Create 5 batches for the DataSourceExec input let mut batches = Vec::new(); for i in 1..=5 { @@ -860,6 +1015,7 @@ mod tests { ParquetCompression::None, 0, // partition_id column_names, + None, // output_schema HashMap::new(), // object_store_options )?; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c179c3b57c5..b35612a88ae 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1866,6 +1866,8 @@ impl PhysicalPlanner { codec, self.partition, writer.column_names.clone(), + (!writer.output_schema.is_empty()) + .then(|| convert_spark_types_to_arrow_schema(&writer.output_schema)), object_store_options, )?); diff --git a/native/core/src/execution/serde.rs b/native/core/src/execution/serde.rs index f31cb9cd35a..89d23c9c3b0 100644 --- a/native/core/src/execution/serde.rs +++ b/native/core/src/execution/serde.rs @@ -31,8 +31,9 @@ use datafusion_comet_proto::{ spark_expression::DataType, spark_operator, }; +use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use prost::Message; -use std::{io::Cursor, sync::Arc}; +use std::{collections::HashMap, io::Cursor, sync::Arc}; /// Deserialize bytes to protobuf type of expression pub fn deserialize_expr(buf: &[u8]) -> Result { @@ -115,10 +116,13 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType { .unwrap() { DatatypeStruct::List(info) => { - let field = Field::new( - "item", - to_arrow_datatype(info.element_type.as_ref().unwrap()), - info.contains_null, + let field = with_parquet_field_id( + Field::new( + "item", + to_arrow_datatype(info.element_type.as_ref().unwrap()), + info.contains_null, + ), + info.element_field_id, ); ArrowDataType::List(Arc::new(field)) } @@ -133,15 +137,21 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType { .unwrap() { DatatypeStruct::Map(info) => { - let key_field = Field::new( - "key", - to_arrow_datatype(info.key_type.as_ref().unwrap()), - false, + let key_field = with_parquet_field_id( + Field::new( + "key", + to_arrow_datatype(info.key_type.as_ref().unwrap()), + false, + ), + info.key_field_id, ); - let value_field = Field::new( - "value", - to_arrow_datatype(info.value_type.as_ref().unwrap()), - info.value_contains_null, + let value_field = with_parquet_field_id( + Field::new( + "value", + to_arrow_datatype(info.value_type.as_ref().unwrap()), + info.value_contains_null, + ), + info.value_field_id, ); let struct_field = Field::new( "entries", @@ -187,3 +197,139 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType { }, } } + +/// Attach a Parquet field ID without changing synthetic fields when Catalyst did not supply one. +fn with_parquet_field_id(field: Field, field_id: Option) -> Field { + match field_id { + Some(id) => field.with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])), + None => field, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_comet_proto::spark_expression::data_type::{DataTypeInfo, ListInfo, MapInfo}; + + fn primitive_type(type_id: DataTypeId) -> DataType { + DataType { + type_id: type_id as i32, + type_info: None, + } + } + + fn list_type(element_type: DataType, element_field_id: Option) -> DataType { + DataType { + type_id: DataTypeId::List as i32, + type_info: Some(Box::new(DataTypeInfo { + datatype_struct: Some(DatatypeStruct::List(Box::new(ListInfo { + element_type: Some(Box::new(element_type)), + contains_null: true, + element_field_id, + }))), + })), + } + } + + fn map_type(key_field_id: Option, value_field_id: Option) -> DataType { + DataType { + type_id: DataTypeId::Map as i32, + type_info: Some(Box::new(DataTypeInfo { + datatype_struct: Some(DatatypeStruct::Map(Box::new(MapInfo { + key_type: Some(Box::new(primitive_type(DataTypeId::Int32))), + value_type: Some(Box::new(primitive_type(DataTypeId::String))), + value_contains_null: true, + key_field_id, + value_field_id, + }))), + })), + } + } + + #[test] + fn list_element_field_id_preserves_zero_and_absence() { + let ArrowDataType::List(field) = + to_arrow_datatype(&list_type(primitive_type(DataTypeId::Int32), Some(0))) + else { + panic!("expected a list data type"); + }; + + assert_eq!(field.name(), "item"); + assert!(field.is_nullable()); + assert_eq!( + field.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"0".to_string()) + ); + + let ArrowDataType::List(field_without_id) = + to_arrow_datatype(&list_type(primitive_type(DataTypeId::Int32), None)) + else { + panic!("expected a list data type"); + }; + assert!(field_without_id.metadata().is_empty()); + } + + #[test] + fn map_key_and_value_field_ids_preserve_nullability() { + let ArrowDataType::Map(entries, _) = to_arrow_datatype(&map_type(Some(0), Some(23))) else { + panic!("expected a map data type"); + }; + let ArrowDataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries to be a struct"); + }; + + assert_eq!(fields[0].name(), "key"); + assert!(!fields[0].is_nullable()); + assert_eq!( + fields[0].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"0".to_string()) + ); + assert_eq!(fields[1].name(), "value"); + assert!(fields[1].is_nullable()); + assert_eq!( + fields[1].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"23".to_string()) + ); + + let ArrowDataType::Map(entries_without_ids, _) = to_arrow_datatype(&map_type(None, None)) + else { + panic!("expected a map data type"); + }; + let ArrowDataType::Struct(fields_without_ids) = entries_without_ids.data_type() else { + panic!("expected map entries to be a struct"); + }; + assert!(fields_without_ids[0].metadata().is_empty()); + assert!(fields_without_ids[1].metadata().is_empty()); + } + + #[test] + fn synthetic_field_ids_are_preserved_at_multiple_nesting_levels() { + let ArrowDataType::List(element) = + to_arrow_datatype(&list_type(map_type(Some(12), Some(13)), Some(11))) + else { + panic!("expected a list data type"); + }; + assert_eq!( + element.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"11".to_string()) + ); + + let ArrowDataType::Map(entries, _) = element.data_type() else { + panic!("expected a nested map data type"); + }; + let ArrowDataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries to be a struct"); + }; + assert_eq!( + fields[0].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"12".to_string()) + ); + assert_eq!( + fields[1].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"13".to_string()) + ); + } +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ced87262f32..d4ba2bf053e 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -475,6 +475,9 @@ message ParquetWriter { // configuration value "spark.hadoop.fs.s3a.access.key" will be stored as "fs.s3a.access.key" in // the map. map object_store_options = 8; + // Catalyst's target schema, including top-level nullability and Parquet field IDs. + // Nested collection field IDs are carried by the corresponding DataType messages. + repeated SparkStructField output_schema = 9; } enum AggregateMode { diff --git a/native/proto/src/proto/types.proto b/native/proto/src/proto/types.proto index 114253ee779..643a9cbabb1 100644 --- a/native/proto/src/proto/types.proto +++ b/native/proto/src/proto/types.proto @@ -83,12 +83,17 @@ message DataType { message ListInfo { DataType element_type = 1; bool contains_null = 2; + // Parquet field ID for the synthetic list element, when provided by Catalyst metadata. + optional int32 element_field_id = 3; } message MapInfo { DataType key_type = 1; DataType value_type = 2; bool value_contains_null = 3; + // Parquet field IDs for the synthetic map fields, when provided by Catalyst metadata. + optional int32 key_field_id = 4; + optional int32 value_field_id = 5; } message StructInfo { diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 85601a9e0c9..6802dfaa646 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -553,7 +553,19 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * doesn't mean it is supported by Comet native execution, i.e., `supportedDataType` may return * false for it. */ - def serializeDataType(dt: org.apache.spark.sql.types.DataType): Option[Types.DataType] = { + def serializeDataType(dt: org.apache.spark.sql.types.DataType): Option[Types.DataType] = + serializeDataType(dt, None, Seq.empty, includeFieldIds = true) + + /** + * Preserves collection field IDs stored on the nearest Catalyst StructField when serializing a + * Parquet write schema. Delta stores synthetic list and map IDs under paths relative to that + * field, and resets the path whenever a nested struct introduces a new StructField. + */ + private[comet] def serializeDataType( + dt: org.apache.spark.sql.types.DataType, + parentField: Option[StructField], + fieldPath: Seq[String], + includeFieldIds: Boolean): Option[Types.DataType] = { val typeId = dt match { case _: BooleanType => 0 case _: ByteType => 1 @@ -595,7 +607,9 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { builder.setTypeInfo(info.build()).build() case a: ArrayType => - val elementType = serializeDataType(a.elementType) + val elementPath = fieldPath :+ "element" + val elementType = + serializeDataType(a.elementType, parentField, elementPath, includeFieldIds) if (elementType.isEmpty) { return None @@ -605,17 +619,21 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val list = ListInfo.newBuilder() list.setElementType(elementType.get) list.setContainsNull(a.containsNull) + nestedParquetFieldId(parentField, elementPath, includeFieldIds) + .foreach(list.setElementFieldId) info.setList(list) builder.setTypeInfo(info.build()).build() case m: MapType => - val keyType = serializeDataType(m.keyType) + val keyPath = fieldPath :+ "key" + val keyType = serializeDataType(m.keyType, parentField, keyPath, includeFieldIds) if (keyType.isEmpty) { return None } - val valueType = serializeDataType(m.valueType) + val valuePath = fieldPath :+ "value" + val valueType = serializeDataType(m.valueType, parentField, valuePath, includeFieldIds) if (valueType.isEmpty) { return None } @@ -625,6 +643,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { map.setKeyType(keyType.get) map.setValueType(valueType.get) map.setValueContainsNull(m.valueContainsNull) + nestedParquetFieldId(parentField, keyPath, includeFieldIds).foreach(map.setKeyFieldId) + nestedParquetFieldId(parentField, valuePath, includeFieldIds).foreach(map.setValueFieldId) info.setMap(map) builder.setTypeInfo(info.build()).build() @@ -634,7 +654,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val struct = StructInfo.newBuilder() val fieldNames = s.map(_.name).asJava - val fieldDatatypes = s.map(f => serializeDataType(f.dataType)) + val fieldDatatypes = s.map { field => + val nestedParentField = parentField.map(_ => field) + val nestedFieldPath = if (nestedParentField.isDefined) Seq(field.name) else Seq.empty + serializeDataType(field.dataType, nestedParentField, nestedFieldPath, includeFieldIds) + } val fieldNullable = s.map(f => Boolean.box(f.nullable)).asJava if (fieldDatatypes.exists(_.isEmpty)) { @@ -646,7 +670,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { struct.addAllFieldNullable(fieldNullable) val fieldIds = s.fields.map { f => - if (ParquetUtils.hasFieldId(f)) Some(ParquetUtils.getFieldId(f)) else None + if (includeFieldIds && ParquetUtils.hasFieldId(f)) Some(ParquetUtils.getFieldId(f)) + else None } if (fieldIds.exists(_.isDefined)) { // Emit one FieldMetadata entry per nested field, parallel to field_names. Entries @@ -668,6 +693,30 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { Some(dataType) } + private def nestedParquetFieldId( + parentField: Option[StructField], + fieldPath: Seq[String], + includeFieldIds: Boolean): Option[Int] = { + val nestedIdsMetadataKey = "parquet.field.nested.ids" + if (!includeFieldIds) { + None + } else { + parentField.flatMap { field => + if (field.metadata.contains(nestedIdsMetadataKey)) { + val nestedIds = field.metadata.getMetadata(nestedIdsMetadataKey) + val nestedPath = fieldPath.mkString(".") + if (nestedIds.contains(nestedPath)) { + Some(Math.toIntExact(nestedIds.getLong(nestedPath))) + } else { + None + } + } else { + None + } + } + } + } + def aggExprToProto( aggExpr: AggregateExpression, inputs: Seq[Attribute], diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index 8157f286825..97e44dff675 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -135,6 +135,10 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec .setOutputPath(outputPath) .setCompression(codec) .addAllColumnNames(cmd.query.output.map(_.name).asJava) + .addAllOutputSchema(schema2Proto( + cmd.query.schema.fields.toIndexedSeq, + Some( + op.session.sessionState.conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))).asJava) // Note: work_dir, job_id, and task_attempt_id will be set at execution time // in CometNativeWriteExec, as they depend on the Spark task context diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index cb7702083b5..cf6e3fabe8d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -30,14 +30,22 @@ import org.apache.comet.shims.ShimFileFormat package object operator { - def schema2Proto(fields: Seq[StructField]): Seq[OperatorOuterClass.SparkStructField] = { + def schema2Proto( + fields: Seq[StructField], + fieldIdWriteEnabled: Option[Boolean] = None): Seq[OperatorOuterClass.SparkStructField] = { val fieldBuilder = OperatorOuterClass.SparkStructField.newBuilder() fields.map { field => fieldBuilder.setName(field.name) - fieldBuilder.setDataType(serializeDataType(field.dataType).get) + val dataType = fieldIdWriteEnabled match { + case Some(includeFieldIds) => + serializeDataType(field.dataType, Some(field), Seq(field.name), includeFieldIds) + case None => + serializeDataType(field.dataType) + } + fieldBuilder.setDataType(dataType.get) fieldBuilder.setNullable(field.nullable) fieldBuilder.clearMetadata() - if (ParquetUtils.hasFieldId(field)) { + if (fieldIdWriteEnabled.getOrElse(true) && ParquetUtils.hasFieldId(field)) { fieldBuilder.putMetadata( CometParquetUtils.PARQUET_FIELD_ID_META_KEY, ParquetUtils.getFieldId(field).toString) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index a1ae1af1d1c..eef77d88246 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -28,12 +28,14 @@ import org.apache.hadoop.fs.{FileSystem, Path} import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.metadata.CompressionCodecName import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.sql.{AnalysisException, CometTestBase, DataFrame, Row, SaveMode} import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan} import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.functions.{array, map, struct, when} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus @@ -143,6 +145,216 @@ 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 nested columns are renamed") { + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val numberMetadata = parquetFieldMetadata(11) + val textMetadata = parquetFieldMetadata(22) + val structMetadata = parquetFieldMetadata(100) + val structChildMetadata = parquetFieldMetadata(101) + val itemsMetadata = parquetFieldMetadata(200, "original_items.element" -> 201L) + val itemChildMetadata = parquetFieldMetadata(202) + val lookupMetadata = + parquetFieldMetadata(300, "original_lookup.key" -> 301L, "original_lookup.value" -> 302L) + val data = spark + .range(1, 3) + .select( + $"id".as("original_number", numberMetadata), + $"id".cast(StringType).as("original_text", textMetadata), + struct($"id".as("original_child", structChildMetadata)) + .as("original_struct", structMetadata), + array(struct($"id".as("original_item_child", itemChildMetadata))) + .as("original_items", itemsMetadata), + map($"id".cast(StringType), $"id") + .as("original_lookup", lookupMetadata)) + + withNativeWriter { + withSQLConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true") { + val plan = captureWritePlan(path => data.write.parquet(path), outputPath) + assertHasCometNativeWriteExec(plan) + } + } + + val renamedSchema = StructType( + Seq( + StructField( + "renamed_lookup", + MapType(StringType, LongType, valueContainsNull = true), + nullable = true, + metadata = parquetFieldMetadata( + 300, + "renamed_lookup.key" -> 301L, + "renamed_lookup.value" -> 302L)), + StructField( + "renamed_items", + ArrayType( + StructType( + Seq( + StructField( + "renamed_item_child", + LongType, + nullable = true, + metadata = itemChildMetadata))), + containsNull = true), + nullable = true, + metadata = parquetFieldMetadata(200, "renamed_items.element" -> 201L)), + StructField( + "renamed_struct", + StructType( + Seq( + StructField( + "renamed_child", + LongType, + nullable = true, + metadata = structChildMetadata))), + nullable = true, + metadata = structMetadata), + StructField("renamed_text", StringType, nullable = true, metadata = textMetadata), + StructField("renamed_number", LongType, nullable = true, metadata = numberMetadata))) + + // Array element and map key/value names are structural in Spark, so rename the + // collection fields and the struct field inside the array element instead. + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") { + checkAnswer( + spark.read.schema(renamedSchema).parquet(outputPath), + Seq( + Row(Map("1" -> 1L), Seq(Row(1L)), Row(1L), "1", 1L), + Row(Map("2" -> 2L), Seq(Row(2L)), Row(2L), "2", 2L))) + } + } + } + test("parquet write with each supported compression codec") { Seq("none", "uncompressed", "snappy", "lz4", "zstd", "gzip").foreach { codec => withTempPath { dir => @@ -828,6 +1040,29 @@ class CometParquetWriterSuite extends CometTestBase { case other => fail(s"unexpected codec: $other") } + private def parquetFieldMetadata(id: Long, nestedIds: (String, Long)*): Metadata = { + val metadata = new MetadataBuilder().putLong("parquet.field.id", id) + if (nestedIds.nonEmpty) { + val nestedMetadata = new MetadataBuilder() + nestedIds.foreach { case (name, nestedId) => nestedMetadata.putLong(name, nestedId) } + metadata.putMetadata("parquet.field.nested.ids", nestedMetadata.build()) + } + metadata.build() + } + + private def assertParquetSchemas(outputPath: String)(verify: MessageType => Unit): Unit = { + val conf = spark.sparkContext.hadoopConfiguration + val partFiles = new File(outputPath).listFiles().filter(_.getName.startsWith("part-")) + assert(partFiles.nonEmpty, s"No part files found under $outputPath") + + partFiles.foreach { partFile => + val inputFile = HadoopInputFile.fromPath(new Path(partFile.getAbsolutePath), conf) + Using.resource(ParquetFileReader.open(inputFile)) { reader => + verify(reader.getFooter.getFileMetaData.getSchema) + } + } + } + /** * Asserts that every column chunk in every part file under `outputPath` reports `expected` as * its compression codec. Reading the data back is not enough on its own: a Parquet reader