Merge branch 'main' into issue-9497-list-json-reader-bench

Author: liamzwbao

arrow-avro/Cargo.toml

@@ -70,7 +70,7 @@ zstd = { version = "0.13", default-features = false, optional = true }
 bzip2 = { version = "0.6.0", optional = true }
 xz = { package = "liblzma", version = "0.4", default-features = false, optional = true }
 crc = { version = "3.0", optional = true }
-strum_macros = "0.27"
+strum_macros = "0.28"
 uuid = "1.17"
 indexmap = "2.10"
 rand = "0.9"

arrow-avro/benches/project_record.rs

@@ -121,7 +121,22 @@ fn gen_double(mut rng: impl Rng, sc: &ApacheSchema, n: usize, prefix: &[u8]) ->
     )
 }
 
-const READER_SCHEMA: &str = r#"
+fn gen_mixed(mut rng: impl Rng, sc: &ApacheSchema, n: usize, prefix: &[u8]) -> Vec<u8> {
+    encode_records_with_prefix(
+        sc,
+        prefix,
+        (0..n).map(|i| {
+            Value::Record(vec![
+                ("f1".into(), Value::Int(rng.random())),
+                ("f2".into(), Value::Long(rng.random())),
+                ("f3".into(), Value::String(format!("name-{i}"))),
+                ("f4".into(), Value::Double(rng.random())),
+            ])
+        }),
+    )
+}
+
+const SKIP_READER_SCHEMA: &str = r#"
     {
         "type":"record",
         "name":"table",
@@ -175,11 +190,42 @@ const DOUBLE_SCHEMA: &str = r#"
     }
     "#;
 
-fn new_decoder(schema_json: &'static str, batch_size: usize) -> Decoder {
+const MIX_SCHEMA: &str = r#"
+    {
+        "type":"record",
+        "name":"Mix",
+        "fields": [
+            { "name": "f1", "type": "int" },
+            { "name": "f2", "type": "long" },
+            { "name": "f3", "type": "string" },
+            { "name": "f4", "type": "double" }
+        ]
+    }
+    "#;
+
+// Project the record type writen to MIX_SCHEMA:
+// skip "f2" and "f4", add "f5" with a default
+const PROJECT_READER_SCHEMA: &str = r#"
+    {
+        "type":"record",
+        "name":"Mix",
+        "fields": [
+            { "name": "f1", "type": "int" },
+            { "name": "f3", "type": "string" },
+            { "name": "f5", "type": "long", "default": 0 }
+        ]
+    }
+    "#;
+
+fn new_decoder(
+    schema_json: &'static str,
+    reader_schema_json: &'static str,
+    batch_size: usize,
+) -> Decoder {
     let schema = AvroSchema::new(schema_json.to_owned());
     let mut store = SchemaStore::new();
     store.register(schema).unwrap();
-    let reader_schema = AvroSchema::new(READER_SCHEMA.to_owned());
+    let reader_schema = AvroSchema::new(reader_schema_json.to_owned());
     ReaderBuilder::new()
         .with_writer_schema_store(store)
         .with_batch_size(batch_size)
@@ -215,19 +261,24 @@ fn bench_with_decoder<F>(
 fn criterion_benches(c: &mut Criterion) {
     let data = gen_avro_data_with(INT_SCHEMA, NUM_ROWS, gen_int);
     bench_with_decoder(c, "skip_int", &data, NUM_ROWS, || {
-        new_decoder(INT_SCHEMA, BATCH_SIZE)
+        new_decoder(INT_SCHEMA, SKIP_READER_SCHEMA, BATCH_SIZE)
     });
     let data = gen_avro_data_with(LONG_SCHEMA, NUM_ROWS, gen_long);
     bench_with_decoder(c, "skip_long", &data, NUM_ROWS, || {
-        new_decoder(LONG_SCHEMA, BATCH_SIZE)
+        new_decoder(LONG_SCHEMA, SKIP_READER_SCHEMA, BATCH_SIZE)
     });
     let data = gen_avro_data_with(FLOAT_SCHEMA, NUM_ROWS, gen_float);
     bench_with_decoder(c, "skip_float", &data, NUM_ROWS, || {
-        new_decoder(FLOAT_SCHEMA, BATCH_SIZE)
+        new_decoder(FLOAT_SCHEMA, SKIP_READER_SCHEMA, BATCH_SIZE)
     });
     let data = gen_avro_data_with(DOUBLE_SCHEMA, NUM_ROWS, gen_double);
     bench_with_decoder(c, "skip_double", &data, NUM_ROWS, || {
-        new_decoder(DOUBLE_SCHEMA, BATCH_SIZE)
+        new_decoder(DOUBLE_SCHEMA, SKIP_READER_SCHEMA, BATCH_SIZE)
+    });
+
+    let data = gen_avro_data_with(MIX_SCHEMA, NUM_ROWS, gen_mixed);
+    bench_with_decoder(c, "project_primitives", &data, NUM_ROWS, || {
+        new_decoder(MIX_SCHEMA, PROJECT_READER_SCHEMA, BATCH_SIZE)
     });
 }
 

arrow-avro/src/codec.rs

@@ -84,14 +84,20 @@ pub(crate) enum AvroLiteral {
 /// Contains the necessary information to resolve a writer's record against a reader's record schema.
 #[derive(Debug, Clone, PartialEq)]
 pub(crate) struct ResolvedRecord {
-    /// Maps a writer's field index to the corresponding reader's field index.
-    /// `None` if the writer's field is not present in the reader's schema.
-    pub(crate) writer_to_reader: Arc<[Option<usize>]>,
+    /// Maps a writer's field index to the field's resolution against the reader's schema.
+    pub(crate) writer_fields: Arc<[ResolvedField]>,
     /// A list of indices in the reader's schema for fields that have a default value.
     pub(crate) default_fields: Arc<[usize]>,
+}
+
+/// Resolution information for record fields in the writer schema.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) enum ResolvedField {
+    /// Resolves to a field indexed in the reader schema.
+    ToReader(usize),
     /// For fields present in the writer's schema but not the reader's, this stores their data type.
     /// This is needed to correctly skip over these fields during deserialization.
-    pub(crate) skip_fields: Arc<[Option<AvroDataType>]>,
+    Skip(AvroDataType),
 }
 
 /// Defines the type of promotion to be applied during schema resolution.
@@ -2281,24 +2287,27 @@ impl<'a> Maker<'a> {
                 data_type: dt,
             });
         }
-        // Build skip_fields in writer order; pre-size and push.
-        let mut skip_fields: Vec<Option<AvroDataType>> =
-            Vec::with_capacity(writer_record.fields.len());
-        for (writer_index, writer_field) in writer_record.fields.iter().enumerate() {
-            if writer_to_reader[writer_index].is_some() {
-                skip_fields.push(None);
-            } else {
-                skip_fields.push(Some(self.parse_type(&writer_field.r#type, writer_ns)?));
-            }
-        }
+        // Build writer field map.
+        let writer_fields = writer_record
+            .fields
+            .iter()
+            .enumerate()
+            .map(|(writer_index, writer_field)| {
+                if let Some(reader_index) = writer_to_reader[writer_index] {
+                    Ok(ResolvedField::ToReader(reader_index))
+                } else {
+                    let dt = self.parse_type(&writer_field.r#type, writer_ns)?;
+                    Ok(ResolvedField::Skip(dt))
+                }
+            })
+            .collect::<Result<_, ArrowError>>()?;
         let resolved = AvroDataType::new_with_resolution(
             Codec::Struct(Arc::from(reader_fields)),
             reader_md,
             None,
             Some(ResolutionInfo::Record(ResolvedRecord {
-                writer_to_reader: Arc::from(writer_to_reader),
+                writer_fields,
                 default_fields: Arc::from(default_fields),
-                skip_fields: Arc::from(skip_fields),
             })),
         );
         // Register a resolved record by reader name+namespace for potential named type refs.
@@ -2792,16 +2801,13 @@ mod tests {
         };
         match resolution {
             ResolutionInfo::Record(ResolvedRecord {
-                writer_to_reader,
+                writer_fields,
                 default_fields,
-                skip_fields,
             }) => {
-                assert_eq!(writer_to_reader.len(), 1);
-                assert_eq!(writer_to_reader[0], Some(0));
+                assert_eq!(writer_fields.len(), 1);
+                assert_eq!(writer_fields[0], ResolvedField::ToReader(0));
                 assert_eq!(default_fields.len(), 1);
                 assert_eq!(default_fields[0], 1);
-                assert_eq!(skip_fields.len(), 1);
-                assert_eq!(skip_fields[0], None);
             }
             other => panic!("unexpected resolution {other:?}"),
         }
@@ -2888,16 +2894,13 @@ mod tests {
         };
         match resolution {
             ResolutionInfo::Record(ResolvedRecord {
-                writer_to_reader,
+                writer_fields,
                 default_fields,
-                skip_fields,
             }) => {
-                assert_eq!(writer_to_reader.len(), 1);
-                assert_eq!(writer_to_reader[0], Some(0));
+                assert_eq!(writer_fields.len(), 1);
+                assert_eq!(writer_fields[0], ResolvedField::ToReader(0));
                 assert_eq!(default_fields.len(), 1);
                 assert_eq!(default_fields[0], 1);
-                assert_eq!(skip_fields.len(), 1);
-                assert_eq!(skip_fields[0], None);
             }
             other => panic!("unexpected resolution {other:?}"),
         }
@@ -3714,11 +3717,18 @@ mod tests {
             Some(ResolutionInfo::Record(ref r)) => r.clone(),
             other => panic!("expected record resolution, got {other:?}"),
         };
-        assert_eq!(rec.writer_to_reader.as_ref(), &[Some(1), None, Some(0)]);
+        assert!(matches!(
+            &rec.writer_fields[..],
+            &[
+                ResolvedField::ToReader(1),
+                ResolvedField::Skip(_),
+                ResolvedField::ToReader(0),
+            ]
+        ));
         assert_eq!(rec.default_fields.as_ref(), &[2usize, 3usize]);
-        assert!(rec.skip_fields[0].is_none());
-        assert!(rec.skip_fields[2].is_none());
-        let skip1 = rec.skip_fields[1].as_ref().expect("skip field present");
+        let ResolvedField::Skip(skip1) = &rec.writer_fields[1] else {
+            panic!("should skip field 1")
+        };
         assert!(matches!(skip1.codec(), Codec::Utf8));
         let name_md = &fields[2].data_type().metadata;
         assert_eq!(