You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2025/09/08 16:41:29 UTC
(arrow-rs) branch main updated: Added List and Struct Encoding to arrow-avro Writer (#8274)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 911940f5a4 Added List and Struct Encoding to arrow-avro Writer (#8274)
911940f5a4 is described below
commit 911940f5a4f09957e893d79b4daf8a8b0a5618c4
Author: Connor Sanders <co...@elastiflow.com>
AuthorDate: Mon Sep 8 11:41:23 2025 -0500
Added List and Struct Encoding to arrow-avro Writer (#8274)
# Which issue does this PR close?
- Part of https://github.com/apache/arrow-rs/issues/4886
# Rationale for this change
This refactor streamlines the `arrow-avro` writer by introducing a
single, schema‑driven `RecordEncoder` that plans writes up front and
encodes rows using consistent, explicit rules for nullability and type
dispatch. It reduces duplication in nested/struct/list handling, makes
the order of Avro union branches (null‑first vs null‑second) an explicit
choice, and aligns header schema generation with value encoding.
This should improve correctness (especially for nested optionals), make
behavior easier to reason about, and pave the way for future
optimizations.
# What changes are included in this PR?
**High‑level:**
* Introduces a unified, schema‑driven `RecordEncoder` with a builder
that walks the Avro record in Avro order and maps each field to its
Arrow column, producing a reusable write plan. The encoder covers
scalars and nested types (struct, (large) lists, maps,
strings/binaries).
* Applies a single model of **nullability** throughout encoding,
including nested sites (list items, fixed‑size list items, map values),
and uses explicit union‑branch indices according to the chosen order.
**API and implementation details:**
* **Writer / encoder refactor**
* Replaces the previous per‑column/child encoding paths with a
**`FieldPlan`** tree (variants for `Scalar`, `Struct { … }`, and `List {
… }`) and per‑site `nullability` carried from the Avro schema.
* Adds encoder variants for `LargeBinary`, `Utf8`, `Utf8Large`, `List`,
`LargeList`, and `Struct`.
* Encodes union branch indices with `write_optional_index` (writes
`0x00/0x02` according to Null‑First/Null‑Second), replacing the old
branch write.
* **Schema generation & metadata**
* Moves the **`Nullability`** enum to `schema.rs` and threads it through
schema generation and writer logic.
* Adds `AvroSchema::from_arrow_with_options(schema,
Option<Nullability>)` to either reuse embedded Avro JSON or build new
Avro JSON that **honors the requested null‑union order at all nullable
sites**.
* Adds `extend_with_passthrough_metadata` so Arrow schema metadata is
copied into Avro JSON while skipping Avro‑reserved and internal Arrow
keys.
* Introduces helpers like `wrap_nullable` and
`arrow_field_to_avro_with_order` to apply ordering consistently for
arrays, fixed‑size lists, maps, structs, and unions.
* **Format and glue**
* Simplifies `writer/format.rs` by removing the `EncoderOptions`
plumbing from the OCF format; `write_long` remains exported for header
writing.
# Are these changes tested?
Yes.
* Adds focused unit tests in `writer/encoder.rs` that verify scalar and
string/binary encodings (e.g., Binary/LargeBinary, Utf8/LargeUtf8) and
validate length/branch encoding primitives used by the writer.
* Round trip integration tests that validate List and Struct decoding in
`writer/mod.rs`.
* Adjusts existing schema tests (e.g., decimal metadata expectations) to
align with the new schema/metadata handling.
# Are there any user-facing changes?
N/A because arrow-avro is not public yet.
---------
Co-authored-by: Ryan Johnson <sc...@users.noreply.github.com>
Co-authored-by: Matthijs Brobbel <m1...@gmail.com>
---
arrow-avro/src/codec.rs | 17 +-
arrow-avro/src/reader/record.rs | 2 +-
arrow-avro/src/schema.rs | 244 ++++++++----
arrow-avro/src/writer/encoder.rs | 777 +++++++++++++++++++++++++++++++--------
arrow-avro/src/writer/format.rs | 22 +-
arrow-avro/src/writer/mod.rs | 177 ++++++---
6 files changed, 928 insertions(+), 311 deletions(-)
diff --git a/arrow-avro/src/codec.rs b/arrow-avro/src/codec.rs
index d19e9b8ccc..8b103ff3b2 100644
--- a/arrow-avro/src/codec.rs
+++ b/arrow-avro/src/codec.rs
@@ -16,8 +16,8 @@
// under the License.
use crate::schema::{
- Attributes, AvroSchema, ComplexType, Enum, PrimitiveType, Record, Schema, Type, TypeName,
- AVRO_ENUM_SYMBOLS_METADATA_KEY,
+ Attributes, AvroSchema, ComplexType, Enum, Nullability, PrimitiveType, Record, Schema, Type,
+ TypeName, AVRO_ENUM_SYMBOLS_METADATA_KEY,
};
use arrow_schema::{
ArrowError, DataType, Field, Fields, IntervalUnit, TimeUnit, DECIMAL128_MAX_PRECISION,
@@ -28,19 +28,6 @@ use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
-/// Avro types are not nullable, with nullability instead encoded as a union
-/// where one of the variants is the null type.
-///
-/// To accommodate this we special case two-variant unions where one of the
-/// variants is the null type, and use this to derive arrow's notion of nullability
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub enum Nullability {
- /// The nulls are encoded as the first union variant
- NullFirst,
- /// The nulls are encoded as the second union variant
- NullSecond,
-}
-
/// Contains information about how to resolve differences between a writer's and a reader's schema.
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum ResolutionInfo {
diff --git a/arrow-avro/src/reader/record.rs b/arrow-avro/src/reader/record.rs
index 6e5756ef41..f443dc0dfe 100644
--- a/arrow-avro/src/reader/record.rs
+++ b/arrow-avro/src/reader/record.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::codec::{AvroDataType, Codec, Nullability, Promotion, ResolutionInfo};
+use crate::codec::{AvroDataType, Codec, Promotion, ResolutionInfo};
use crate::reader::block::{Block, BlockDecoder};
use crate::reader::cursor::AvroCursor;
use crate::reader::header::Header;
diff --git a/arrow-avro/src/schema.rs b/arrow-avro/src/schema.rs
index 46ac30b495..6e343736c1 100644
--- a/arrow-avro/src/schema.rs
+++ b/arrow-avro/src/schema.rs
@@ -59,6 +59,20 @@ pub fn compare_schemas(writer: &Schema, reader: &Schema) -> Result<bool, ArrowEr
Ok(canon_writer == canon_reader)
}
+/// Avro types are not nullable, with nullability instead encoded as a union
+/// where one of the variants is the null type.
+///
+/// To accommodate this, we specially case two-variant unions where one of the
+/// variants is the null type, and use this to derive arrow's notion of nullability
+#[derive(Debug, Copy, Clone, PartialEq, Default)]
+pub enum Nullability {
+ /// The nulls are encoded as the first union variant
+ #[default]
+ NullFirst,
+ /// The nulls are encoded as the second union variant
+ NullSecond,
+}
+
/// Either a [`PrimitiveType`] or a reference to a previously defined named type
///
/// <https://avro.apache.org/docs/1.11.1/specification/#names>
@@ -113,7 +127,7 @@ pub struct Attributes<'a> {
/// Additional JSON attributes
#[serde(flatten)]
- pub additional: HashMap<&'a str, serde_json::Value>,
+ pub additional: HashMap<&'a str, Value>,
}
impl Attributes<'_> {
@@ -309,51 +323,11 @@ pub struct AvroSchema {
impl TryFrom<&ArrowSchema> for AvroSchema {
type Error = ArrowError;
+ /// Converts an `ArrowSchema` to `AvroSchema`, delegating to
+ /// `AvroSchema::from_arrow_with_options` with `None` so that the
+ /// union null ordering is decided by `Nullability::default()`.
fn try_from(schema: &ArrowSchema) -> Result<Self, Self::Error> {
- // Fast‑path: schema already contains Avro JSON
- if let Some(json) = schema.metadata.get(SCHEMA_METADATA_KEY) {
- return Ok(AvroSchema::new(json.clone()));
- }
- let mut name_gen = NameGenerator::default();
- let fields_json = schema
- .fields()
- .iter()
- .map(|f| arrow_field_to_avro(f, &mut name_gen))
- .collect::<Result<Vec<_>, _>>()?;
- // Assemble top‑level record
- let record_name = schema
- .metadata
- .get(AVRO_NAME_METADATA_KEY)
- .map_or("topLevelRecord", |s| s.as_str());
- let mut record = JsonMap::with_capacity(schema.metadata.len() + 4);
- record.insert("type".into(), Value::String("record".into()));
- record.insert(
- "name".into(),
- Value::String(sanitise_avro_name(record_name)),
- );
- if let Some(ns) = schema.metadata.get(AVRO_NAMESPACE_METADATA_KEY) {
- record.insert("namespace".into(), Value::String(ns.clone()));
- }
- if let Some(doc) = schema.metadata.get(AVRO_DOC_METADATA_KEY) {
- record.insert("doc".into(), Value::String(doc.clone()));
- }
- record.insert("fields".into(), Value::Array(fields_json));
- let schema_prefix = format!("{SCHEMA_METADATA_KEY}.");
- for (meta_key, meta_val) in &schema.metadata {
- // Skip keys already handled or internal
- if meta_key.starts_with("avro.")
- || meta_key.starts_with(schema_prefix.as_str())
- || is_internal_arrow_key(meta_key)
- {
- continue;
- }
- let json_val =
- serde_json::from_str(meta_val).unwrap_or_else(|_| Value::String(meta_val.clone()));
- record.insert(meta_key.clone(), json_val);
- }
- let json_string = serde_json::to_string(&Value::Object(record))
- .map_err(|e| ArrowError::SchemaError(format!("Serialising Avro JSON failed: {e}")))?;
- Ok(AvroSchema::new(json_string))
+ AvroSchema::from_arrow_with_options(schema, None)
}
}
@@ -453,6 +427,49 @@ impl AvroSchema {
pub fn generate_canonical_form(schema: &Schema) -> Result<String, ArrowError> {
build_canonical(schema, None)
}
+
+ /// Build Avro JSON from an Arrow [`ArrowSchema`], applying the given null‑union order.
+ ///
+ /// If the input Arrow schema already contains Avro JSON in
+ /// [`SCHEMA_METADATA_KEY`], that JSON is returned verbatim to preserve
+ /// the exact header encoding alignment; otherwise, a new JSON is generated
+ /// honoring `null_union_order` at **all nullable sites**.
+ pub fn from_arrow_with_options(
+ schema: &ArrowSchema,
+ null_order: Option<Nullability>,
+ ) -> Result<AvroSchema, ArrowError> {
+ if let Some(json) = schema.metadata.get(SCHEMA_METADATA_KEY) {
+ return Ok(AvroSchema::new(json.clone()));
+ }
+ let order = null_order.unwrap_or_default();
+ let mut name_gen = NameGenerator::default();
+ let fields_json = schema
+ .fields()
+ .iter()
+ .map(|f| arrow_field_to_avro(f, &mut name_gen, order))
+ .collect::<Result<Vec<_>, _>>()?;
+ let record_name = schema
+ .metadata
+ .get(AVRO_NAME_METADATA_KEY)
+ .map_or("topLevelRecord", |s| s.as_str());
+ let mut record = JsonMap::with_capacity(schema.metadata.len() + 4);
+ record.insert("type".into(), Value::String("record".into()));
+ record.insert(
+ "name".into(),
+ Value::String(sanitise_avro_name(record_name)),
+ );
+ if let Some(ns) = schema.metadata.get(AVRO_NAMESPACE_METADATA_KEY) {
+ record.insert("namespace".into(), Value::String(ns.clone()));
+ }
+ if let Some(doc) = schema.metadata.get(AVRO_DOC_METADATA_KEY) {
+ record.insert("doc".into(), Value::String(doc.clone()));
+ }
+ record.insert("fields".into(), Value::Array(fields_json));
+ extend_with_passthrough_metadata(&mut record, &schema.metadata);
+ let json_string = serde_json::to_string(&Value::Object(record))
+ .map_err(|e| ArrowError::SchemaError(format!("Serializing Avro JSON failed: {e}")))?;
+ Ok(AvroSchema::new(json_string))
+ }
}
/// Supported fingerprint algorithms for Avro schema identification.
@@ -862,6 +879,24 @@ fn is_internal_arrow_key(key: &str) -> bool {
key.starts_with("ARROW:") || key == SCHEMA_METADATA_KEY
}
+/// Copies Arrow schema metadata entries to the provided JSON map,
+/// skipping keys that are Avro-reserved, internal Arrow keys, or
+/// nested under the `avro.schema.` namespace. Values that parse as
+/// JSON are inserted as JSON; otherwise the raw string is preserved.
+fn extend_with_passthrough_metadata(
+ target: &mut JsonMap<String, Value>,
+ metadata: &HashMap<String, String>,
+) {
+ for (meta_key, meta_val) in metadata {
+ if meta_key.starts_with("avro.") || is_internal_arrow_key(meta_key) {
+ continue;
+ }
+ let json_val =
+ serde_json::from_str(meta_val).unwrap_or_else(|_| Value::String(meta_val.clone()));
+ target.insert(meta_key.clone(), json_val);
+ }
+}
+
// Sanitize an arbitrary string so it is a valid Avro field or type name
fn sanitise_avro_name(base_name: &str) -> String {
if base_name.is_empty() {
@@ -932,12 +967,21 @@ fn merge_extras(schema: Value, mut extras: JsonMap<String, Value>) -> Value {
}
}
-// Convert an Arrow `DataType` into an Avro schema `Value`.
+fn wrap_nullable(inner: Value, null_order: Nullability) -> Value {
+ let null = Value::String("null".into());
+ let elements = match null_order {
+ Nullability::NullFirst => vec![null, inner],
+ Nullability::NullSecond => vec![inner, null],
+ };
+ Value::Array(elements)
+}
+
fn datatype_to_avro(
dt: &DataType,
field_name: &str,
metadata: &HashMap<String, String>,
name_gen: &mut NameGenerator,
+ null_order: Nullability,
) -> Result<(Value, JsonMap<String, Value>), ArrowError> {
let mut extras = JsonMap::new();
let val = match dt {
@@ -1051,20 +1095,32 @@ fn datatype_to_avro(
if matches!(dt, DataType::LargeList(_)) {
extras.insert("arrowLargeList".into(), Value::Bool(true));
}
- let (items, ie) =
- datatype_to_avro(child.data_type(), child.name(), child.metadata(), name_gen)?;
+ let items_schema = process_datatype(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ null_order,
+ child.is_nullable(),
+ )?;
json!({
"type": "array",
- "items": merge_extras(items, ie)
+ "items": items_schema
})
}
DataType::FixedSizeList(child, len) => {
extras.insert("arrowFixedSize".into(), json!(len));
- let (items, ie) =
- datatype_to_avro(child.data_type(), child.name(), child.metadata(), name_gen)?;
+ let items_schema = process_datatype(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ null_order,
+ child.is_nullable(),
+ )?;
json!({
"type": "array",
- "items": merge_extras(items, ie)
+ "items": items_schema
})
}
DataType::Map(entries, _) => {
@@ -1076,21 +1132,23 @@ fn datatype_to_avro(
))
}
};
- let (val_schema, value_entry) = datatype_to_avro(
+ let values_schema = process_datatype(
value_field.data_type(),
value_field.name(),
value_field.metadata(),
name_gen,
+ null_order,
+ value_field.is_nullable(),
)?;
json!({
"type": "map",
- "values": merge_extras(val_schema, value_entry)
+ "values": values_schema
})
}
DataType::Struct(fields) => {
let avro_fields = fields
.iter()
- .map(|field| arrow_field_to_avro(field, name_gen))
+ .map(|field| arrow_field_to_avro(field, name_gen, null_order))
.collect::<Result<Vec<_>, _>>()?;
json!({
"type": "record",
@@ -1108,19 +1166,24 @@ fn datatype_to_avro(
"symbols": symbols
})
} else {
- let (inner, ie) = datatype_to_avro(value.as_ref(), field_name, metadata, name_gen)?;
- merge_extras(inner, ie)
+ process_datatype(
+ value.as_ref(),
+ field_name,
+ metadata,
+ name_gen,
+ null_order,
+ false,
+ )?
}
}
- DataType::RunEndEncoded(_, values) => {
- let (inner, ie) = datatype_to_avro(
- values.data_type(),
- values.name(),
- values.metadata(),
- name_gen,
- )?;
- merge_extras(inner, ie)
- }
+ DataType::RunEndEncoded(_, values) => process_datatype(
+ values.data_type(),
+ values.name(),
+ values.metadata(),
+ name_gen,
+ null_order,
+ false,
+ )?,
DataType::Union(_, _) => {
return Err(ArrowError::NotYetImplemented(
"Arrow Union to Avro Union not yet supported".into(),
@@ -1135,27 +1198,40 @@ fn datatype_to_avro(
Ok((val, extras))
}
+fn process_datatype(
+ dt: &DataType,
+ field_name: &str,
+ metadata: &HashMap<String, String>,
+ name_gen: &mut NameGenerator,
+ null_order: Nullability,
+ is_nullable: bool,
+) -> Result<Value, ArrowError> {
+ let (schema, extras) = datatype_to_avro(dt, field_name, metadata, name_gen, null_order)?;
+ let mut merged = merge_extras(schema, extras);
+ if is_nullable {
+ merged = wrap_nullable(merged, null_order)
+ }
+ Ok(merged)
+}
+
fn arrow_field_to_avro(
field: &ArrowField,
name_gen: &mut NameGenerator,
+ null_order: Nullability,
) -> Result<Value, ArrowError> {
- // Sanitize field name to ensure Avro validity but store the original in metadata
let avro_name = sanitise_avro_name(field.name());
- let (schema, extras) =
- datatype_to_avro(field.data_type(), &avro_name, field.metadata(), name_gen)?;
- // If nullable, wrap `[ "null", <type> ]`, NOTE: second order nullability to be added in a follow-up
- let mut schema = if field.is_nullable() {
- Value::Array(vec![
- Value::String("null".into()),
- merge_extras(schema, extras),
- ])
- } else {
- merge_extras(schema, extras)
- };
+ let schema_value = process_datatype(
+ field.data_type(),
+ &avro_name,
+ field.metadata(),
+ name_gen,
+ null_order,
+ field.is_nullable(),
+ )?;
// Build the field map
let mut map = JsonMap::with_capacity(field.metadata().len() + 3);
map.insert("name".into(), Value::String(avro_name));
- map.insert("type".into(), schema);
+ map.insert("type".into(), schema_value);
// Transfer selected metadata
for (meta_key, meta_val) in field.metadata() {
if is_internal_arrow_key(meta_key) {
@@ -1728,7 +1804,7 @@ mod tests {
r#type: Schema::Type(Type {
r#type: TypeName::Primitive(PrimitiveType::Bytes),
attributes: Attributes {
- logical_type: Some("decimal"),
+ logical_type: None,
additional: HashMap::from([("precision", json!(4))]),
},
}),
@@ -2056,4 +2132,12 @@ mod tests {
assert_eq!(arrow_field, expected);
}
+
+ #[test]
+ fn default_order_is_consistent() {
+ let arrow_schema = ArrowSchema::new(vec![ArrowField::new("s", DataType::Utf8, true)]);
+ let a = AvroSchema::try_from(&arrow_schema).unwrap().json_string;
+ let b = AvroSchema::from_arrow_with_options(&arrow_schema, None);
+ assert_eq!(a, b.unwrap().json_string);
+ }
}
diff --git a/arrow-avro/src/writer/encoder.rs b/arrow-avro/src/writer/encoder.rs
index c45aa6cfcf..ccf80fd8d1 100644
--- a/arrow-avro/src/writer/encoder.rs
+++ b/arrow-avro/src/writer/encoder.rs
@@ -17,31 +17,25 @@
//! Avro Encoder for Arrow types.
+use crate::codec::{AvroDataType, AvroField, Codec};
+use crate::schema::Nullability;
use arrow_array::cast::AsArray;
use arrow_array::types::{
ArrowPrimitiveType, Float32Type, Float64Type, Int32Type, Int64Type, TimestampMicrosecondType,
};
-use arrow_array::OffsetSizeTrait;
-use arrow_array::{Array, GenericBinaryArray, PrimitiveArray, RecordBatch};
+use arrow_array::{
+ Array, GenericBinaryArray, GenericListArray, GenericStringArray, LargeListArray, ListArray,
+ OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray,
+};
use arrow_buffer::NullBuffer;
-use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit};
+use arrow_schema::{ArrowError, DataType, Field, Schema as ArrowSchema, TimeUnit};
use std::io::Write;
-/// Behavior knobs for the Avro encoder.
-///
-/// When `impala_mode` is `true`, optional/nullable values are encoded
-/// as Avro unions with **null second** (`[T, "null"]`). When `false`
-/// (default), we use **null first** (`["null", T]`).
-#[derive(Debug, Clone, Copy, Default)]
-pub struct EncoderOptions {
- impala_mode: bool, // Will be fully implemented in a follow-up PR
-}
-
/// Encode a single Avro-`long` using ZigZag + variable length, buffered.
///
/// Spec: <https://avro.apache.org/docs/1.11.1/specification/#binary-encoding>
#[inline]
-pub fn write_long<W: Write + ?Sized>(writer: &mut W, value: i64) -> Result<(), ArrowError> {
+pub fn write_long<W: Write + ?Sized>(out: &mut W, value: i64) -> Result<(), ArrowError> {
let mut zz = ((value << 1) ^ (value >> 63)) as u64;
// At most 10 bytes for 64-bit varint
let mut buf = [0u8; 10];
@@ -53,28 +47,25 @@ pub fn write_long<W: Write + ?Sized>(writer: &mut W, value: i64) -> Result<(), A
}
buf[i] = (zz & 0x7F) as u8;
i += 1;
- writer
- .write_all(&buf[..i])
+ out.write_all(&buf[..i])
.map_err(|e| ArrowError::IoError(format!("write long: {e}"), e))
}
#[inline]
-fn write_int<W: Write + ?Sized>(writer: &mut W, value: i32) -> Result<(), ArrowError> {
- write_long(writer, value as i64)
+fn write_int<W: Write + ?Sized>(out: &mut W, value: i32) -> Result<(), ArrowError> {
+ write_long(out, value as i64)
}
#[inline]
-fn write_len_prefixed<W: Write + ?Sized>(writer: &mut W, bytes: &[u8]) -> Result<(), ArrowError> {
- write_long(writer, bytes.len() as i64)?;
- writer
- .write_all(bytes)
+fn write_len_prefixed<W: Write + ?Sized>(out: &mut W, bytes: &[u8]) -> Result<(), ArrowError> {
+ write_long(out, bytes.len() as i64)?;
+ out.write_all(bytes)
.map_err(|e| ArrowError::IoError(format!("write bytes: {e}"), e))
}
#[inline]
-fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowError> {
- writer
- .write_all(&[if v { 1 } else { 0 }])
+fn write_bool<W: Write + ?Sized>(out: &mut W, v: bool) -> Result<(), ArrowError> {
+ out.write_all(&[if v { 1 } else { 0 }])
.map_err(|e| ArrowError::IoError(format!("write bool: {e}"), e))
}
@@ -83,146 +74,385 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// Branch index is 0-based per Avro unions:
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
-#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
- writer: &mut W,
+fn write_optional_index<W: Write + ?Sized>(
+ out: &mut W,
is_null: bool,
- impala_mode: bool,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ let byte = union_value_branch_byte(null_order, is_null);
+ out.write_all(&[byte])
+ .map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Encode a `RecordBatch` in Avro binary format using **default options**.
-pub fn encode_record_batch<W: Write>(batch: &RecordBatch, out: &mut W) -> Result<(), ArrowError> {
- encode_record_batch_with_options(batch, out, &EncoderOptions::default())
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ union_value_byte: u8,
+ },
+ Nullable {
+ nulls: NullBuffer,
+ null_order: Nullability,
+ },
}
-/// Encode a `RecordBatch` with explicit `EncoderOptions`.
-pub fn encode_record_batch_with_options<W: Write>(
- batch: &RecordBatch,
- out: &mut W,
- opts: &EncoderOptions,
-) -> Result<(), ArrowError> {
- let mut encoders = batch
- .schema()
- .fields()
- .iter()
- .zip(batch.columns())
- .map(|(field, array)| Ok((field.is_nullable(), make_encoder(array.as_ref())?)))
- .collect::<Result<Vec<_>, ArrowError>>()?;
- (0..batch.num_rows()).try_for_each(|row| {
- encoders.iter_mut().try_for_each(|(is_nullable, enc)| {
- if *is_nullable {
- let is_null = enc.is_null(row);
- write_optional_branch(out, is_null, opts.impala_mode)?;
- if is_null {
- return Ok(());
+/// Arrow to Avro FieldEncoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Carries the per-site nullability **state** as a single enum that enforces invariants
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ null_state: NullState,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: &FieldPlan,
+ nullability: Option<Nullability>,
+ ) -> Result<Self, ArrowError> {
+ let encoder = match plan {
+ FieldPlan::Struct { encoders } => {
+ let arr = array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| ArrowError::SchemaError("Expected StructArray".into()))?;
+ Encoder::Struct(Box::new(StructEncoder::try_new(arr, encoders)?))
+ }
+ FieldPlan::List {
+ items_nullability,
+ item_plan,
+ } => match array.data_type() {
+ DataType::List(_) => {
+ let arr = array
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .ok_or_else(|| ArrowError::SchemaError("Expected ListArray".into()))?;
+ Encoder::List(Box::new(ListEncoder32::try_new(
+ arr,
+ *items_nullability,
+ item_plan.as_ref(),
+ )?))
+ }
+ DataType::LargeList(_) => {
+ let arr = array
+ .as_any()
+ .downcast_ref::<LargeListArray>()
+ .ok_or_else(|| ArrowError::SchemaError("Expected LargeListArray".into()))?;
+ Encoder::LargeList(Box::new(ListEncoder64::try_new(
+ arr,
+ *items_nullability,
+ item_plan.as_ref(),
+ )?))
+ }
+ other => {
+ return Err(ArrowError::SchemaError(format!(
+ "Avro array site requires Arrow List/LargeList, found: {other:?}"
+ )))
+ }
+ },
+ FieldPlan::Scalar => match array.data_type() {
+ DataType::Boolean => Encoder::Boolean(BooleanEncoder(array.as_boolean())),
+ DataType::Utf8 => {
+ Encoder::Utf8(Utf8GenericEncoder::<i32>(array.as_string::<i32>()))
+ }
+ DataType::LargeUtf8 => {
+ Encoder::Utf8Large(Utf8GenericEncoder::<i64>(array.as_string::<i64>()))
+ }
+ DataType::Int32 => Encoder::Int(IntEncoder(array.as_primitive::<Int32Type>())),
+ DataType::Int64 => Encoder::Long(LongEncoder(array.as_primitive::<Int64Type>())),
+ DataType::Float32 => {
+ Encoder::Float32(F32Encoder(array.as_primitive::<Float32Type>()))
+ }
+ DataType::Float64 => {
+ Encoder::Float64(F64Encoder(array.as_primitive::<Float64Type>()))
+ }
+ DataType::Binary => Encoder::Binary(BinaryEncoder(array.as_binary::<i32>())),
+ DataType::LargeBinary => {
+ Encoder::LargeBinary(BinaryEncoder(array.as_binary::<i64>()))
}
+ DataType::Timestamp(TimeUnit::Microsecond, _) => Encoder::Timestamp(LongEncoder(
+ array.as_primitive::<TimestampMicrosecondType>(),
+ )),
+ other => {
+ return Err(ArrowError::NotYetImplemented(format!(
+ "Avro scalar type not yet supported: {other:?}"
+ )));
+ }
+ },
+ other => {
+ return Err(ArrowError::NotYetImplemented(format!(
+ "Avro writer: {other:?} not yet supported",
+ )));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, array.null_count() > 0) {
+ (None, false) => NullState::NonNullable,
+ (None, true) => {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Avro site '{}' is non-nullable, but array contains nulls",
+ field.name()
+ )));
+ }
+ (Some(order), false) => {
+ // Optimization: drop any bitmap; emit a constant "value" branch byte.
+ NullState::NullableNoNulls {
+ union_value_byte: union_value_branch_byte(order, false),
+ }
+ }
+ (Some(null_order), true) => {
+ let Some(nulls) = array.nulls().cloned() else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ )));
+ };
+ NullState::Nullable { nulls, null_order }
}
- enc.encode(row, out)
+ };
+ Ok(Self {
+ encoder,
+ null_state,
})
- })
-}
+ }
-/// Enum for static dispatch of concrete encoders.
-enum Encoder<'a> {
- Boolean(BooleanEncoder<'a>),
- Int(IntEncoder<'a, Int32Type>),
- Long(LongEncoder<'a, Int64Type>),
- Timestamp(LongEncoder<'a, TimestampMicrosecondType>),
- Float32(F32Encoder<'a>),
- Float64(F64Encoder<'a>),
- Binary(BinaryEncoder<'a, i32>),
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => {}
+ NullState::NullableNoNulls { union_value_byte } => out
+ .write_all(&[*union_value_byte])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))?,
+ NullState::Nullable { nulls, null_order } if nulls.is_null(idx) => {
+ return write_optional_index(out, true, *null_order); // no value to write
+ }
+ NullState::Nullable { null_order, .. } => {
+ write_optional_index(out, false, *null_order)?;
+ }
+ }
+ self.encoder.encode(out, idx)
+ }
}
-impl<'a> Encoder<'a> {
- /// Encode the value at `idx`.
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
- match self {
- Encoder::Boolean(e) => e.encode(idx, out),
- Encoder::Int(e) => e.encode(idx, out),
- Encoder::Long(e) => e.encode(idx, out),
- Encoder::Timestamp(e) => e.encode(idx, out),
- Encoder::Float32(e) => e.encode(idx, out),
- Encoder::Float64(e) => e.encode(idx, out),
- Encoder::Binary(e) => e.encode(idx, out),
- }
+fn union_value_branch_byte(null_order: Nullability, is_null: bool) -> u8 {
+ let nulls_first = null_order == Nullability::default();
+ if nulls_first == is_null {
+ 0x00
+ } else {
+ 0x02
}
}
-/// An encoder + a null buffer for nullable fields.
-pub struct NullableEncoder<'a> {
- encoder: Encoder<'a>,
- nulls: Option<NullBuffer>,
+/// Per‑site encoder plan for a field. This mirrors the Avro structure, so nested
+/// optional branch order can be honored exactly as declared by the schema.
+#[derive(Debug, Clone)]
+enum FieldPlan {
+ /// Non-nested scalar/logical type
+ Scalar,
+ /// Record/Struct with Avro‑ordered children
+ Struct { encoders: Vec<FieldBinding> },
+ /// Array with item‑site nullability and nested plan
+ List {
+ items_nullability: Option<Nullability>,
+ item_plan: Box<FieldPlan>,
+ },
}
-impl<'a> NullableEncoder<'a> {
- /// Create a new nullable encoder, wrapping a non-null encoder and a null buffer.
- #[inline]
- fn new(encoder: Encoder<'a>, nulls: Option<NullBuffer>) -> Self {
- Self { encoder, nulls }
- }
+#[derive(Debug, Clone)]
+struct FieldBinding {
+ /// Index of the Arrow field/column associated with this Avro field site
+ arrow_index: usize,
+ /// Nullability/order for this site (None for required fields)
+ nullability: Option<Nullability>,
+ /// Nested plan for this site
+ plan: FieldPlan,
+}
- /// Encode the value at `idx`, assuming it's not-null.
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
- self.encoder.encode(idx, out)
+/// Builder for `RecordEncoder` write plan
+#[derive(Debug)]
+pub struct RecordEncoderBuilder<'a> {
+ avro_root: &'a AvroField,
+ arrow_schema: &'a ArrowSchema,
+}
+
+impl<'a> RecordEncoderBuilder<'a> {
+ /// Create a new builder from the Avro root and Arrow schema.
+ pub fn new(avro_root: &'a AvroField, arrow_schema: &'a ArrowSchema) -> Self {
+ Self {
+ avro_root,
+ arrow_schema,
+ }
}
- /// Check if the value at `idx` is null.
- #[inline]
- fn is_null(&self, idx: usize) -> bool {
- self.nulls.as_ref().is_some_and(|nulls| nulls.is_null(idx))
+ /// Build the `RecordEncoder` by walking the Avro **record** root in Avro order,
+ /// resolving each field to an Arrow index by name.
+ pub fn build(self) -> Result<RecordEncoder, ArrowError> {
+ let avro_root_dt = self.avro_root.data_type();
+ let Codec::Struct(root_fields) = avro_root_dt.codec() else {
+ return Err(ArrowError::SchemaError(
+ "Top-level Avro schema must be a record/struct".into(),
+ ));
+ };
+ let mut columns = Vec::with_capacity(root_fields.len());
+ for root_field in root_fields.as_ref() {
+ let name = root_field.name();
+ let arrow_index = self.arrow_schema.index_of(name).map_err(|e| {
+ ArrowError::SchemaError(format!("Schema mismatch for field '{name}': {e}"))
+ })?;
+ columns.push(FieldBinding {
+ arrow_index,
+ nullability: root_field.data_type().nullability(),
+ plan: FieldPlan::build(
+ root_field.data_type(),
+ self.arrow_schema.field(arrow_index),
+ )?,
+ });
+ }
+ Ok(RecordEncoder { columns })
}
}
-/// Creates an Avro encoder for the given `array`.
-pub fn make_encoder<'a>(array: &'a dyn Array) -> Result<NullableEncoder<'a>, ArrowError> {
- let nulls = array.nulls().cloned();
- let enc = match array.data_type() {
- DataType::Boolean => {
- let arr = array.as_boolean();
- NullableEncoder::new(Encoder::Boolean(BooleanEncoder(arr)), nulls)
- }
- DataType::Int32 => {
- let arr = array.as_primitive::<Int32Type>();
- NullableEncoder::new(Encoder::Int(IntEncoder(arr)), nulls)
- }
- DataType::Int64 => {
- let arr = array.as_primitive::<Int64Type>();
- NullableEncoder::new(Encoder::Long(LongEncoder(arr)), nulls)
- }
- DataType::Float32 => {
- let arr = array.as_primitive::<Float32Type>();
- NullableEncoder::new(Encoder::Float32(F32Encoder(arr)), nulls)
- }
- DataType::Float64 => {
- let arr = array.as_primitive::<Float64Type>();
- NullableEncoder::new(Encoder::Float64(F64Encoder(arr)), nulls)
+/// A pre-computed plan for encoding a `RecordBatch` to Avro.
+///
+/// Derived from an Avro schema and an Arrow schema. It maps
+/// top-level Avro fields to Arrow columns and contains a nested encoding plan
+/// for each column.
+#[derive(Debug, Clone)]
+pub struct RecordEncoder {
+ columns: Vec<FieldBinding>,
+}
+
+impl RecordEncoder {
+ fn prepare_for_batch<'a>(
+ &'a self,
+ batch: &'a RecordBatch,
+ ) -> Result<Vec<FieldEncoder<'a>>, ArrowError> {
+ let schema_binding = batch.schema();
+ let fields = schema_binding.fields();
+ let arrays = batch.columns();
+ let mut out = Vec::with_capacity(self.columns.len());
+ for col_plan in self.columns.iter() {
+ let arrow_index = col_plan.arrow_index;
+ let array = arrays.get(arrow_index).ok_or_else(|| {
+ ArrowError::SchemaError(format!("Column index {arrow_index} out of range"))
+ })?;
+ let field = fields[arrow_index].as_ref();
+ let encoder = prepare_value_site_encoder(
+ array.as_ref(),
+ field,
+ col_plan.nullability,
+ &col_plan.plan,
+ )?;
+ out.push(encoder);
}
- DataType::Binary => {
- let arr = array.as_binary::<i32>();
- NullableEncoder::new(Encoder::Binary(BinaryEncoder(arr)), nulls)
+ Ok(out)
+ }
+
+ /// Encode a `RecordBatch` using this encoder plan.
+ ///
+ /// Tip: Wrap `out` in a `std::io::BufWriter` to reduce the overhead of many small writes.
+ pub fn encode<W: Write>(&self, out: &mut W, batch: &RecordBatch) -> Result<(), ArrowError> {
+ let mut column_encoders = self.prepare_for_batch(batch)?;
+ for row in 0..batch.num_rows() {
+ for encoder in column_encoders.iter_mut() {
+ encoder.encode(out, row)?;
+ }
}
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- let arr = array.as_primitive::<TimestampMicrosecondType>();
- NullableEncoder::new(Encoder::Timestamp(LongEncoder(arr)), nulls)
+ Ok(())
+ }
+}
+
+fn find_struct_child_index(fields: &arrow_schema::Fields, name: &str) -> Option<usize> {
+ fields.iter().position(|f| f.name() == name)
+}
+
+impl FieldPlan {
+ fn build(avro_dt: &AvroDataType, arrow_field: &Field) -> Result<Self, ArrowError> {
+ match avro_dt.codec() {
+ Codec::Struct(avro_fields) => {
+ let fields = match arrow_field.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ other => {
+ return Err(ArrowError::SchemaError(format!(
+ "Avro struct maps to Arrow Struct, found: {other:?}"
+ )))
+ }
+ };
+ let mut encoders = Vec::with_capacity(avro_fields.len());
+ for avro_field in avro_fields.iter() {
+ let name = avro_field.name().to_string();
+ let idx = find_struct_child_index(fields, &name).ok_or_else(|| {
+ ArrowError::SchemaError(format!(
+ "Struct field '{name}' not present in Arrow field '{}'",
+ arrow_field.name()
+ ))
+ })?;
+ encoders.push(FieldBinding {
+ arrow_index: idx,
+ nullability: avro_field.data_type().nullability(),
+ plan: FieldPlan::build(avro_field.data_type(), fields[idx].as_ref())?,
+ });
+ }
+ Ok(FieldPlan::Struct { encoders })
+ }
+ Codec::List(items_dt) => match arrow_field.data_type() {
+ DataType::List(field_ref) => Ok(FieldPlan::List {
+ items_nullability: items_dt.nullability(),
+ item_plan: Box::new(FieldPlan::build(items_dt.as_ref(), field_ref.as_ref())?),
+ }),
+ DataType::LargeList(field_ref) => Ok(FieldPlan::List {
+ items_nullability: items_dt.nullability(),
+ item_plan: Box::new(FieldPlan::build(items_dt.as_ref(), field_ref.as_ref())?),
+ }),
+ other => Err(ArrowError::SchemaError(format!(
+ "Avro array maps to Arrow List/LargeList, found: {other:?}"
+ ))),
+ },
+ _ => Ok(FieldPlan::Scalar),
}
- other => {
- return Err(ArrowError::NotYetImplemented(format!(
- "Unsupported data type for Avro encoding in slim build: {other:?}"
- )))
+ }
+}
+
+enum Encoder<'a> {
+ Boolean(BooleanEncoder<'a>),
+ Int(IntEncoder<'a, Int32Type>),
+ Long(LongEncoder<'a, Int64Type>),
+ Timestamp(LongEncoder<'a, TimestampMicrosecondType>),
+ Float32(F32Encoder<'a>),
+ Float64(F64Encoder<'a>),
+ Binary(BinaryEncoder<'a, i32>),
+ LargeBinary(BinaryEncoder<'a, i64>),
+ Utf8(Utf8Encoder<'a>),
+ Utf8Large(Utf8LargeEncoder<'a>),
+ List(Box<ListEncoder32<'a>>),
+ LargeList(Box<ListEncoder64<'a>>),
+ Struct(Box<StructEncoder<'a>>),
+}
+
+impl<'a> Encoder<'a> {
+ /// Encode the value at `idx`.
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
+ match self {
+ Encoder::Boolean(e) => e.encode(out, idx),
+ Encoder::Int(e) => e.encode(out, idx),
+ Encoder::Long(e) => e.encode(out, idx),
+ Encoder::Timestamp(e) => e.encode(out, idx),
+ Encoder::Float32(e) => e.encode(out, idx),
+ Encoder::Float64(e) => e.encode(out, idx),
+ Encoder::Binary(e) => e.encode(out, idx),
+ Encoder::LargeBinary(e) => e.encode(out, idx),
+ Encoder::Utf8(e) => e.encode(out, idx),
+ Encoder::Utf8Large(e) => e.encode(out, idx),
+ Encoder::List(e) => e.encode(out, idx),
+ Encoder::LargeList(e) => e.encode(out, idx),
+ Encoder::Struct(e) => e.encode(out, idx),
}
- };
- Ok(enc)
+ }
}
struct BooleanEncoder<'a>(&'a arrow_array::BooleanArray);
impl BooleanEncoder<'_> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
write_bool(out, self.0.value(idx))
}
}
@@ -230,8 +460,7 @@ impl BooleanEncoder<'_> {
/// Generic Avro `int` encoder for primitive arrays with `i32` native values.
struct IntEncoder<'a, P: ArrowPrimitiveType<Native = i32>>(&'a PrimitiveArray<P>);
impl<'a, P: ArrowPrimitiveType<Native = i32>> IntEncoder<'a, P> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
write_int(out, self.0.value(idx))
}
}
@@ -239,8 +468,7 @@ impl<'a, P: ArrowPrimitiveType<Native = i32>> IntEncoder<'a, P> {
/// Generic Avro `long` encoder for primitive arrays with `i64` native values.
struct LongEncoder<'a, P: ArrowPrimitiveType<Native = i64>>(&'a PrimitiveArray<P>);
impl<'a, P: ArrowPrimitiveType<Native = i64>> LongEncoder<'a, P> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
write_long(out, self.0.value(idx))
}
}
@@ -248,16 +476,14 @@ impl<'a, P: ArrowPrimitiveType<Native = i64>> LongEncoder<'a, P> {
/// Unified binary encoder generic over offset size (i32/i64).
struct BinaryEncoder<'a, O: OffsetSizeTrait>(&'a GenericBinaryArray<O>);
impl<'a, O: OffsetSizeTrait> BinaryEncoder<'a, O> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
write_len_prefixed(out, self.0.value(idx))
}
}
struct F32Encoder<'a>(&'a arrow_array::Float32Array);
impl F32Encoder<'_> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
// Avro float: 4 bytes, IEEE-754 little-endian
let bits = self.0.value(idx).to_bits();
out.write_all(&bits.to_le_bytes())
@@ -267,11 +493,274 @@ impl F32Encoder<'_> {
struct F64Encoder<'a>(&'a arrow_array::Float64Array);
impl F64Encoder<'_> {
- #[inline]
- fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
// Avro double: 8 bytes, IEEE-754 little-endian
let bits = self.0.value(idx).to_bits();
out.write_all(&bits.to_le_bytes())
.map_err(|e| ArrowError::IoError(format!("write f64: {e}"), e))
}
}
+
+struct Utf8GenericEncoder<'a, O: OffsetSizeTrait>(&'a GenericStringArray<O>);
+
+impl<'a, O: OffsetSizeTrait> Utf8GenericEncoder<'a, O> {
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+struct StructEncoder<'a> {
+ encoders: Vec<FieldEncoder<'a>>,
+}
+
+impl<'a> StructEncoder<'a> {
+ fn try_new(
+ array: &'a StructArray,
+ field_bindings: &[FieldBinding],
+ ) -> Result<Self, ArrowError> {
+ let DataType::Struct(fields) = array.data_type() else {
+ return Err(ArrowError::SchemaError("Expected Struct".into()));
+ };
+ let mut encoders = Vec::with_capacity(field_bindings.len());
+ for field_binding in field_bindings {
+ let idx = field_binding.arrow_index;
+ let column = array.columns().get(idx).ok_or_else(|| {
+ ArrowError::SchemaError(format!("Struct child index {idx} out of range"))
+ })?;
+ let field = fields.get(idx).ok_or_else(|| {
+ ArrowError::SchemaError(format!("Struct child index {idx} out of range"))
+ })?;
+ let encoder = prepare_value_site_encoder(
+ column.as_ref(),
+ field,
+ field_binding.nullability,
+ &field_binding.plan,
+ )?;
+ encoders.push(encoder);
+ }
+ Ok(Self { encoders })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
+ for encoder in self.encoders.iter_mut() {
+ encoder.encode(out, idx)?;
+ }
+ Ok(())
+ }
+}
+
+/// Encode a blocked range of items with Avro array block framing.
+///
+/// `write_item` must take `(out, index)` to maintain the "out-first" convention.
+fn encode_blocked_range<W: Write + ?Sized, F>(
+ out: &mut W,
+ start: usize,
+ end: usize,
+ mut write_item: F,
+) -> Result<(), ArrowError>
+where
+ F: FnMut(&mut W, usize) -> Result<(), ArrowError>,
+{
+ let len = end.saturating_sub(start);
+ if len == 0 {
+ // Zero-length terminator per Avro spec.
+ write_long(out, 0)?;
+ return Ok(());
+ }
+ // Emit a single positive block for performance, then the end marker.
+ write_long(out, len as i64)?;
+ for row in start..end {
+ write_item(out, row)?;
+ }
+ write_long(out, 0)?;
+ Ok(())
+}
+
+struct ListEncoder<'a, O: OffsetSizeTrait> {
+ list: &'a GenericListArray<O>,
+ values: FieldEncoder<'a>,
+ values_offset: usize,
+}
+
+type ListEncoder32<'a> = ListEncoder<'a, i32>;
+type ListEncoder64<'a> = ListEncoder<'a, i64>;
+
+impl<'a, O: OffsetSizeTrait> ListEncoder<'a, O> {
+ fn try_new(
+ list: &'a GenericListArray<O>,
+ items_nullability: Option<Nullability>,
+ item_plan: &FieldPlan,
+ ) -> Result<Self, ArrowError> {
+ let child_field = match list.data_type() {
+ DataType::List(field) => field.as_ref(),
+ DataType::LargeList(field) => field.as_ref(),
+ _ => {
+ return Err(ArrowError::SchemaError(
+ "Expected List or LargeList for ListEncoder".into(),
+ ))
+ }
+ };
+ let values_enc = prepare_value_site_encoder(
+ list.values().as_ref(),
+ child_field,
+ items_nullability,
+ item_plan,
+ )?;
+ Ok(Self {
+ list,
+ values: values_enc,
+ values_offset: list.values().offset(),
+ })
+ }
+
+ fn encode_list_range<W: Write + ?Sized>(
+ &mut self,
+ out: &mut W,
+ start: usize,
+ end: usize,
+ ) -> Result<(), ArrowError> {
+ encode_blocked_range(out, start, end, |out, row| {
+ self.values
+ .encode(out, row.saturating_sub(self.values_offset))
+ })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), ArrowError> {
+ let offsets = self.list.offsets();
+ let start = offsets[idx].to_usize().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!("Error converting offset[{idx}] to usize"))
+ })?;
+ let end = offsets[idx + 1].to_usize().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Error converting offset[{}] to usize",
+ idx + 1
+ ))
+ })?;
+ self.encode_list_range(out, start, end)
+ }
+}
+
+fn prepare_value_site_encoder<'a>(
+ values_array: &'a dyn Array,
+ value_field: &Field,
+ nullability: Option<Nullability>,
+ plan: &FieldPlan,
+) -> Result<FieldEncoder<'a>, ArrowError> {
+ // Effective nullability is computed here from the writer-declared site nullability and data.
+ FieldEncoder::make_encoder(values_array, value_field, plan, nullability)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::types::Int32Type;
+ use arrow_array::{
+ Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array,
+ Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, StringArray,
+ TimestampMicrosecondArray,
+ };
+ use arrow_schema::{DataType, Field, Fields};
+
+ fn zigzag_i64(v: i64) -> u64 {
+ ((v << 1) ^ (v >> 63)) as u64
+ }
+
+ fn varint(mut x: u64) -> Vec<u8> {
+ let mut out = Vec::new();
+ while (x & !0x7f) != 0 {
+ out.push(((x & 0x7f) as u8) | 0x80);
+ x >>= 7;
+ }
+ out.push((x & 0x7f) as u8);
+ out
+ }
+
+ fn avro_long_bytes(v: i64) -> Vec<u8> {
+ varint(zigzag_i64(v))
+ }
+
+ fn avro_len_prefixed_bytes(payload: &[u8]) -> Vec<u8> {
+ let mut out = avro_long_bytes(payload.len() as i64);
+ out.extend_from_slice(payload);
+ out
+ }
+
+ fn encode_all(
+ array: &dyn Array,
+ plan: &FieldPlan,
+ nullability: Option<Nullability>,
+ ) -> Vec<u8> {
+ let field = Field::new("f", array.data_type().clone(), true);
+ let mut enc = FieldEncoder::make_encoder(array, &field, plan, nullability).unwrap();
+ let mut out = Vec::new();
+ for i in 0..array.len() {
+ enc.encode(&mut out, i).unwrap();
+ }
+ out
+ }
+
+ fn assert_bytes_eq(actual: &[u8], expected: &[u8]) {
+ if actual != expected {
+ let to_hex = |b: &[u8]| {
+ b.iter()
+ .map(|x| format!("{:02X}", x))
+ .collect::<Vec<_>>()
+ .join(" ")
+ };
+ panic!(
+ "mismatch\n expected: [{}]\n actual: [{}]",
+ to_hex(expected),
+ to_hex(actual)
+ );
+ }
+ }
+
+ #[test]
+ fn binary_encoder() {
+ let values: Vec<&[u8]> = vec![b"", b"ab", b"\x00\xFF"];
+ let arr = BinaryArray::from_vec(values);
+ let mut expected = Vec::new();
+ for payload in [b"" as &[u8], b"ab", b"\x00\xFF"] {
+ expected.extend(avro_len_prefixed_bytes(payload));
+ }
+ let got = encode_all(&arr, &FieldPlan::Scalar, None);
+ assert_bytes_eq(&got, &expected);
+ }
+
+ #[test]
+ fn large_binary_encoder() {
+ let values: Vec<&[u8]> = vec![b"xyz", b""];
+ let arr = LargeBinaryArray::from_vec(values);
+ let mut expected = Vec::new();
+ for payload in [b"xyz" as &[u8], b""] {
+ expected.extend(avro_len_prefixed_bytes(payload));
+ }
+ let got = encode_all(&arr, &FieldPlan::Scalar, None);
+ assert_bytes_eq(&got, &expected);
+ }
+
+ #[test]
+ fn utf8_encoder() {
+ let arr = StringArray::from(vec!["", "A", "BC"]);
+ let mut expected = Vec::new();
+ for s in ["", "A", "BC"] {
+ expected.extend(avro_len_prefixed_bytes(s.as_bytes()));
+ }
+ let got = encode_all(&arr, &FieldPlan::Scalar, None);
+ assert_bytes_eq(&got, &expected);
+ }
+
+ #[test]
+ fn large_utf8_encoder() {
+ let arr = LargeStringArray::from(vec!["hello", ""]);
+ let mut expected = Vec::new();
+ for s in ["hello", ""] {
+ expected.extend(avro_len_prefixed_bytes(s.as_bytes()));
+ }
+ let got = encode_all(&arr, &FieldPlan::Scalar, None);
+ assert_bytes_eq(&got, &expected);
+ }
+}
diff --git a/arrow-avro/src/writer/format.rs b/arrow-avro/src/writer/format.rs
index 0ebc7a64b4..6fac9e8286 100644
--- a/arrow-avro/src/writer/format.rs
+++ b/arrow-avro/src/writer/format.rs
@@ -17,17 +17,15 @@
use crate::compression::{CompressionCodec, CODEC_METADATA_KEY};
use crate::schema::{AvroSchema, SCHEMA_METADATA_KEY};
-use crate::writer::encoder::{write_long, EncoderOptions};
+use crate::writer::encoder::write_long;
use arrow_schema::{ArrowError, Schema};
use rand::RngCore;
-use serde_json::{Map as JsonMap, Value as JsonValue};
use std::fmt::Debug;
use std::io::Write;
/// Format abstraction implemented by each container‐level writer.
pub trait AvroFormat: Debug + Default {
/// Write any bytes required at the very beginning of the output stream
- /// (file header, etc.).
/// Implementations **must not** write any record data.
fn start_stream<W: Write>(
&mut self,
@@ -44,24 +42,6 @@ pub trait AvroFormat: Debug + Default {
#[derive(Debug, Default)]
pub struct AvroOcfFormat {
sync_marker: [u8; 16],
- /// Optional encoder behavior hints to keep file header schema ordering
- /// consistent with value encoding (e.g. Impala null-second).
- encoder_options: EncoderOptions,
-}
-
-impl AvroOcfFormat {
- /// Optional helper to attach encoder options (i.e., Impala null-second) to the format.
- #[allow(dead_code)]
- pub fn with_encoder_options(mut self, opts: EncoderOptions) -> Self {
- self.encoder_options = opts;
- self
- }
-
- /// Access the options used by this format.
- #[allow(dead_code)]
- pub fn encoder_options(&self) -> &EncoderOptions {
- &self.encoder_options
- }
}
impl AvroFormat for AvroOcfFormat {
diff --git a/arrow-avro/src/writer/mod.rs b/arrow-avro/src/writer/mod.rs
index 4c46289b52..a5b2691bb8 100644
--- a/arrow-avro/src/writer/mod.rs
+++ b/arrow-avro/src/writer/mod.rs
@@ -32,13 +32,14 @@ pub mod encoder;
/// Logic for different Avro container file formats.
pub mod format;
+use crate::codec::AvroFieldBuilder;
use crate::compression::CompressionCodec;
-use crate::schema::AvroSchema;
-use crate::writer::encoder::{encode_record_batch, write_long};
+use crate::schema::{AvroSchema, SCHEMA_METADATA_KEY};
+use crate::writer::encoder::{write_long, RecordEncoder, RecordEncoderBuilder};
use crate::writer::format::{AvroBinaryFormat, AvroFormat, AvroOcfFormat};
use arrow_array::RecordBatch;
use arrow_schema::{ArrowError, Schema};
-use std::io::{self, Write};
+use std::io::Write;
use std::sync::Arc;
/// Builder to configure and create a `Writer`.
@@ -46,6 +47,7 @@ use std::sync::Arc;
pub struct WriterBuilder {
schema: Schema,
codec: Option<CompressionCodec>,
+ capacity: usize,
}
impl WriterBuilder {
@@ -54,6 +56,7 @@ impl WriterBuilder {
Self {
schema,
codec: None,
+ capacity: 1024,
}
}
@@ -63,19 +66,41 @@ impl WriterBuilder {
self
}
+ /// Sets the capacity for the given object and returns the modified instance.
+ pub fn with_capacity(mut self, capacity: usize) -> Self {
+ self.capacity = capacity;
+ self
+ }
+
/// Create a new `Writer` with specified `AvroFormat` and builder options.
- pub fn build<W, F>(self, writer: W) -> Writer<W, F>
+ /// Performs one‑time startup (header/stream init, encoder plan).
+ pub fn build<W, F>(self, mut writer: W) -> Result<Writer<W, F>, ArrowError>
where
W: Write,
F: AvroFormat,
{
- Writer {
+ let mut format = F::default();
+ let avro_schema = match self.schema.metadata.get(SCHEMA_METADATA_KEY) {
+ Some(json) => AvroSchema::new(json.clone()),
+ None => AvroSchema::try_from(&self.schema)?,
+ };
+ let mut md = self.schema.metadata().clone();
+ md.insert(
+ SCHEMA_METADATA_KEY.to_string(),
+ avro_schema.clone().json_string,
+ );
+ let schema = Arc::new(Schema::new_with_metadata(self.schema.fields().clone(), md));
+ format.start_stream(&mut writer, &schema, self.codec)?;
+ let avro_root = AvroFieldBuilder::new(&avro_schema.schema()?).build()?;
+ let encoder = RecordEncoderBuilder::new(&avro_root, schema.as_ref()).build()?;
+ Ok(Writer {
writer,
- schema: Arc::from(self.schema),
- format: F::default(),
+ schema,
+ format,
compression: self.codec,
- started: false,
- }
+ capacity: self.capacity,
+ encoder,
+ })
}
}
@@ -86,7 +111,8 @@ pub struct Writer<W: Write, F: AvroFormat> {
schema: Arc<Schema>,
format: F,
compression: Option<CompressionCodec>,
- started: bool,
+ capacity: usize,
+ encoder: RecordEncoder,
}
/// Alias for an Avro **Object Container File** writer.
@@ -95,15 +121,9 @@ pub type AvroWriter<W> = Writer<W, AvroOcfFormat>;
pub type AvroStreamWriter<W> = Writer<W, AvroBinaryFormat>;
impl<W: Write> Writer<W, AvroOcfFormat> {
- /// Convenience constructor – same as
+ /// Convenience constructor – same as [`WriterBuilder::build`] with `AvroOcfFormat`.
pub fn new(writer: W, schema: Schema) -> Result<Self, ArrowError> {
- Ok(WriterBuilder::new(schema).build::<W, AvroOcfFormat>(writer))
- }
-
- /// Change the compression codec after construction.
- pub fn with_compression(mut self, codec: Option<CompressionCodec>) -> Self {
- self.compression = codec;
- self
+ WriterBuilder::new(schema).build::<W, AvroOcfFormat>(writer)
}
/// Return a reference to the 16‑byte sync marker generated for this file.
@@ -115,19 +135,14 @@ impl<W: Write> Writer<W, AvroOcfFormat> {
impl<W: Write> Writer<W, AvroBinaryFormat> {
/// Convenience constructor to create a new [`AvroStreamWriter`].
pub fn new(writer: W, schema: Schema) -> Result<Self, ArrowError> {
- Ok(WriterBuilder::new(schema).build::<W, AvroBinaryFormat>(writer))
+ WriterBuilder::new(schema).build::<W, AvroBinaryFormat>(writer)
}
}
impl<W: Write, F: AvroFormat> Writer<W, F> {
/// Serialize one [`RecordBatch`] to the output.
pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
- if !self.started {
- self.format
- .start_stream(&mut self.writer, &self.schema, self.compression)?;
- self.started = true;
- }
- if batch.schema() != self.schema {
+ if batch.schema().fields() != self.schema.fields() {
return Err(ArrowError::SchemaError(
"Schema of RecordBatch differs from Writer schema".to_string(),
));
@@ -150,11 +165,6 @@ impl<W: Write, F: AvroFormat> Writer<W, F> {
/// Flush remaining buffered data and (for OCF) ensure the header is present.
pub fn finish(&mut self) -> Result<(), ArrowError> {
- if !self.started {
- self.format
- .start_stream(&mut self.writer, &self.schema, self.compression)?;
- self.started = true;
- }
self.writer
.flush()
.map_err(|e| ArrowError::IoError(format!("Error flushing writer: {e}"), e))
@@ -167,7 +177,7 @@ impl<W: Write, F: AvroFormat> Writer<W, F> {
fn write_ocf_block(&mut self, batch: &RecordBatch, sync: &[u8; 16]) -> Result<(), ArrowError> {
let mut buf = Vec::<u8>::with_capacity(1024);
- encode_record_batch(batch, &mut buf)?;
+ self.encoder.encode(&mut buf, batch)?;
let encoded = match self.compression {
Some(codec) => codec.compress(&buf)?,
None => buf,
@@ -184,19 +194,22 @@ impl<W: Write, F: AvroFormat> Writer<W, F> {
}
fn write_stream(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
- encode_record_batch(batch, &mut self.writer)
+ self.encoder.encode(&mut self.writer, batch)
}
}
#[cfg(test)]
mod tests {
use super::*;
+ use crate::compression::CompressionCodec;
use crate::reader::ReaderBuilder;
+ use crate::schema::{AvroSchema, SchemaStore};
use crate::test_util::arrow_test_data;
- use arrow_array::{ArrayRef, BinaryArray, Int32Array, RecordBatch, StringArray};
- use arrow_schema::{DataType, Field, Schema};
+ use arrow_array::{ArrayRef, BinaryArray, Int32Array, RecordBatch};
+ use arrow_schema::{DataType, Field, IntervalUnit, Schema};
use std::fs::File;
- use std::io::BufReader;
+ use std::io::{BufReader, Cursor};
+ use std::path::PathBuf;
use std::sync::Arc;
use tempfile::NamedTempFile;
@@ -217,10 +230,6 @@ mod tests {
.expect("failed to build test RecordBatch")
}
- fn contains_ascii(haystack: &[u8], needle: &[u8]) -> bool {
- haystack.windows(needle.len()).any(|w| w == needle)
- }
-
#[test]
fn test_ocf_writer_generates_header_and_sync() -> Result<(), ArrowError> {
let batch = make_batch();
@@ -230,12 +239,8 @@ mod tests {
writer.finish()?;
let out = writer.into_inner();
assert_eq!(&out[..4], b"Obj\x01", "OCF magic bytes missing/incorrect");
- let sync = AvroWriter::new(Vec::new(), make_schema())?
- .sync_marker()
- .cloned();
let trailer = &out[out.len() - 16..];
assert_eq!(trailer.len(), 16, "expected 16‑byte sync marker");
- let _ = sync;
Ok(())
}
@@ -309,16 +314,20 @@ mod tests {
let tmp = NamedTempFile::new().expect("create temp file");
let out_path = tmp.into_temp_path();
let out_file = File::create(&out_path).expect("create temp avro");
- let mut writer = AvroWriter::new(out_file, original.schema().as_ref().clone())?;
- if rel.contains(".snappy.") {
- writer = writer.with_compression(Some(CompressionCodec::Snappy));
+ let codec = if rel.contains(".snappy.") {
+ Some(CompressionCodec::Snappy)
} else if rel.contains(".zstandard.") {
- writer = writer.with_compression(Some(CompressionCodec::ZStandard));
+ Some(CompressionCodec::ZStandard)
} else if rel.contains(".bzip2.") {
- writer = writer.with_compression(Some(CompressionCodec::Bzip2));
+ Some(CompressionCodec::Bzip2)
} else if rel.contains(".xz.") {
- writer = writer.with_compression(Some(CompressionCodec::Xz));
- }
+ Some(CompressionCodec::Xz)
+ } else {
+ None
+ };
+ let mut writer = WriterBuilder::new(original.schema().as_ref().clone())
+ .with_compression(codec)
+ .build::<_, AvroOcfFormat>(out_file)?;
writer.write(&original)?;
writer.finish()?;
drop(writer);
@@ -338,4 +347,72 @@ mod tests {
}
Ok(())
}
+
+ #[test]
+ fn test_roundtrip_nested_records_writer() -> Result<(), ArrowError> {
+ let path = arrow_test_data("avro/nested_records.avro");
+ let rdr_file = File::open(&path).expect("open nested_records.avro");
+ let mut reader = ReaderBuilder::new()
+ .build(BufReader::new(rdr_file))
+ .expect("build reader for nested_records.avro");
+ let schema = reader.schema();
+ let batches = reader.collect::<Result<Vec<_>, _>>()?;
+ let original = arrow::compute::concat_batches(&schema, &batches).expect("concat original");
+ let tmp = NamedTempFile::new().expect("create temp file");
+ let out_path = tmp.into_temp_path();
+ {
+ let out_file = File::create(&out_path).expect("create output avro");
+ let mut writer = AvroWriter::new(out_file, original.schema().as_ref().clone())?;
+ writer.write(&original)?;
+ writer.finish()?;
+ }
+ let rt_file = File::open(&out_path).expect("open round_trip avro");
+ let mut rt_reader = ReaderBuilder::new()
+ .build(BufReader::new(rt_file))
+ .expect("build round_trip reader");
+ let rt_schema = rt_reader.schema();
+ let rt_batches = rt_reader.collect::<Result<Vec<_>, _>>()?;
+ let round_trip =
+ arrow::compute::concat_batches(&rt_schema, &rt_batches).expect("concat round_trip");
+ assert_eq!(
+ round_trip, original,
+ "Round-trip batch mismatch for nested_records.avro"
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn test_roundtrip_nested_lists_writer() -> Result<(), ArrowError> {
+ let path = arrow_test_data("avro/nested_lists.snappy.avro");
+ let rdr_file = File::open(&path).expect("open nested_lists.snappy.avro");
+ let mut reader = ReaderBuilder::new()
+ .build(BufReader::new(rdr_file))
+ .expect("build reader for nested_lists.snappy.avro");
+ let schema = reader.schema();
+ let batches = reader.collect::<Result<Vec<_>, _>>()?;
+ let original = arrow::compute::concat_batches(&schema, &batches).expect("concat original");
+ let tmp = NamedTempFile::new().expect("create temp file");
+ let out_path = tmp.into_temp_path();
+ {
+ let out_file = File::create(&out_path).expect("create output avro");
+ let mut writer = WriterBuilder::new(original.schema().as_ref().clone())
+ .with_compression(Some(CompressionCodec::Snappy))
+ .build::<_, AvroOcfFormat>(out_file)?;
+ writer.write(&original)?;
+ writer.finish()?;
+ }
+ let rt_file = File::open(&out_path).expect("open round_trip avro");
+ let mut rt_reader = ReaderBuilder::new()
+ .build(BufReader::new(rt_file))
+ .expect("build round_trip reader");
+ let rt_schema = rt_reader.schema();
+ let rt_batches = rt_reader.collect::<Result<Vec<_>, _>>()?;
+ let round_trip =
+ arrow::compute::concat_batches(&rt_schema, &rt_batches).expect("concat round_trip");
+ assert_eq!(
+ round_trip, original,
+ "Round-trip batch mismatch for nested_lists.snappy.avro"
+ );
+ Ok(())
+ }
}