You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "jecsand838 (via GitHub)" <gi...@apache.org> on 2025/09/03 04:17:17 UTC
[PR] Refactor arrow-avro writer to introduce unified `RecordEncoder` and s… [arrow-rs]
jecsand838 opened a new pull request, #8274:
URL: https://github.com/apache/arrow-rs/pull/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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2326404190
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,181 +78,417 @@ 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 has_nulls = array.null_count() > 0;
Review Comment:
```suggestion
```
no longer used?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322694024
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
+}
+
fn datatype_to_avro(
dt: &DataType,
field_name: &str,
metadata: &HashMap<String, String>,
name_gen: &mut NameGenerator,
+) -> Result<(Value, JsonMap<String, Value>), ArrowError> {
+ datatype_to_avro_with_order(dt, field_name, metadata, name_gen, Nullability::NullFirst)
+}
+
+fn arrow_field_to_avro(
+ field: &ArrowField,
+ name_gen: &mut NameGenerator,
+) -> Result<Value, ArrowError> {
+ arrow_field_to_avro_with_order(field, name_gen, Nullability::NullFirst)
+}
+
+fn datatype_to_avro_with_order(
+ dt: &DataType,
+ field_name: &str,
+ metadata: &HashMap<String, String>,
+ name_gen: &mut NameGenerator,
+ order: Nullability,
Review Comment:
Thats also a very valid call out. I'll clean this up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322745292
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
+ }
+}
+
+struct StructEncoder<'a> {
+ encoders: Vec<FieldEncoder<'a>>,
+}
+
+impl<'a> StructEncoder<'a> {
+ fn try_new(
+ array: &'a StructArray,
+ field_bindings: &[FieldBinding],
+ ) -> Result<Self, ArrowError> {
+ let fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => 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"))
+ })?
+ .as_ref();
+ let encoder = prepare_value_site_encoder(
+ column.as_ref(),
+ field,
+ field_binding.nullability,
+ &field_binding.plan,
+ )?;
+ encoders.push(encoder);
+ }
+ Ok(Self { encoders })
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ for encoder in self.encoders.iter_mut() {
+ encoder.encode(idx, out)?;
+ }
+ Ok(())
+ }
+}
+
+#[inline]
+fn encode_blocked_range<W: Write + ?Sized, F>(
+ out: &mut W,
+ start: usize,
+ end: usize,
+ mut write_item: F,
+) -> Result<(), ArrowError>
+where
+ F: FnMut(usize, &mut W) -> 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 j in start..end {
+ write_item(j, out)?;
+ }
+ 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(),
+ })
+ }
+
+ #[inline]
+ fn encode_list_range<W: Write + ?Sized>(
+ &mut self,
+ out: &mut W,
+ start: usize,
+ end: usize,
+ ) -> Result<(), ArrowError> {
+ encode_blocked_range(out, start, end, |row, out| {
+ self.values
+ .encode(row.saturating_sub(self.values_offset), out)
+ })
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> 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)
+ }
+}
+
+#[inline]
+fn prepare_value_site_encoder<'a>(
+ values_array: &'a dyn Array,
+ value_field: &Field,
+ site_nullability: Option<Nullability>,
+ plan: PlanRef<'_>,
+) -> Result<FieldEncoder<'a>, ArrowError> {
+ // Effective nullability is exactly the site's Avro-declared nullability.
+ Ok(FieldEncoder::make_encoder(values_array, value_field, plan)?
+ .with_effective_nullability(site_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, site: Option<Nullability>) -> Vec<u8> {
Review Comment:
I'll rename to `nullability`, that's a good recommendation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2323709400
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors 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>,
+ },
}
-/// 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(());
+/// Unified binding used for both top‑level columns and struct children.
+///
+/// This replaces the previous duplication between `StructChildPlan` and `ColumnPlan`.
+#[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 if not optional)
Review Comment:
> Does avro use `optional` and `required` as complementary terms, the way parquet does?
Avro does seem to follow that convention.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322774626
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
I wouldn't be shocked if one or two inline markers actually matter... but no idea which two. Can we just remove all the markers and track a follow-up issue to make profile-guided decisions on which ones to bring back?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325187158
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
Review Comment:
```suggestion
union_value_byte: u8,
```
(`byte` by itself adds nothing -- `u8` is a byte by definition)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
+ },
+ Nullable {
+ nulls: NullBuffer,
+ null_order: Nullability,
+ },
+}
+
+/// 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
Review Comment:
This is the only use site for `has_nulls`, just fold it in?
```suggestion
let null_state = match (nullability, array.null_count() > 0) {
```
(or at least define it closer to its use site)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
+ },
+ Nullable {
+ nulls: NullBuffer,
+ null_order: Nullability,
+ },
+}
+
+/// 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
+ NullState::Nullable {
+ nulls: null_buffer,
+ null_order,
+ }
+ }
+ };
+ Ok(Self {
+ encoder,
+ null_state,
+ })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => self.encoder.encode(idx, out),
+ NullState::NullableNoNulls { byte } => {
+ // Constant non-null branch byte, then value.
+ out.write_all(&[*byte]).map_err(|e| {
+ ArrowError::IoError(format!("write union value branch: {e}"), e)
+ })?;
+ self.encoder.encode(idx, out)
+ }
+ NullState::Nullable { nulls, null_order } => {
+ let is_null = nulls.is_null(idx);
+ write_optional_index(out, is_null, *null_order)?;
+ if is_null {
+ Ok(())
+ } else {
+ self.encoder.encode(idx, out)
+ }
+ }
+ }
Review Comment:
nit -- We could remove some duplication by:
```suggestion
NullState::NonNullable => {}
NullState::NullableNoNulls { byte } => out.write_all(&[*byte]).map_err(|e| {
ArrowError::IoError(format!("write union value branch: {e}"), e)
})?,
NullState::Nullable { nulls, null_order } => {
let is_null = nulls.is_null(idx);
write_optional_index(out, is_null, *null_order)?;
if is_null {
return Ok(()); // no value to write
}
}
}
self.encoder.encode(idx, out)
```
or even
```suggestion
NullState::NonNullable => {}
NullState::NullableNoNulls { byte } => out.write_all(&[*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(idx, out)
```
(splitting apart the `write_optional_index` calls will likely trigger a bunch of jump threading and inlining optimizations that more than makes up for the apparent duplication at the call site -- the compiler will basically put half of the function in each location, thanks to the hard-wired `is_null` arg value)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,49 +81,358 @@ 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>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &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);
+ writer
+ .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 {
+ 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
+ NullState::Nullable {
+ nulls: null_buffer,
+ null_order,
}
}
- enc.encode(row, out)
+ };
+ Ok(Self {
+ encoder,
+ null_state,
})
- })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => self.encoder.encode(idx, out),
+ NullState::NullableNoNulls { byte } => {
+ // Constant non-null branch byte, then value.
+ out.write_all(&[*byte]).map_err(|e| {
+ ArrowError::IoError(format!("write union value branch: {e}"), e)
+ })?;
+ self.encoder.encode(idx, out)
+ }
+ NullState::Nullable { nulls, null_order } => {
+ let is_null = nulls.is_null(idx);
+ write_optional_index(out, is_null, *null_order)?;
+ if is_null {
+ Ok(())
+ } else {
+ self.encoder.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
+ }
+}
+
+/// 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>,
+ },
+}
+
+#[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 if required)
+ nullability: Option<Nullability>,
+ /// Nested plan for this site
+ plan: FieldPlan,
+}
+
+/// 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,
+ }
+ }
+
+ /// 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 root_fields = match avro_root_dt.codec() {
+ Codec::Struct(fields) => fields,
+ _ => {
+ 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.iter() {
Review Comment:
nit: the `iter()` call should be unnecssary (thanks to `impl IntoIterator for &Self`)?
```suggestion
for root_field in root_fields {
```
##########
arrow-avro/src/schema.rs:
##########
@@ -1118,26 +1056,37 @@ fn datatype_to_avro_with_order(
Ok((val, extras))
}
-fn arrow_field_to_avro_with_order(
+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 merged = merge_extras(schema, extras);
+ Ok(if is_nullable {
+ wrap_nullable(merged, null_order)
+ } else {
+ merged
+ })
Review Comment:
nit:
```suggestion
let mut merged = merge_extras(schema, extras);
if is_nullable {
merged = wrap_nullable(merged, null_order)
}
Ok(merged)
```
(`mut` is underutilized in situations like this, and IMO `if` inside parens is hard to read+maintain)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
+ },
+ Nullable {
+ nulls: NullBuffer,
+ null_order: Nullability,
+ },
+}
+
+/// 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
Review Comment:
```suggestion
let nulls = array.nulls().cloned().ok_or_else(|| {
```
(simplifies the code below, and one less name to grok)
##########
arrow-avro/src/schema.rs:
##########
@@ -2056,4 +2133,15 @@ mod tests {
assert_eq!(arrow_field, expected);
}
+
+ #[test]
+ fn default_order_is_consistent() {
+ // Ensure TryFrom delegates to from_arrow_with_options(None)
+ 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)
+ .unwrap()
+ .json_string;
+ assert_eq!(a, b);
Review Comment:
nit
```suggestion
let b = AvroSchema::from_arrow_with_options(&arrow_schema, None)
assert_eq!(a, b.unwrap().json_string);
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,49 +81,358 @@ 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>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &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);
+ writer
+ .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 {
+ 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
+ NullState::Nullable {
+ nulls: null_buffer,
+ null_order,
}
}
- enc.encode(row, out)
+ };
+ Ok(Self {
+ encoder,
+ null_state,
})
- })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => self.encoder.encode(idx, out),
+ NullState::NullableNoNulls { byte } => {
+ // Constant non-null branch byte, then value.
+ out.write_all(&[*byte]).map_err(|e| {
+ ArrowError::IoError(format!("write union value branch: {e}"), e)
+ })?;
+ self.encoder.encode(idx, out)
+ }
+ NullState::Nullable { nulls, null_order } => {
+ let is_null = nulls.is_null(idx);
+ write_optional_index(out, is_null, *null_order)?;
+ if is_null {
+ Ok(())
+ } else {
+ self.encoder.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
+ }
+}
+
+/// 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>,
+ },
+}
+
+#[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 if required)
Review Comment:
nit:
```suggestion
/// Nullability/order for this site (None for required fields)
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -267,11 +513,275 @@ 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> {
// 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, idx: usize, out: &mut W) -> 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 fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => 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"))
+ })?
+ .as_ref();
Review Comment:
nit:
```suggestion
})?;
```
(just pass `field.as_ref()` to the function below)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
+ },
+ Nullable {
+ nulls: NullBuffer,
+ null_order: Nullability,
+ },
+}
+
+/// 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
Review Comment:
aside: I'm not sure `ok_or_else` plus `?` is actually helpful... unless the thirteen extra chars from `return Err(...);` makes the inner part spill to multiple lines (which it doesn't, in this case)
```suggestion
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()
)));
};
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -267,11 +513,275 @@ 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> {
// 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, idx: usize, out: &mut W) -> 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 fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => return Err(ArrowError::SchemaError("Expected Struct".into())),
Review Comment:
```suggestion
let DataType::Struct(fields) = array.data_type() else {
return Err(ArrowError::SchemaError("Expected Struct".into()));
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,49 +81,358 @@ 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>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &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);
+ writer
+ .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 {
+ 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
+ NullState::Nullable {
+ nulls: null_buffer,
+ null_order,
}
}
- enc.encode(row, out)
+ };
+ Ok(Self {
+ encoder,
+ null_state,
})
- })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => self.encoder.encode(idx, out),
+ NullState::NullableNoNulls { byte } => {
+ // Constant non-null branch byte, then value.
+ out.write_all(&[*byte]).map_err(|e| {
+ ArrowError::IoError(format!("write union value branch: {e}"), e)
+ })?;
+ self.encoder.encode(idx, out)
+ }
+ NullState::Nullable { nulls, null_order } => {
+ let is_null = nulls.is_null(idx);
+ write_optional_index(out, is_null, *null_order)?;
+ if is_null {
+ Ok(())
+ } else {
+ self.encoder.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
+ }
+}
+
+/// 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>,
+ },
+}
+
+#[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 if required)
+ nullability: Option<Nullability>,
+ /// Nested plan for this site
+ plan: FieldPlan,
+}
+
+/// 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,
+ }
+ }
+
+ /// 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 root_fields = match avro_root_dt.codec() {
+ Codec::Struct(fields) => fields,
+ _ => {
+ return Err(ArrowError::SchemaError(
+ "Top-level Avro schema must be a record/struct".into(),
+ ))
+ }
Review Comment:
```suggestion
let Codec::Struct(root_fields) = avro_root_dt.codec() else {
return Err(ArrowError::SchemaError(
"Top-level Avro schema must be a record/struct".into(),
));
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -267,11 +513,275 @@ 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> {
// 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, idx: usize, out: &mut W) -> 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 fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => 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"))
+ })?
+ .as_ref();
+ 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, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ for encoder in self.encoders.iter_mut() {
+ encoder.encode(idx, out)?;
+ }
+ Ok(())
+ }
+}
+
+fn encode_blocked_range<W: Write + ?Sized, F>(
+ out: &mut W,
+ start: usize,
+ end: usize,
+ mut write_item: F,
+) -> Result<(), ArrowError>
+where
+ F: FnMut(usize, &mut W) -> 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 j in start..end {
+ write_item(j, out)?;
Review Comment:
aside: It's a bit disorienting that some methods (like this one and encoder `encode`) take `out` as the second param, while others (like `write_long` below) take `out` as the first param? Can we harmonize them? I would suggest `out` is always the first param, but consistency is the main concern.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2326491468
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,181 +78,417 @@ 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 has_nulls = array.null_count() > 0;
Review Comment:
Ty for catching that. It's definitely no longer being used.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322622541
##########
arrow-avro/src/schema.rs:
##########
@@ -370,6 +371,49 @@ impl AvroSchema {
pub fn fingerprint(&self) -> Result<Fingerprint, ArrowError> {
generate_fingerprint_rabin(&self.schema()?)
}
+
+ /// 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_union_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_union_order.unwrap_or(Nullability::NullFirst);
+ let mut name_gen = NameGenerator::default();
+ let fields_json = schema
+ .fields()
+ .iter()
+ .map(|f| arrow_field_to_avro_with_order(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());
Review Comment:
> aside: Is this a well-known default name? Or just an arbitrary naming choice by this package?
While not an Avro‑spec default, `topLevelRecord` is used as a de‑facto default because several popular tools (notably Spark/Databricks) default to the same name when they synthesize an Avro schema from a struct/row.
> And does it actually matter in practice? (I guess if it mattered, the schema metadata would say so)?
Avro requires that a record have a name to be valid. Also because the record name participates in canonical form parsing, changing a record's name will change it's fingerprint.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322847044
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
One thought that struck me this morning -- do we actually need `has_nulls` at all?
Or is it actually just an intermediate calculation?
In other words, if the constructor sees `!has_nulls` then:
* set `nulls = None` even if the input array had a null buffer (because we just proved it's a no-op)
* generate `pre`
|nulls|nullability|pre|note
|:-:|:-:|:-:|-
|None|None|None|Non-nullable data with non-nullable writer
|None|None|Some|**INVALID** -- Some `pre` implies Some `nullability`
|None|Some|None|**Missed optimization** -- we should have Some `pre`
|None|Some|Some|Non-nullable data with nullable writer (leverage `pre`)
|Some|None|None|**INVALID** -- Some `nulls` but writer is non-nullable
|Some|None|Some|**INVALID** -- Some `pre` implies Some `nullability`<br>**INVALID** -- Some `nulls` but writer is non-nullable<br>**INVALID** -- Some `pre` implies None `nulls`
|Some|Some|None|Nullable data with nulls and nullable writer (cannot leverage `pre`)
|Some|Some|Some|**INVALID** -- Some `pre` implies None `nulls`
... but it's still a huge truth table with only three valid states...
Maybe better to just define an enum with three variants?
```rust
enum NullState {
// non-nullable writer with non-nullable data
NonNullable,
// nullable writer with non-nullable data
NullableButNoNulls(NullOrdering, u8),
// nullable writer with nullable data
Nullable(NullBuffer, NullOrdering),
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322698161
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
+}
+
fn datatype_to_avro(
Review Comment:
> can we keep the nice name and just add the extra arg
This is the way. I'll makes those changes. You're 100% correct.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322847044
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
One thought that struck me this morning -- do we actually need `has_nulls` at all?
Or is it actually just an intermediate calculation?
In other words, if the constructor sees `!has_nulls` then:
* set `nulls = None` even if the input array had a null buffer (because we just proved it's a no-op)
* generate `pre`
|nulls|nullability|pre|note
|:-:|:-:|:-:|-
|None|None|None|Non-nullable data with non-nullable writer
|None|None|Some|**INVALID** -- Some `pre` implies Some `nullability`
|None|Some|None|**Missed optimization** -- we should have Some `pre`
|None|Some|Some|Non-nullable data with nullable writer (leverage `pre`)
|Some|None|None|**INVALID** -- Some `nulls` but writer is non-nullable
|Some|None|Some|**INVALID** -- Some `pre` implies Some `nullability`<br>**INVALID** -- Some `nulls` but writer is non-nullable<br>**INVALID** -- Some `pre` implies None `nulls`
|Some|Some|None|Nullable data with nulls and nullable writer (cannot leverage `pre`)
|Some|Some|Some|**INVALID** -- Some `pre` implies None `nulls`
... but it's still a huge truth table with only three valid states...
Maybe better to just define an enum with three variants?
```rust
enum NullState {
// non-nullable data with non-nullable writer
NonNullable,
// non-nullable data with nullable writer
NullableButNoNulls(NullOrdering, u8),
// nullable data with nullable writer
Nullable(NullBuffer, NullOrdering),
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322665974
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
@scovich I think you're really onto something here. I'll dive into this further to see how this can be simplified and re-thought out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325734395
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,36 +81,186 @@ 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_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- order: Nullability,
+ null_order: Nullability,
) -> Result<(), ArrowError> {
- // For NullFirst: null => 0x00, value => 0x02
- // For NullSecond: value => 0x00, null => 0x02
- let byte = match order {
- Nullability::NullFirst => {
- if is_null {
- 0x00
- } else {
- 0x02
- }
- }
- Nullability::NullSecond => {
- if is_null {
- 0x02
- } else {
- 0x00
- }
- }
- };
+ let byte = union_value_branch_byte(null_order, is_null);
writer
.write_all(&[byte])
.map_err(|e| ArrowError::IoError(format!("write union branch: {e}"), e))
}
-/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+#[derive(Debug, Clone)]
+enum NullState {
+ NonNullable,
+ NullableNoNulls {
+ byte: u8,
Review Comment:
100%, much more descriptive.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322702309
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -17,25 +17,26 @@
//! 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,
+ ArrowPrimitiveType, Float32Type, Float64Type, Int32Type, Int64Type, IntervalMonthDayNanoType,
+ TimestampMicrosecondType,
+};
+use arrow_array::{
+ Array, Decimal128Array, Decimal256Array, Decimal32Array, Decimal64Array, DictionaryArray,
+ FixedSizeBinaryArray, GenericBinaryArray, GenericListArray, GenericStringArray, LargeListArray,
+ ListArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StringArray, StructArray,
};
-use arrow_array::OffsetSizeTrait;
-use arrow_array::{Array, GenericBinaryArray, PrimitiveArray, RecordBatch};
use arrow_buffer::NullBuffer;
-use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit};
+use arrow_schema::{ArrowError, DataType, Field, IntervalUnit, Schema as ArrowSchema, TimeUnit};
use std::io::Write;
+use std::sync::Arc;
+use uuid::Uuid;
-/// 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
-}
+/// Plan reference passed to the unified encoder constructor (required).
+type PlanRef<'p> = &'p FieldPlan;
Review Comment:
Good catch!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Decoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2317718792
##########
arrow-avro/src/codec.rs:
##########
@@ -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,
-}
Review Comment:
I moved this to `src/schema.rs` to keep the imports clean and non-circular.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "mbrobbel (via GitHub)" <gi...@apache.org>.
mbrobbel commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2329419436
##########
arrow-avro/src/schema.rs:
##########
@@ -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!("Serialising Avro JSON failed: {e}")))?;
Review Comment:
```suggestion
.map_err(|e| ArrowError::SchemaError(format!("Serializing Avro JSON failed: {e}")))?;
```
##########
arrow-avro/src/schema.rs:
##########
@@ -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
Review Comment:
```suggestion
/// the exact header encoding alignment; otherwise, a new JSON is generated
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,181 +74,416 @@ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
Review Comment:
```suggestion
return Err(ArrowError::NotYetImplemented(format!(
"Avro writer: {other:?} not yet supported".into(),
)));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3259422115
> Bunch of nits to simplify the code, but otherwise ready to merge.
>
> > There are two items I wanted to follow-up on in future PRs:
> > ```
> > 1. Further `arrow-avro/src/schema.rs` file refactoring (if possible). As I updated the file I saw other potential areas to improve. For example, I think if we introduced some sort of struct to hold context for `datatype_to_avro`, we could simplify and improve the logic.
> > ```
>
> Definitely future work.
>
> > ```
> > 2. Renaming `Nullability`. The name change will impact other areas of the code not related to the `Writer`. I'd also want to consider the impact `Union` type support will have as well when we get there. But I do think @scovich has a valid point about the current naming convention.
> > ```
>
> Yes. Mixing refactors with new functionality is always messy.
>
> I would recommend two other follow-ups: 3. Profile-guided placement of `#[inline]` markers 4. Harmonize arg ordering for the various encoding methods that take `out`; some pass it as the first arg, others as the second arg.
@scovich All great call outs. I accepted your latest suggestions and pushed up changes that harmonized the arg ordering for the methods taking `out` as a parameter. As always thank you for the solid and thorough review 😄. We'll also be sure to follow-up on those PRs.
@alamb @mbrobbel Would either of you be able to look over this PR if you get a chance?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3266907272
I'll plan to merge this PR once the CI passes
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3267117205
Thanks @jecsand838
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325128883
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
I bet profile-guided optimizations would give back far more than 4%, since it covers far more than just inlining decisions. If we care about optimizing that much, we should probably go that route? But that requires benchmarking and profiling infra that arrow-rs probably lacks. An hybrid (and probably almost optimal) approach would be to manually gather .profdata, and then see what `llvm-profdata show` has to say about these functions.
Humans are provably terrible at guessing where to sprinkle hints like inlining or branch prediction, and so I would expect that some of the inline markers are badly-placed and actually _hurting_ performance, which eats into the performance gains of the well-placed markers. Automatic or manual profile-guided optimization at least gives some concrete data to replace those guesses.
Overall, my recommended approach would be to remove all the inline markers for now (focus on getting correct code in), and if we want to chase the 4% that could be a follow-up item to do it right. Preferably after the code has stabilized, since it's still evolving rapidly right now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325742520
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -267,11 +513,275 @@ 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> {
// 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, idx: usize, out: &mut W) -> 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 fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => 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"))
+ })?
+ .as_ref();
+ 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, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ for encoder in self.encoders.iter_mut() {
+ encoder.encode(idx, out)?;
+ }
+ Ok(())
+ }
+}
+
+fn encode_blocked_range<W: Write + ?Sized, F>(
+ out: &mut W,
+ start: usize,
+ end: usize,
+ mut write_item: F,
+) -> Result<(), ArrowError>
+where
+ F: FnMut(usize, &mut W) -> 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 j in start..end {
+ write_item(j, out)?;
Review Comment:
That's very fair, I'll clean that up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2318689900
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
Review Comment:
nit: Worthwhile to sacrifice a couple lines to reduce redundancy?
```suggestion
let null = Value::String("null".into());
let elements = match order {
Nullability::NullFirst => vec![null, inner],
Nullability::NullSecond => vec![inner, null],
};
Value::Array(elements)
```
##########
arrow-avro/src/schema.rs:
##########
@@ -370,6 +371,49 @@ impl AvroSchema {
pub fn fingerprint(&self) -> Result<Fingerprint, ArrowError> {
generate_fingerprint_rabin(&self.schema()?)
}
+
+ /// 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_union_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_union_order.unwrap_or(Nullability::NullFirst);
+ let mut name_gen = NameGenerator::default();
+ let fields_json = schema
+ .fields()
+ .iter()
+ .map(|f| arrow_field_to_avro_with_order(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());
Review Comment:
aside: Is this a well-known default name? Or just an arbitrary naming choice by this package? And does it actually matter in practice? (I guess if it mattered, the `schema` metadata would say so)?
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
+}
+
fn datatype_to_avro(
dt: &DataType,
field_name: &str,
metadata: &HashMap<String, String>,
name_gen: &mut NameGenerator,
+) -> Result<(Value, JsonMap<String, Value>), ArrowError> {
+ datatype_to_avro_with_order(dt, field_name, metadata, name_gen, Nullability::NullFirst)
+}
+
+fn arrow_field_to_avro(
+ field: &ArrowField,
+ name_gen: &mut NameGenerator,
+) -> Result<Value, ArrowError> {
+ arrow_field_to_avro_with_order(field, name_gen, Nullability::NullFirst)
+}
Review Comment:
Why not define this down below, similar to what you did here with `data_type_to_avro[_with_order]`?
##########
arrow-avro/src/schema.rs:
##########
@@ -370,6 +371,49 @@ impl AvroSchema {
pub fn fingerprint(&self) -> Result<Fingerprint, ArrowError> {
generate_fingerprint_rabin(&self.schema()?)
}
+
+ /// 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_union_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_union_order.unwrap_or(Nullability::NullFirst);
Review Comment:
I see at least three places that use `Nullability::NullFirst` as a default when no order was requested. Would it make sense to derive `Default` for the enum to make this decision publicly explicit?
```suggestion
let order = null_union_order.unwrap_or_default();
```
##########
arrow-avro/src/writer/format.rs:
##########
@@ -44,24 +43,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
- }
Review Comment:
This code is removed in favor of the new approach, correct?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
Review Comment:
nit: is this cleaner?
```suggestion
let byte = match (order, is_null) {
(Nullability::NullFirst, true) | (Nullablility::NullSecond, false) => 0x00,
(Nullability::NullFirst, false) | (Nullability::NullSecond, true) => 0x02,
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -17,25 +17,26 @@
//! 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,
+ ArrowPrimitiveType, Float32Type, Float64Type, Int32Type, Int64Type, IntervalMonthDayNanoType,
+ TimestampMicrosecondType,
+};
+use arrow_array::{
+ Array, Decimal128Array, Decimal256Array, Decimal32Array, Decimal64Array, DictionaryArray,
+ FixedSizeBinaryArray, GenericBinaryArray, GenericListArray, GenericStringArray, LargeListArray,
+ ListArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StringArray, StructArray,
};
-use arrow_array::OffsetSizeTrait;
-use arrow_array::{Array, GenericBinaryArray, PrimitiveArray, RecordBatch};
use arrow_buffer::NullBuffer;
-use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit};
+use arrow_schema::{ArrowError, DataType, Field, IntervalUnit, Schema as ArrowSchema, TimeUnit};
use std::io::Write;
+use std::sync::Arc;
+use uuid::Uuid;
-/// 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
-}
+/// Plan reference passed to the unified encoder constructor (required).
+type PlanRef<'p> = &'p FieldPlan;
Review Comment:
I'm not sure this is a helpful type alias? It only has two use sites, both as function arg types with elided lifetimes. With the alias they look like this, because lifetime elision can't fully simplify the signature:
```rust
plan: PlanRef<'_>,
```
vs. without the alias and complete lifetime elision:
```rust
plan: &Plan,
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
I see a lot of these `#[inline]` markers, but the compiler already inlines code pretty aggressively.
Do we have some evidence that these markers actually improve performance?
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
+}
+
fn datatype_to_avro(
Review Comment:
Do we need to preserve existing callsites after this PR? Or can we keep the nice name and just add the extra arg? Asking because the function is private, I can only find six callsites in my local arrow-rs checkout, and this PR changes all six of them.
A similar question applies to `arrow_field_to_avro`, except this PR only seems to update one of the two callsites; the callsite in `impl TryFrom<&ArrowSchema> for AvroSchema` remains unchanged. Seems reasonable to just have one function (with the extra arg) and pass `NullFirst` manually as needed?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors 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>,
+ },
}
-/// 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(());
+/// Unified binding used for both top‑level columns and struct children.
+///
+/// This replaces the previous duplication between `StructChildPlan` and `ColumnPlan`.
+#[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 if not optional)
Review Comment:
Also (tiny nit), it might be slightly more intuitive to put the plan (noun) first and nullability (adjective describing noun) second, both here and in FieldPlan::List above?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
Staring more at the code, I don't think `with_effective_nullability` should exist as a separate method. The `nullability` and `pre` are too closely tied up with `nulls` and `has_nulls`.
Also, there are too many Option here, which explodes the number of possible states we have to consider. Is there a way to simplify it?
|nulls|has_nulls|nullability|pre|note
|:-:|:-:|:-:|:-:|-
|None|true|*|*|**INVALID** -- `has_null` implies Some `nulls`
|*|true|None|*|**INVALID** -- `has_nulls` but writer is non-nullable
|*|*|None|Some|**INVALID** -- Some `pre` implies Some `nullability`
|*|true|*|Some|**INVALID** -- Some `pre` implies not `has_nulls`
|*|false|Some|None|**Missed optimization** -- we should have Some `pre`
|None|false|None|None|Everyone agrees this is a non-nullable column
|None|false|Some|Some|Non-nullable data with nullable writer (leverage `pre`)
|Some|false|None|None|Nullable data with no nulls and non-nullable writer
|Some|false|Some|Some|Nullable data with no nulls and nullable writer (leverage `pre`)
|Some|true|Some|None|Nullable data with nulls and nullable writer (cannot leverage `pre`)
##########
arrow-avro/src/schema.rs:
##########
@@ -790,12 +856,35 @@ 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, order: Nullability) -> Value {
+ match order {
+ Nullability::NullFirst => Value::Array(vec![Value::String("null".into()), inner]),
+ Nullability::NullSecond => Value::Array(vec![inner, Value::String("null".into())]),
+ }
+}
+
fn datatype_to_avro(
dt: &DataType,
field_name: &str,
metadata: &HashMap<String, String>,
name_gen: &mut NameGenerator,
+) -> Result<(Value, JsonMap<String, Value>), ArrowError> {
+ datatype_to_avro_with_order(dt, field_name, metadata, name_gen, Nullability::NullFirst)
+}
+
+fn arrow_field_to_avro(
+ field: &ArrowField,
+ name_gen: &mut NameGenerator,
+) -> Result<Value, ArrowError> {
+ arrow_field_to_avro_with_order(field, name_gen, Nullability::NullFirst)
+}
+
+fn datatype_to_avro_with_order(
+ dt: &DataType,
+ field_name: &str,
+ metadata: &HashMap<String, String>,
+ name_gen: &mut NameGenerator,
+ order: Nullability,
Review Comment:
In isolation, "order" is neither intuitive nor self-describing. For that matter, `Nullability` isn't, either. Perhaps `null_order` and `NullOrder` would be better names, along with e.g. `datatype_to_avro_with_null_order` (tho see other comment about the function's name)?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+/// optional branch order can be honored exactly as declared by the schema.
+#[derive(Debug, Clone)]
+enum FieldPlan {
Review Comment:
No `Map` variant needed because maps are just a fancy way of interpreting a struct?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors 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>,
+ },
}
-/// 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(());
+/// Unified binding used for both top‑level columns and struct children.
+///
+/// This replaces the previous duplication between `StructChildPlan` and `ColumnPlan`.
+#[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 if not optional)
Review Comment:
Does avro use `optional` and `required` as complementary terms, the way parquet does?
```suggestion
/// Nullability/order for this site (None if required)
```
##########
arrow-avro/src/schema.rs:
##########
@@ -909,20 +998,42 @@ 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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
+ let items_schema = if child.is_nullable() {
+ wrap_nullable(items_with_meta, order)
+ } else {
+ items_with_meta
+ };
Review Comment:
This exact pattern shows up four times (List/LargeList, FixedSizedList, Map, and struct field). Can we pull out a helper that all four call sites can use? It might look something like this:
```rust
fn process_datatype_with_nullability(
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_with_order(dt, field_name, metadata, name_gen, null_order)?;
let schema = merge_extras(schema, extras);
if is_nullable {
wrap_nullable(schema, null_union_order)
} else {
Ok(schema)
}
}
```
(terrible name for the helper, but you get the idea)
There are two additional sites (Dictionary and RunEndEncoded) that could also use the helper -- they'd just need to pass `is_nullable: false`.
If you're looking for ways to split up the PR, a useful "prerefactor" might:
* plumb through the new (largely unused) `order` arg in `datatype_to_avro_with_order` and friends
* add+use this new helper at the 4-6 call sites that benefit from it
##########
arrow-avro/src/writer/mod.rs:
##########
@@ -63,19 +66,42 @@ 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 = if let Some(json) = self.schema.metadata.get(SCHEMA_METADATA_KEY) {
+ AvroSchema::new(json.clone())
+ } else {
+ AvroSchema::try_from(&self.schema)?
+ };
Review Comment:
aside: I'm always torn whether to use if let Some vs. match at times like this...
```suggestion
let avro_schema = match self.schema.metadata.get(SCHEMA_METADATA_KEY) {
Some(json) => AvroSchema::new(json.clone()),
None => AvroSchema::try_from(&self.schema)?,
};
```
(I think it _might_ be a net win this time?)
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
Review Comment:
There seems to be some overlap with `write_optional_index`; is there a way to factor out the common logic (and magic constants) in a way that both can use?
<details>
<summary>Mabye something like this?</summary>
```rust
fn union_value_branch_byte(null_order: NullOrder, is_null: bool) -> u8 {
let nulls_first = null_order == NullOrder::NullFirst;
nulls_first == is_null { 0x00 } else { 0x02 }
}
```
and then the code in `write_optional_index` changes to:
```rust
let byte = union_value_branch_byte(null_order, is_null);
```
and this code here:
```suggestion
// ... explanatory comment here ...
let null_order = nullability?;
(!self.has_nulls()).then(|| union_value_branch_byte(null_order, false))
```
</details>
Also, would `(None, true)` be an error, or at least an invalid combination?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
This seems a bit complex/redundant (two `self.encode_inner` calls). Can it be simplified while still preserving the `pre` shortcut?
```suggestion
if let Some(b) = self.pre {
out
.write_all(&[b])
.map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))?;
} else if let Some(null_order) = self.nullability {
let is_null = self.is_null(idx);
write_optional_index(out, is_null, null_order)?;
if is_null {
return Ok(());
}
}
self.encode_inner(idx, out)
```
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
Review Comment:
Could also keep a similar approach to the existing code:
```rust
let nulls_first = order == Nullability::NullFirst;
let byte = if nulls_first == is_null { 0x00 } else { 0x02 };
```
... very compact but perhaps harder to grok
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
Review Comment:
Also -- What should happen if `self.is_null(idx) && self.nullability.is_none()`?
Should `with_effective_nullability` block that combination by returning an error?
Or do we need to somehow handle it here?
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
+ }
+
+ #[inline]
+ fn encode_inner<W: Write + ?Sized>(
+ &mut self,
+ idx: usize,
+ out: &mut W,
+ ) -> Result<(), ArrowError> {
+ self.encoder.encode(idx, out)
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ if let Some(b) = self.pre {
+ return out
+ .write_all(&[b])
+ .map_err(|e| ArrowError::IoError(format!("write union value branch: {e}"), e))
+ .and_then(|_| self.encode_inner(idx, out));
+ }
+ if let Some(order) = self.nullability {
+ let is_null = self.is_null(idx);
+ write_optional_index(out, is_null, order)?;
+ if is_null {
+ return Ok(());
+ }
+ }
+ self.encode_inner(idx, out)
+ }
+}
+
+struct StructEncoder<'a> {
+ encoders: Vec<FieldEncoder<'a>>,
+}
+
+impl<'a> StructEncoder<'a> {
+ fn try_new(
+ array: &'a StructArray,
+ field_bindings: &[FieldBinding],
+ ) -> Result<Self, ArrowError> {
+ let fields = match array.data_type() {
+ DataType::Struct(struct_fields) => struct_fields,
+ _ => 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"))
+ })?
+ .as_ref();
+ let encoder = prepare_value_site_encoder(
+ column.as_ref(),
+ field,
+ field_binding.nullability,
+ &field_binding.plan,
+ )?;
+ encoders.push(encoder);
+ }
+ Ok(Self { encoders })
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ for encoder in self.encoders.iter_mut() {
+ encoder.encode(idx, out)?;
+ }
+ Ok(())
+ }
+}
+
+#[inline]
+fn encode_blocked_range<W: Write + ?Sized, F>(
+ out: &mut W,
+ start: usize,
+ end: usize,
+ mut write_item: F,
+) -> Result<(), ArrowError>
+where
+ F: FnMut(usize, &mut W) -> 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 j in start..end {
+ write_item(j, out)?;
+ }
+ 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(),
+ })
+ }
+
+ #[inline]
+ fn encode_list_range<W: Write + ?Sized>(
+ &mut self,
+ out: &mut W,
+ start: usize,
+ end: usize,
+ ) -> Result<(), ArrowError> {
+ encode_blocked_range(out, start, end, |row, out| {
+ self.values
+ .encode(row.saturating_sub(self.values_offset), out)
+ })
+ }
+
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> 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)
+ }
+}
+
+#[inline]
+fn prepare_value_site_encoder<'a>(
+ values_array: &'a dyn Array,
+ value_field: &Field,
+ site_nullability: Option<Nullability>,
+ plan: PlanRef<'_>,
+) -> Result<FieldEncoder<'a>, ArrowError> {
+ // Effective nullability is exactly the site's Avro-declared nullability.
+ Ok(FieldEncoder::make_encoder(values_array, value_field, plan)?
+ .with_effective_nullability(site_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, site: Option<Nullability>) -> Vec<u8> {
Review Comment:
`site` is a really unintuitive name for this... earlier code used `site_nullability` which seems better?
But why not just `nullability` for both?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "nathaniel-d-ef (via GitHub)" <gi...@apache.org>.
nathaniel-d-ef commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2319092823
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
Review Comment:
This is explained elsewhere, I think we can eliminate the comments. Maybe alias the values as constants globally.
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors 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>,
+ },
}
-/// 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(());
+/// Unified binding used for both top‑level columns and struct children.
+///
+/// This replaces the previous duplication between `StructChildPlan` and `ColumnPlan`.
Review Comment:
I would omit this for committed code. The context is good but only as far as the purposes of this PR IMO
##########
arrow-avro/src/schema.rs:
##########
@@ -909,20 +998,42 @@ 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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
+ let items_schema = if child.is_nullable() {
+ wrap_nullable(items_with_meta, order)
+ } else {
+ items_with_meta
+ };
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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
Review Comment:
Perhaps worth it to put this nullable logic in `datatype_to_avro_with_order` as it's duplicated here and in `Map`?
##########
arrow-avro/src/writer/format.rs:
##########
@@ -85,7 +66,7 @@ impl AvroFormat for AvroOcfFormat {
Some(CompressionCodec::Xz) => "xz",
None => "null",
};
- write_long(writer, 2)?;
+ write_long(writer, 2)?; // two entries
Review Comment:
This note seems unnecessary
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3256030690
@scovich @nathaniel-d-ef I really appreciate the solid reviews.
The changes I just pushed should address most of your feedback. There are two items I wanted to follow-up on in future PRs:
1. Further `arrow-avro/src/schema.rs` file refactoring. As I update the file I saw other potential areas to improve. For example, I think if we introduced some sort of struct to hold context for `datatype_to_avro`, we could simplify and improve the logic.
2. Renaming `Nullability`. The name change will impact other areas of the code not related to the `Writer`. I'd also want to consider the impact `Union` type support will have as well when we get there. But I do think @scovich has a valid point about the current naming convention.
Let me know what you all think and what else needs to be done to get this PR into solid shape whenever you get a chance.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322657655
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
That's a good callout. Before I push up my changes, I'll benchmark with and without the `#[inline]` markers to determine whether they add value or not.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322688001
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
Review Comment:
I like that third approach. Also @nathaniel-d-ef recommended adding a constant for `0x00` and `0x02` to improve readability.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Decoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3247646856
@scovich @alamb Would either of you be able to review this one? This should be the last large PR for `arrow-avro` as it lays the full foundation for the Writer. I tried to slim it down as much as possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#issuecomment-3267115930
Integration test failure is unrelated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322628559
##########
arrow-avro/src/schema.rs:
##########
@@ -909,20 +998,42 @@ 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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
+ let items_schema = if child.is_nullable() {
+ wrap_nullable(items_with_meta, order)
+ } else {
+ items_with_meta
+ };
Review Comment:
Solid catch, I'll add that in!
##########
arrow-avro/src/schema.rs:
##########
@@ -909,20 +998,42 @@ 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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
+ let items_schema = if child.is_nullable() {
+ wrap_nullable(items_with_meta, order)
+ } else {
+ items_with_meta
+ };
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_inner, ie) = datatype_to_avro_with_order(
+ child.data_type(),
+ child.name(),
+ child.metadata(),
+ name_gen,
+ order,
+ )?;
+ let items_with_meta = merge_extras(items_inner, ie);
Review Comment:
Solid catch, @scovich caught this one too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322647719
##########
arrow-avro/src/writer/format.rs:
##########
@@ -44,24 +43,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
- }
Review Comment:
That's correct. I realized having encoder options created two sources of truth informing encoder behavior: The schema and the encoder options. The correct way is to have the schema be the only source of truth, otherwise we risk having records written that deviate form the associated writer schema, i.e. rows written to an Avro file that cannot be decoding using the writer schema provided in the file's header.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322743187
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
+ fn has_nulls(&self) -> bool {
+ self.has_nulls
+ }
+
+ #[inline]
+ fn is_null(&self, idx: usize) -> bool {
+ self.nulls
+ .as_ref()
+ .is_some_and(|null_buffer| null_buffer.is_null(idx))
+ }
+
+ #[inline]
+ fn with_effective_nullability(mut self, order: Option<Nullability>) -> Self {
+ self.nullability = order;
+ self.pre = self.precomputed_union_value_branch(order);
+ self
+ }
+
+ #[inline]
+ fn precomputed_union_value_branch(&self, order: Option<Nullability>) -> Option<u8> {
+ match (order, self.has_nulls()) {
+ (Some(Nullability::NullFirst), false) => Some(0x02), // value branch index 1
+ (Some(Nullability::NullSecond), false) => Some(0x00), // value branch index 0
+ _ => None,
+ }
Review Comment:
> There seems to be some overlap with write_optional_index; is there a way to factor out the common logic (and magic constants) in a way that both can use?
100%.
> Also, would (None, true) be an error, or at least an invalid combination?
That's also correct.
I'll make those changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322835463
##########
arrow-avro/src/schema.rs:
##########
@@ -370,6 +371,49 @@ impl AvroSchema {
pub fn fingerprint(&self) -> Result<Fingerprint, ArrowError> {
generate_fingerprint_rabin(&self.schema()?)
}
+
+ /// 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_union_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_union_order.unwrap_or(Nullability::NullFirst);
+ let mut name_gen = NameGenerator::default();
+ let fields_json = schema
+ .fields()
+ .iter()
+ .map(|f| arrow_field_to_avro_with_order(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());
Review Comment:
I think it's good enough in practice, but we can always follow-up on this if needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325251775
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
100% fair point. I went ahead and removed them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325740170
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -83,49 +81,358 @@ 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>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &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);
+ writer
+ .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 {
+ 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 has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ // Compute the effective null state from writer-declared nullability and data nulls.
+ let null_state = match (nullability, has_nulls) {
+ (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.
+ let byte = union_value_branch_byte(order, false);
+ NullState::NullableNoNulls { byte }
+ }
+ (Some(null_order), true) => {
+ let null_buffer = array.nulls().cloned().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Array for Avro site '{}' reports nulls but has no null buffer",
+ field.name()
+ ))
+ })?;
+ NullState::Nullable {
+ nulls: null_buffer,
+ null_order,
}
}
- enc.encode(row, out)
+ };
+ Ok(Self {
+ encoder,
+ null_state,
})
- })
+ }
+
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ match &self.null_state {
+ NullState::NonNullable => self.encoder.encode(idx, out),
+ NullState::NullableNoNulls { byte } => {
+ // Constant non-null branch byte, then value.
+ out.write_all(&[*byte]).map_err(|e| {
+ ArrowError::IoError(format!("write union value branch: {e}"), e)
+ })?;
+ self.encoder.encode(idx, out)
+ }
+ NullState::Nullable { nulls, null_order } => {
+ let is_null = nulls.is_null(idx);
+ write_optional_index(out, is_null, *null_order)?;
+ if is_null {
+ Ok(())
+ } else {
+ self.encoder.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
+ }
+}
+
+/// 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>,
+ },
+}
+
+#[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 if required)
+ nullability: Option<Nullability>,
+ /// Nested plan for this site
+ plan: FieldPlan,
+}
+
+/// 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,
+ }
+ }
+
+ /// 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 root_fields = match avro_root_dt.codec() {
+ Codec::Struct(fields) => fields,
+ _ => {
+ 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.iter() {
Review Comment:
You're right, this is a good catch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2325128883
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
I bet profile-guided optimizations would give back far more than 4%, since it covers far more than just inlining decisions. If we care about optimizing that much, we should probably go that route? But that requires benchmarking and profiling infra that arrow-rs probably lacks. An hybrid (and probably almost optimal) approach would be to manually gather .profdata, and then see what `llvm-profdata show` has to say about these functions.
Humans are provably terrible at guessing where to sprinkle hints like inlining or branch prediction, and so I would expect that some of the inline markers are badly-placed and actually _hurting_ performance -- especially for real-world (non-microbenchmark) scenarios where CPU instruction cache capacity matters -- which eats into the performance gains of the well-placed markers. Automatic or manual profile-guided optimization at least gives some concrete data to replace those guesses.
Overall, my recommended approach would be to remove all the inline markers for now (focus on getting correct code in), and if we want to chase the 4% that could be a follow-up item to do it right. Preferably after the code has stabilized, since it's still evolving rapidly right now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322649671
##########
arrow-avro/src/writer/mod.rs:
##########
@@ -63,19 +66,42 @@ 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 = if let Some(json) = self.schema.metadata.get(SCHEMA_METADATA_KEY) {
+ AvroSchema::new(json.clone())
+ } else {
+ AvroSchema::try_from(&self.schema)?
+ };
Review Comment:
That is much cleaner and idiomatic.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322761117
##########
arrow-avro/src/schema.rs:
##########
@@ -370,6 +371,49 @@ impl AvroSchema {
pub fn fingerprint(&self) -> Result<Fingerprint, ArrowError> {
generate_fingerprint_rabin(&self.schema()?)
}
+
+ /// 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_union_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_union_order.unwrap_or(Nullability::NullFirst);
+ let mut name_gen = NameGenerator::default();
+ let fields_json = schema
+ .fields()
+ .iter()
+ .map(|f| arrow_field_to_avro_with_order(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());
Review Comment:
So do we need to make this tunable somehow? Or the de facto default is good enough in practice?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322647719
##########
arrow-avro/src/writer/format.rs:
##########
@@ -44,24 +43,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
- }
Review Comment:
That's correct. I realized having encoder options created two sources of truth informing encoder behavior: The schema and the encoder options. The correct way is to have the schema be the only source of truth, otherwise we risk having records written that deviate form the associated writer schema, i.e. rows written to an Avro file that cannot be decoded using the writer schema provided in the file's header.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322672654
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
+ };
+ writer
+ .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())
+/// Per‑site encoder plan for a field. This mirrors Avro structure so nested
+/// optional branch order can be honored exactly as declared by the schema.
+#[derive(Debug, Clone)]
+enum FieldPlan {
Review Comment:
@scovich I have additional plans for maps, enums, decimals, etc on another branch. I just removed them from this PR to reduce the size.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "scovich (via GitHub)" <gi...@apache.org>.
scovich commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2322867316
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -84,48 +85,211 @@ fn write_bool<W: Write + ?Sized>(writer: &mut W, v: bool) -> Result<(), ArrowErr
/// - Null-first (default): null => 0, value => 1
/// - Null-second (Impala): value => 0, null => 1
#[inline]
-fn write_optional_branch<W: Write + ?Sized>(
+fn write_optional_index<W: Write + ?Sized>(
writer: &mut W,
is_null: bool,
- impala_mode: bool,
+ order: Nullability,
) -> Result<(), ArrowError> {
- let branch = if impala_mode == is_null { 1 } else { 0 };
- write_int(writer, branch)
+ // For NullFirst: null => 0x00, value => 0x02
+ // For NullSecond: value => 0x00, null => 0x02
+ let byte = match order {
+ Nullability::NullFirst => {
+ if is_null {
+ 0x00
+ } else {
+ 0x02
+ }
+ }
+ Nullability::NullSecond => {
+ if is_null {
+ 0x02
+ } else {
+ 0x00
+ }
+ }
Review Comment:
Tho if there's only one place the two constants appear (see my other comment), then naming them matters less and we just need a decent code comment to explain it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Added List and Struct Encoding to arrow-avro Writer [arrow-rs]
Posted by "jecsand838 (via GitHub)" <gi...@apache.org>.
jecsand838 commented on code in PR #8274:
URL: https://github.com/apache/arrow-rs/pull/8274#discussion_r2323656246
##########
arrow-avro/src/writer/encoder.rs:
##########
@@ -275,3 +384,434 @@ impl F64Encoder<'_> {
.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> {
+ #[inline]
+ fn encode<W: Write + ?Sized>(&mut self, idx: usize, out: &mut W) -> Result<(), ArrowError> {
+ write_len_prefixed(out, self.0.value(idx).as_bytes())
+ }
+}
+
+type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
+type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
+
+/// Unified field encoder:
+/// - Holds the inner `Encoder` (by value)
+/// - Tracks the column/site null buffer and whether any nulls exist
+/// - Carries per-site Avro `Nullability` and precomputed union branch (fast path)
+pub struct FieldEncoder<'a> {
+ encoder: Encoder<'a>,
+ nulls: Option<NullBuffer>,
+ has_nulls: bool,
+ nullability: Option<Nullability>,
+ /// Precomputed constant branch byte if the site is nullable but contains no nulls
+ pre: Option<u8>,
+}
+
+impl<'a> FieldEncoder<'a> {
+ fn make_encoder(
+ array: &'a dyn Array,
+ field: &Field,
+ plan: PlanRef<'_>,
+ ) -> Result<Self, ArrowError> {
+ let nulls = array.nulls().cloned();
+ let has_nulls = array.null_count() > 0;
+ 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(
+ "Avro writer: {other:?} not yet supported".into(),
+ ));
+ }
+ };
+ Ok(Self {
+ encoder,
+ nulls,
+ has_nulls,
+ nullability: None,
+ pre: None,
+ })
+ }
+
+ #[inline]
Review Comment:
Removing the inlines caused about a ~4% decrease in throughput performance.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org