Add row-by-row encoding support and tests to arrow-avro

- Introduced `RecordEncoder::encode_rows` to buffer encoded rows as contiguous slices with per-row offsets using `BytesMut`.
- Added `Encoder` for row-by-row Avro encoding, including zero-copy `Bytes` row access via `EncodedRows`.
- Integrated `bytes` crate for efficient encoding operations.
- Updated writer API to offer `build_encoder` for stream formats (e.g., SOE) alongside row-capacity configuration support.
- Adjusted docs to highlight new encoder capabilities.
- Comprehensive tests added to validate single/multi-column, nullable, prefix-based, and empty batch encoding scenarios.

Author: jecsand838

arrow-avro/Cargo.toml

@@ -66,6 +66,7 @@ indexmap = "2.10"
 rand = "0.9"
 md5 = { version = "0.8", optional = true }
 sha2 = { version = "0.10", optional = true }
+bytes = "1.11"
 
 [dev-dependencies]
 arrow-data = { workspace = true }
@@ -76,9 +77,8 @@ rand = { version = "0.9.1", default-features = false, features = [
 ] }
 criterion = { workspace = true, default-features = false }
 tempfile = "3.3"
-arrow = { workspace = true }
+arrow = { workspace = true, features = ["prettyprint"] }
 futures = "0.3.31"
-bytes = "1.10.1"
 async-stream = "0.3.6"
 apache-avro = "0.21.0"
 num-bigint = "0.4"

arrow-avro/src/writer/encoder.rs

@@ -43,6 +43,7 @@ use arrow_buffer::{ArrowNativeType, NullBuffer};
 use arrow_schema::{
     ArrowError, DataType, Field, IntervalUnit, Schema as ArrowSchema, TimeUnit, UnionMode,
 };
+use bytes::{BufMut, BytesMut};
 use std::io::Write;
 use std::sync::Arc;
 use uuid::Uuid;
@@ -826,6 +827,69 @@ impl RecordEncoder {
         }
         Ok(())
     }
+
+    /// Encode rows into a single contiguous `BytesMut` and append row-end offsets.
+    ///
+    /// # Invariants
+    ///
+    /// * `offsets` must be non-empty and seeded with `0` at index 0.
+    /// * `offsets.last()` must equal `out.len()` on entry.
+    /// * On success, exactly `batch.num_rows()` additional offsets are pushed, and
+    ///   `offsets.last()` equals the new `out.len()`.
+    pub(crate) fn encode_rows(
+        &self,
+        batch: &RecordBatch,
+        row_capacity: usize,
+        out: &mut BytesMut,
+        offsets: &mut Vec<u64>,
+    ) -> Result<(), ArrowError> {
+        // Validate invariants once per call (cheap vs. per-row allocations).
+        if offsets.is_empty() {
+            return Err(ArrowError::InvalidArgumentError(
+                "encode_rows requires offsets to be seeded with a 0 sentinel".to_string(),
+            ));
+        }
+        if offsets[0] != 0 {
+            return Err(ArrowError::InvalidArgumentError(
+                "encode_rows requires offsets[0] == 0".to_string(),
+            ));
+        }
+        let expected_last = out.len() as u64;
+        if *offsets.last().unwrap() != expected_last {
+            return Err(ArrowError::InvalidArgumentError(format!(
+                "encode_rows requires offsets.last() == out.len() ({} != {})",
+                offsets.last().unwrap(),
+                expected_last
+            )));
+        }
+        let mut column_encoders = self.prepare_for_batch(batch)?;
+        let n = batch.num_rows();
+        offsets.reserve(n);
+        out.reserve(n.saturating_mul(row_capacity));
+        let mut w = out.writer();
+        match &self.prefix {
+            Some(prefix) => {
+                let prefix_bytes = prefix.as_slice();
+                for row in 0..n {
+                    w.write_all(prefix_bytes)
+                        .map_err(|e| ArrowError::IoError(format!("write prefix: {e}"), e))?;
+                    for enc in column_encoders.iter_mut() {
+                        enc.encode(&mut w, row)?;
+                    }
+                    offsets.push((*w.get_ref()).len() as u64);
+                }
+            }
+            None => {
+                for row in 0..n {
+                    for enc in column_encoders.iter_mut() {
+                        enc.encode(&mut w, row)?;
+                    }
+                    offsets.push((*w.get_ref()).len() as u64);
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 fn find_struct_child_index(fields: &arrow_schema::Fields, name: &str) -> Option<usize> {
@@ -1977,6 +2041,12 @@ mod tests {
         }
     }
 
+    fn row_slice<'a>(buf: &'a [u8], offsets: &[u64], row: usize) -> &'a [u8] {
+        let start = offsets[row] as usize;
+        let end = offsets[row + 1] as usize;
+        &buf[start..end]
+    }
+
     #[test]
     fn binary_encoder() {
         let values: Vec<&[u8]> = vec![b"", b"ab", b"\x00\xFF"];
@@ -3045,4 +3115,222 @@ mod tests {
             other => panic!("expected NullableNoNulls, got {other:?}"),
         }
     }
+
+    #[test]
+    fn encode_rows_single_column_int32() {
+        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
+        let arr = Int32Array::from(vec![1, 2, 3]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![FieldBinding {
+                arrow_index: 0,
+                nullability: None,
+                plan: FieldPlan::Scalar,
+            }],
+            prefix: None,
+        };
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 16, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 4);
+        assert_eq!(*offsets.last().unwrap(), out.len() as u64);
+        assert_bytes_eq(row_slice(&out, &offsets, 0), &avro_long_bytes(1));
+        assert_bytes_eq(row_slice(&out, &offsets, 1), &avro_long_bytes(2));
+        assert_bytes_eq(row_slice(&out, &offsets, 2), &avro_long_bytes(3));
+    }
+
+    #[test]
+    fn encode_rows_multiple_columns() {
+        let schema = ArrowSchema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Utf8, false),
+        ]);
+        let int_arr = Int32Array::from(vec![10, 20]);
+        let str_arr = StringArray::from(vec!["hello", "world"]);
+        let batch = RecordBatch::try_new(
+            Arc::new(schema.clone()),
+            vec![Arc::new(int_arr), Arc::new(str_arr)],
+        )
+        .unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![
+                FieldBinding {
+                    arrow_index: 0,
+                    nullability: None,
+                    plan: FieldPlan::Scalar,
+                },
+                FieldBinding {
+                    arrow_index: 1,
+                    nullability: None,
+                    plan: FieldPlan::Scalar,
+                },
+            ],
+            prefix: None,
+        };
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 32, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 3);
+        assert_eq!(*offsets.last().unwrap(), out.len() as u64);
+        let mut expected_row0 = Vec::new();
+        expected_row0.extend(avro_long_bytes(10));
+        expected_row0.extend(avro_len_prefixed_bytes(b"hello"));
+        assert_bytes_eq(row_slice(&out, &offsets, 0), &expected_row0);
+        let mut expected_row1 = Vec::new();
+        expected_row1.extend(avro_long_bytes(20));
+        expected_row1.extend(avro_len_prefixed_bytes(b"world"));
+        assert_bytes_eq(row_slice(&out, &offsets, 1), &expected_row1);
+    }
+
+    #[test]
+    fn encode_rows_with_prefix() {
+        use crate::codec::AvroFieldBuilder;
+        use crate::schema::AvroSchema;
+        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
+        let arr = Int32Array::from(vec![42]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
+        let avro_schema = AvroSchema::try_from(&schema).unwrap();
+        let fingerprint = avro_schema
+            .fingerprint(crate::schema::FingerprintAlgorithm::Rabin)
+            .unwrap();
+        let avro_root = AvroFieldBuilder::new(&avro_schema.schema().unwrap())
+            .build()
+            .unwrap();
+        let encoder = RecordEncoderBuilder::new(&avro_root, &schema)
+            .with_fingerprint(Some(fingerprint))
+            .build()
+            .unwrap();
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 32, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 2);
+        let row0 = row_slice(&out, &offsets, 0);
+        assert!(row0.len() > 10, "Row should contain prefix + encoded value");
+        assert_eq!(row0[0], 0xC3);
+        assert_eq!(row0[1], 0x01);
+    }
+
+    #[test]
+    fn encode_rows_empty_batch() {
+        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
+        let arr = Int32Array::from(Vec::<i32>::new());
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![FieldBinding {
+                arrow_index: 0,
+                nullability: None,
+                plan: FieldPlan::Scalar,
+            }],
+            prefix: None,
+        };
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 16, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets, vec![0u64]);
+        assert!(out.is_empty());
+    }
+
+    #[test]
+    fn encode_rows_matches_encode_output() {
+        let schema = ArrowSchema::new(vec![
+            Field::new("a", DataType::Int64, false),
+            Field::new("b", DataType::Float64, false),
+        ]);
+        let int_arr = Int64Array::from(vec![100i64, 200, 300]);
+        let float_arr = Float64Array::from(vec![1.5, 2.5, 3.5]);
+        let batch = RecordBatch::try_new(
+            Arc::new(schema.clone()),
+            vec![Arc::new(int_arr), Arc::new(float_arr)],
+        )
+        .unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![
+                FieldBinding {
+                    arrow_index: 0,
+                    nullability: None,
+                    plan: FieldPlan::Scalar,
+                },
+                FieldBinding {
+                    arrow_index: 1,
+                    nullability: None,
+                    plan: FieldPlan::Scalar,
+                },
+            ],
+            prefix: None,
+        };
+        let mut stream_buf = Vec::new();
+        encoder.encode(&mut stream_buf, &batch).unwrap();
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 32, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 1 + batch.num_rows());
+        assert_bytes_eq(&out[..], &stream_buf);
+    }
+
+    #[test]
+    fn encode_rows_appends_to_existing_buffer() {
+        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
+        let arr = Int32Array::from(vec![5, 6]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![FieldBinding {
+                arrow_index: 0,
+                nullability: None,
+                plan: FieldPlan::Scalar,
+            }],
+            prefix: None,
+        };
+        let mut out = BytesMut::new();
+        out.extend_from_slice(&[0xAA, 0xBB]);
+        let mut offsets: Vec<u64> = vec![0, out.len() as u64];
+        encoder
+            .encode_rows(&batch, 16, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 4);
+        assert_eq!(*offsets.last().unwrap(), out.len() as u64);
+        assert_bytes_eq(row_slice(&out, &offsets, 0), &[0xAA, 0xBB]);
+        assert_bytes_eq(row_slice(&out, &offsets, 1), &avro_long_bytes(5));
+        assert_bytes_eq(row_slice(&out, &offsets, 2), &avro_long_bytes(6));
+    }
+
+    #[test]
+    fn encode_rows_nullable_column() {
+        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, true)]);
+        let arr = Int32Array::from(vec![Some(1), None, Some(3)]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
+        let encoder = RecordEncoder {
+            columns: vec![FieldBinding {
+                arrow_index: 0,
+                nullability: Some(Nullability::NullFirst),
+                plan: FieldPlan::Scalar,
+            }],
+            prefix: None,
+        };
+        let mut out = BytesMut::new();
+        let mut offsets: Vec<u64> = vec![0];
+        encoder
+            .encode_rows(&batch, 16, &mut out, &mut offsets)
+            .unwrap();
+        assert_eq!(offsets.len(), 4);
+        let mut expected_row0 = Vec::new();
+        expected_row0.extend(avro_long_bytes(1)); // union branch for value
+        expected_row0.extend(avro_long_bytes(1)); // value
+        assert_bytes_eq(row_slice(&out, &offsets, 0), &expected_row0);
+        let expected_row1 = avro_long_bytes(0); // union branch for null
+        assert_bytes_eq(row_slice(&out, &offsets, 1), &expected_row1);
+        let mut expected_row2 = Vec::new();
+        expected_row2.extend(avro_long_bytes(1)); // union branch for value
+        expected_row2.extend(avro_long_bytes(3)); // value
+        assert_bytes_eq(row_slice(&out, &offsets, 2), &expected_row2);
+    }
 }

arrow-avro/src/writer/format.rs

@@ -135,6 +135,40 @@ impl AvroFormat for AvroSoeFormat {
     }
 }
 
+/// Unframed Avro binary streaming format ("raw Avro record body bytes (no prefix, no OCF header)").
+///
+/// Each record written by the stream writer contains only the raw Avro
+/// record body bytes (i.e., the Avro binary encoding of the datum) with **no**
+/// per-record prefix and **no** Object Container File (OCF) header.
+///
+/// This format is useful when another transport provides framing (for example,
+/// length-delimited buffers) or when embedding Avro record payloads inside a
+/// larger envelope.
+#[derive(Debug, Default)]
+pub struct AvroBinaryFormat;
+
+impl AvroFormat for AvroBinaryFormat {
+    const NEEDS_PREFIX: bool = false;
+
+    fn start_stream<W: Write>(
+        &mut self,
+        _writer: &mut W,
+        _schema: &Schema,
+        compression: Option<CompressionCodec>,
+    ) -> Result<(), ArrowError> {
+        if compression.is_some() {
+            return Err(ArrowError::InvalidArgumentError(
+                "Compression not supported for Avro binary streaming".to_string(),
+            ));
+        }
+        Ok(())
+    }
+
+    fn sync_marker(&self) -> Option<&[u8; 16]> {
+        None
+    }
+}
+
 #[inline]
 fn write_string<W: Write>(writer: &mut W, s: &str) -> Result<(), ArrowError> {
     write_bytes(writer, s.as_bytes())