Apply suggestions from code review

getChan · jecsand838 · web-flow · commit d55f945eef0d · 2025-11-26T11:42:39.000+09:00
Co-authored-by: Connor Sanders &lt;170039284+jecsand838@users.noreply.github.com&gt;
diff --git a/datafusion/datasource-avro/src/mod.rs b/datafusion/datasource-avro/src/mod.rs
@@ -33,12 +33,187 @@ use arrow::datatypes::Schema;
 pub use arrow_avro;
 use arrow_avro::reader::ReaderBuilder;
 pub use file_format::*;
+use arrow_avro::schema::SCHEMA_METADATA_KEY;
+use datafusion_common::DataFusionError;
 use std::io::{BufReader, Read};
+use std::sync::Arc;
 
 /// Read Avro schema given a reader
 pub fn read_avro_schema_from_reader<R: Read>(
     reader: &mut R,
 ) -> datafusion_common::Result<Schema> {
     let avro_reader = ReaderBuilder::new().build(BufReader::new(reader))?;
-    Ok(avro_reader.schema().as_ref().clone())
+    let schema_ref = avro_reader.schema();
+    // Extract the raw Avro JSON schema from the OCF header.
+    let raw_json = avro_reader
+        .avro_header()
+        .get(SCHEMA_METADATA_KEY.as_bytes())
+        .map(|bytes| {
+            std::str::from_utf8(bytes).map_err(|e| {
+                DataFusionError::Execution(format!(
+                    "Invalid UTF-8 in Avro schema metadata ({SCHEMA_METADATA_KEY}): {e}"
+                ))
+            })
+        })
+        .transpose()?
+        .map(str::to_owned);
+    drop(avro_reader);
+    if let Some(raw_json) = raw_json {
+        let mut schema = Arc::unwrap_or_clone(schema_ref);
+        // Insert the raw Avro JSON schema using `SCHEMA_METADATA_KEY`.
+        // This should enable the avro schema metadata to be picked downstream.
+        schema
+            .metadata
+            .insert(SCHEMA_METADATA_KEY.to_string(), raw_json);
+        Ok(schema)
+    } else {
+        // Return error because Avro spec requires the Avro schema metadata to be present in the OCF header.
+        Err(DataFusionError::Execution(format!(
+            "Avro schema metadata ({SCHEMA_METADATA_KEY}) is missing from OCF header"
+        )))
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use arrow::array::{BinaryArray, BooleanArray, Float64Array};
+    use arrow::datatypes::DataType;
+    use arrow_avro::reader::ReaderBuilder;
+    use arrow_avro::schema::{AvroSchema, SCHEMA_METADATA_KEY};
+    use datafusion_common::test_util::arrow_test_data;
+    use datafusion_common::{DataFusionError, Result as DFResult};
+    use serde_json::Value;
+    use std::collections::HashMap;
+    use std::fs::File;
+    use std::io::BufReader;
+
+    fn avro_test_file(name: &str) -> String {
+        format!("{}/avro/{name}", arrow_test_data())
+    }
+
+    #[test]
+    fn read_avro_schema_includes_avro_json_metadata() -> DFResult<()> {
+        let path = avro_test_file("alltypes_plain.avro");
+        let mut file = File::open(&path)?;
+        let schema = read_avro_schema_from_reader(&mut file)?;
+        let meta_json = schema
+            .metadata()
+            .get(SCHEMA_METADATA_KEY)
+            .expect("schema metadata missing avro.schema entry");
+        assert!(
+            !meta_json.is_empty(),
+            "avro.schema metadata should not be empty"
+        );
+        let mut raw = File::open(&path)?;
+        let avro_reader = ReaderBuilder::new().build(BufReader::new(&mut raw))?;
+        let header_json = avro_reader
+            .avro_header()
+            .get(SCHEMA_METADATA_KEY.as_bytes())
+            .and_then(|bytes| std::str::from_utf8(bytes).ok())
+            .expect("missing avro.schema metadata in OCF header");
+        assert_eq!(
+            meta_json, header_json,
+            "schema metadata avro.schema should match OCF header"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn read_and_project_using_schema_metadata() -> DFResult<()> {
+        let path = avro_test_file("alltypes_dictionary.avro");
+        let mut file = File::open(&path)?;
+        let file_schema = read_avro_schema_from_reader(&mut file)?;
+        let projected_field_names = vec!["string_col", "double_col", "bool_col"];
+        let avro_json = file_schema
+            .metadata()
+            .get(SCHEMA_METADATA_KEY)
+            .expect("schema metadata missing avro.schema entry");
+        let projected_avro_schema =
+            build_projected_reader_schema(avro_json, &projected_field_names)?;
+        let mut reader = ReaderBuilder::new()
+            .with_reader_schema(projected_avro_schema)
+            .with_batch_size(64)
+            .build(BufReader::new(File::open(&path)?))?;
+        let batch = reader.next().expect("no batch produced")?;
+        assert_eq!(3, batch.num_columns());
+        assert_eq!(2, batch.num_rows());
+        let schema = batch.schema();
+        assert_eq!("string_col", schema.field(0).name());
+        assert_eq!(&DataType::Binary, schema.field(0).data_type());
+        let col = batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<BinaryArray>()
+            .expect("column 0 not BinaryArray");
+        assert_eq!("0".as_bytes(), col.value(0));
+        assert_eq!("1".as_bytes(), col.value(1));
+        assert_eq!("double_col", schema.field(1).name());
+        assert_eq!(&DataType::Float64, schema.field(1).data_type());
+        let col = batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .expect("column 1 not Float64Array");
+        assert_eq!(0.0, col.value(0));
+        assert_eq!(10.1, col.value(1));
+        assert_eq!("bool_col", schema.field(2).name());
+        assert_eq!(&DataType::Boolean, schema.field(2).data_type());
+        let col = batch
+            .column(2)
+            .as_any()
+            .downcast_ref::<BooleanArray>()
+            .expect("column 2 not BooleanArray");
+        assert!(col.value(0));
+        assert!(!col.value(1));
+        Ok(())
+    }
+
+    fn build_projected_reader_schema(
+        avro_json: &str,
+        projected_field_names: &[&str],
+    ) -> DFResult<AvroSchema> {
+        let mut schema_json: Value = serde_json::from_str(avro_json).map_err(|e| {
+            DataFusionError::Execution(format!(
+                "Failed to parse Avro schema JSON from metadata: {e}"
+            ))
+        })?;
+        let obj = schema_json.as_object_mut().ok_or_else(|| {
+            DataFusionError::Execution(
+                "Top-level Avro schema JSON is not an object".to_string(),
+            )
+        })?;
+        let fields_val = obj.get_mut("fields").ok_or_else(|| {
+            DataFusionError::Execution(
+                "Top-level Avro schema JSON has no `fields` key".to_string(),
+            )
+        })?;
+        let fields = fields_val.as_array_mut().ok_or_else(|| {
+            DataFusionError::Execution(
+                "Top-level Avro schema `fields` is not an array".to_string(),
+            )
+        })?;
+        let mut by_name: HashMap<String, Value> = HashMap::new();
+        for field in fields.iter() {
+            if let Some(name) = field.get("name").and_then(|v| v.as_str()) {
+                by_name.insert(name.to_string(), field.clone());
+            }
+        }
+        let mut projected_fields = Vec::with_capacity(projected_field_names.len());
+        for name in projected_field_names {
+            let Some(field) = by_name.get(*name) else {
+                return Err(DataFusionError::Execution(format!(
+                    "Projected field `{name}` not found in Avro writer schema"
+                )));
+            };
+            projected_fields.push(field.clone());
+        }
+        *fields_val = Value::Array(projected_fields);
+        let projected_json = serde_json::to_string(&schema_json).map_err(|e| {
+            DataFusionError::Execution(format!(
+                "Failed to serialize projected Avro schema JSON: {e}"
+            ))
+        })?;
+        Ok(AvroSchema::new(projected_json))
+    }
 }
diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs
@@ -58,22 +58,99 @@ impl AvroSource {
     }
 
     fn open<R: std::io::BufRead>(&self, reader: R) -> Result<Reader<R>> {
-        let schema = self.table_schema.file_schema().as_ref(); // todo - avro metadata loading
-
-        let projected_schema = if let Some(projection) = &self.file_projection {
-            &schema.project(projection)?
-        } else {
-            schema
-        };
-
-        let avro_schema = AvroSchema::try_from(projected_schema)?;
-
+        // TODO: Once `ReaderBuilder::with_projection` is available, we should use it instead.
+        //  This should be an easy change. We'd simply need to:
+        //      1. Use the full file schema to generate the reader `AvroSchema`.
+        //      2. Pass `&self.file_projection` into `ReaderBuilder::with_projection`.
+        //      3. Remove the `build_projected_reader_schema` methods.
         ReaderBuilder::new()
-            .with_reader_schema(avro_schema) // Used for projection on read.
+            .with_reader_schema(self.build_projected_reader_schema()?)
             .with_batch_size(self.batch_size.expect("Batch size must set before open"))
             .build(reader)
             .map_err(Into::into)
     }
+
+    fn build_projected_reader_schema(&self) -> Result<AvroSchema> {
+        let file_schema = self.table_schema.file_schema().as_ref();
+        // Fast path: no projection. If we have the original writer schema JSON
+        // in metadata, just reuse it as-is without parsing.
+        if self.file_projection.is_none() {
+            return if let Some(avro_json) =
+                file_schema.metadata().get(SCHEMA_METADATA_KEY)
+            {
+                Ok(AvroSchema::new(avro_json.clone()))
+            } else {
+                // Fall back to deriving Avro from the full Arrow file schema, should be ok
+                // if not using projection.
+                Ok(AvroSchema::try_from(file_schema)
+                    .map_err(Into::<DataFusionError>::into)?)
+            };
+        }
+        // Use the writer Avro schema JSON tagged upstream to build a projected reader schema
+        match file_schema.metadata().get(SCHEMA_METADATA_KEY) {
+            Some(avro_json) => {
+                let mut schema_json: Value =
+                    serde_json::from_str(avro_json).map_err(|e| {
+                        DataFusionError::Execution(format!(
+                            "Failed to parse Avro schema JSON from metadata: {e}"
+                        ))
+                    })?;
+                let obj = schema_json.as_object_mut().ok_or_else(|| {
+                    DataFusionError::Execution(
+                        "Top-level Avro schema JSON must be an object".to_string(),
+                    )
+                })?;
+                let fields_val = obj.get_mut("fields").ok_or_else(|| {
+                    DataFusionError::Execution(
+                        "Top-level Avro schema JSON must contain a `fields` array"
+                            .to_string(),
+                    )
+                })?;
+                let fields_arr = fields_val.as_array_mut().ok_or_else(|| {
+                    DataFusionError::Execution(
+                        "Top-level Avro schema `fields` must be an array".to_string(),
+                    )
+                })?;
+                // Move existing fields out so we can rebuild them in projected order.
+                let original_fields = std::mem::take(fields_arr);
+                let mut by_name: HashMap<String, Value> =
+                    HashMap::with_capacity(original_fields.len());
+                for field in original_fields {
+                    if let Some(name) = field.get("name").and_then(|v| v.as_str()) {
+                        by_name.insert(name.to_string(), field);
+                    }
+                }
+                // Rebuild `fields` in the same order as the projected Arrow schema.
+                let projection = self.file_projection.as_ref().ok_or_else(|| {
+                    DataFusionError::Internal("checked file_projection is Some above".to_string())
+                })?;
+                let projected_schema = file_schema.project(projection)?;
+                let mut projected_fields =
+                    Vec::with_capacity(projected_schema.fields().len());
+                for arrow_field in projected_schema.fields() {
+                    let name = arrow_field.name();
+                    let field = by_name.remove(name).ok_or_else(|| {
+                        DataFusionError::Execution(format!(
+                            "Projected field `{name}` not found in Avro writer schema"
+                        ))
+                    })?;
+                    projected_fields.push(field);
+                }
+                *fields_val = Value::Array(projected_fields);
+                let projected_json =
+                    serde_json::to_string(&schema_json).map_err(|e| {
+                        DataFusionError::Execution(format!(
+                            "Failed to serialize projected Avro schema JSON: {e}"
+                        ))
+                    })?;
+                Ok(AvroSchema::new(projected_json))
+            }
+            None => Err(DataFusionError::Execution(format!(
+                "Avro schema metadata ({SCHEMA_METADATA_KEY}) is missing from file schema, but is required for projection"
+            ))),
+        }
+    }
+}
 }
 
 impl FileSource for AvroSource {
