1use crate::codec::{AvroDataType, AvroField, Codec};
21use crate::errors::AvroError;
22use crate::schema::{Fingerprint, Nullability, Prefix};
23use arrow_array::Float16Array;
24use arrow_array::cast::AsArray;
25use arrow_array::types::{
26 ArrowPrimitiveType, Date32Type, Date64Type, DurationMicrosecondType, DurationMillisecondType,
27 DurationNanosecondType, DurationSecondType, Float16Type, Float32Type, Float64Type, Int8Type,
28 Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType,
29 IntervalYearMonthType, RunEndIndexType, Time32MillisecondType, Time32SecondType,
30 Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType,
31 TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type,
32 UInt32Type, UInt64Type,
33};
34use arrow_array::{
35 Array, BinaryViewArray, Decimal128Array, Decimal256Array, DictionaryArray,
36 FixedSizeBinaryArray, FixedSizeListArray, GenericBinaryArray, GenericListArray,
37 GenericListViewArray, GenericStringArray, LargeListArray, LargeListViewArray, ListArray,
38 ListViewArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, RunArray, StringArray,
39 StringViewArray, StructArray, UnionArray,
40};
41#[cfg(feature = "small_decimals")]
42use arrow_array::{Decimal32Array, Decimal64Array};
43use arrow_buffer::{ArrowNativeType, NullBuffer};
44use arrow_schema::{DataType, Field, IntervalUnit, Schema as ArrowSchema, TimeUnit, UnionMode};
45use bytes::{BufMut, BytesMut};
46use std::io::Write;
47use std::sync::Arc;
48use uuid::Uuid;
49
50macro_rules! for_rows_with_prefix {
51 ($n:expr, $prefix:expr, $out:ident, |$row:ident| $body:block) => {{
52 match $prefix {
53 Some(prefix) => {
54 for $row in 0..$n {
55 $out.write_all(prefix)
56 .map_err(|e| AvroError::IoError(format!("write prefix: {e}"), e))?;
57 $body
58 }
59 }
60 None => {
61 for $row in 0..$n {
62 $body
63 }
64 }
65 }
66 }};
67}
68
69#[inline]
73pub(crate) fn write_long<W: Write + ?Sized>(out: &mut W, value: i64) -> Result<(), AvroError> {
74 let mut zz = ((value << 1) ^ (value >> 63)) as u64;
75 let mut buf = [0u8; 10];
77 let mut i = 0;
78 while (zz & !0x7F) != 0 {
79 buf[i] = ((zz & 0x7F) as u8) | 0x80;
80 i += 1;
81 zz >>= 7;
82 }
83 buf[i] = (zz & 0x7F) as u8;
84 i += 1;
85 out.write_all(&buf[..i])
86 .map_err(|e| AvroError::IoError(format!("write long: {e}"), e))
87}
88
89#[inline]
90fn write_int<W: Write + ?Sized>(out: &mut W, value: i32) -> Result<(), AvroError> {
91 write_long(out, value as i64)
92}
93
94#[inline]
95fn write_len_prefixed<W: Write + ?Sized>(out: &mut W, bytes: &[u8]) -> Result<(), AvroError> {
96 write_long(out, bytes.len() as i64)?;
97 out.write_all(bytes)
98 .map_err(|e| AvroError::IoError(format!("write bytes: {e}"), e))
99}
100
101#[inline]
102fn write_bool<W: Write + ?Sized>(out: &mut W, v: bool) -> Result<(), AvroError> {
103 out.write_all(&[u8::from(v)])
104 .map_err(|e| AvroError::IoError(format!("write bool: {e}"), e))
105}
106
107#[inline]
116fn minimal_twos_complement(be: &[u8]) -> &[u8] {
117 if be.is_empty() {
118 return be;
119 }
120 let sign_byte = if (be[0] & 0x80) != 0 { 0xFF } else { 0x00 };
121 let mut k = 0usize;
122 while k < be.len() && be[k] == sign_byte {
123 k += 1;
124 }
125 if k == 0 {
126 return be;
127 }
128 if k == be.len() {
129 return &be[be.len() - 1..];
130 }
131 let drop = if ((be[k] ^ sign_byte) & 0x80) == 0 {
132 k
133 } else {
134 k - 1
135 };
136 &be[drop..]
137}
138
139#[inline]
151fn write_sign_extended<W: Write + ?Sized>(
152 out: &mut W,
153 src_be: &[u8],
154 n: usize,
155) -> Result<(), AvroError> {
156 let len = src_be.len();
157 if len == n {
158 out.write_all(src_be)?;
159 return Ok(());
160 }
161 let sign_byte = if len > 0 && (src_be[0] & 0x80) != 0 {
162 0xFF
163 } else {
164 0x00
165 };
166 if len > n {
167 let extra = len - n;
168 if n == 0 && src_be.iter().all(|&b| b == sign_byte) {
169 return Ok(());
170 }
171 if src_be[..extra].iter().any(|&b| b != sign_byte)
174 || ((src_be[extra] ^ sign_byte) & 0x80) != 0
175 {
176 return Err(AvroError::InvalidArgument(format!(
177 "Decimal value with {len} bytes cannot be represented in {n} bytes without overflow",
178 )));
179 }
180 return out
181 .write_all(&src_be[extra..])
182 .map_err(|e| AvroError::IoError(format!("write decimal fixed: {e}"), e));
183 }
184 let pad_len = n - len;
186 const ZPAD: [u8; 64] = [0x00; 64];
188 const FPAD: [u8; 64] = [0xFF; 64];
189 let pad = if sign_byte == 0x00 {
190 &ZPAD[..]
191 } else {
192 &FPAD[..]
193 };
194 let mut rem = pad_len;
197 while rem >= pad.len() {
198 out.write_all(pad)
199 .map_err(|e| AvroError::IoError(format!("write decimal fixed: {e}"), e))?;
200 rem -= pad.len();
201 }
202 if rem > 0 {
203 out.write_all(&pad[..rem])
204 .map_err(|e| AvroError::IoError(format!("write decimal fixed: {e}"), e))?;
205 }
206 out.write_all(src_be)
207 .map_err(|e| AvroError::IoError(format!("write decimal fixed: {e}"), e))
208}
209
210fn write_optional_index<W: Write + ?Sized>(
216 out: &mut W,
217 is_null: bool,
218 null_order: Nullability,
219) -> Result<(), AvroError> {
220 let byte = union_value_branch_byte(null_order, is_null);
221 out.write_all(&[byte])
222 .map_err(|e| AvroError::IoError(format!("write union branch: {e}"), e))
223}
224
225#[derive(Debug, Clone)]
226enum NullState<'a> {
227 NonNullable,
228 NullableNoNulls {
229 union_value_byte: u8,
230 },
231 Nullable {
232 nulls: &'a NullBuffer,
233 null_order: Nullability,
234 },
235}
236
237pub(crate) struct FieldEncoder<'a> {
241 encoder: Encoder<'a>,
242 null_state: NullState<'a>,
243}
244
245impl<'a> FieldEncoder<'a> {
246 fn make_encoder(
247 array: &'a dyn Array,
248 plan: &FieldPlan,
249 nullability: Option<Nullability>,
250 ) -> Result<Self, AvroError> {
251 let encoder = match plan {
252 FieldPlan::Scalar => match array.data_type() {
253 DataType::Null => Encoder::Null,
254 DataType::Boolean => Encoder::Boolean(BooleanEncoder(array.as_boolean())),
255 DataType::Utf8 => {
256 Encoder::Utf8(Utf8GenericEncoder::<i32>(array.as_string::<i32>()))
257 }
258 DataType::LargeUtf8 => {
259 Encoder::Utf8Large(Utf8GenericEncoder::<i64>(array.as_string::<i64>()))
260 }
261 DataType::Utf8View => {
262 let arr = array
263 .as_any()
264 .downcast_ref::<StringViewArray>()
265 .ok_or_else(|| AvroError::SchemaError("Expected StringViewArray".into()))?;
266 Encoder::Utf8View(Utf8ViewEncoder(arr))
267 }
268 DataType::BinaryView => {
269 let arr = array
270 .as_any()
271 .downcast_ref::<BinaryViewArray>()
272 .ok_or_else(|| AvroError::SchemaError("Expected BinaryViewArray".into()))?;
273 Encoder::BinaryView(BinaryViewEncoder(arr))
274 }
275 DataType::Int32 => Encoder::Int(IntEncoder(array.as_primitive::<Int32Type>())),
276 DataType::Int64 => Encoder::Long(LongEncoder(array.as_primitive::<Int64Type>())),
277 #[cfg(feature = "avro_custom_types")]
278 DataType::Int8 => Encoder::Int8(Int8Encoder(array.as_primitive::<Int8Type>())),
279 #[cfg(not(feature = "avro_custom_types"))]
280 DataType::Int8 => {
281 Encoder::Int8ToInt(Int8ToIntEncoder(array.as_primitive::<Int8Type>()))
282 }
283 #[cfg(feature = "avro_custom_types")]
284 DataType::Int16 => Encoder::Int16(Int16Encoder(array.as_primitive::<Int16Type>())),
285 #[cfg(not(feature = "avro_custom_types"))]
286 DataType::Int16 => {
287 Encoder::Int16ToInt(Int16ToIntEncoder(array.as_primitive::<Int16Type>()))
288 }
289 #[cfg(feature = "avro_custom_types")]
290 DataType::UInt8 => Encoder::UInt8(UInt8Encoder(array.as_primitive::<UInt8Type>())),
291 #[cfg(not(feature = "avro_custom_types"))]
292 DataType::UInt8 => {
293 Encoder::UInt8ToInt(UInt8ToIntEncoder(array.as_primitive::<UInt8Type>()))
294 }
295 #[cfg(feature = "avro_custom_types")]
296 DataType::UInt16 => {
297 Encoder::UInt16(UInt16Encoder(array.as_primitive::<UInt16Type>()))
298 }
299 #[cfg(not(feature = "avro_custom_types"))]
300 DataType::UInt16 => {
301 Encoder::UInt16ToInt(UInt16ToIntEncoder(array.as_primitive::<UInt16Type>()))
302 }
303 #[cfg(feature = "avro_custom_types")]
304 DataType::UInt32 => {
305 Encoder::UInt32(UInt32Encoder(array.as_primitive::<UInt32Type>()))
306 }
307 #[cfg(not(feature = "avro_custom_types"))]
308 DataType::UInt32 => {
309 Encoder::UInt32ToLong(UInt32ToLongEncoder(array.as_primitive::<UInt32Type>()))
310 }
311 #[cfg(feature = "avro_custom_types")]
312 DataType::UInt64 => {
313 Encoder::UInt64Fixed(UInt64FixedEncoder(array.as_primitive::<UInt64Type>()))
314 }
315 #[cfg(not(feature = "avro_custom_types"))]
316 DataType::UInt64 => {
317 Encoder::UInt64ToLong(UInt64ToLongEncoder(array.as_primitive::<UInt64Type>()))
318 }
319 #[cfg(feature = "avro_custom_types")]
320 DataType::Float16 => {
321 Encoder::Float16Fixed(Float16FixedEncoder(array.as_primitive::<Float16Type>()))
322 }
323 #[cfg(not(feature = "avro_custom_types"))]
324 DataType::Float16 => Encoder::Float16ToFloat(Float16ToFloatEncoder(
325 array.as_primitive::<Float16Type>(),
326 )),
327 DataType::Date32 => Encoder::Date32(IntEncoder(array.as_primitive::<Date32Type>())),
328 #[cfg(not(feature = "avro_custom_types"))]
329 DataType::Date64 => {
330 Encoder::Date64ToLong(Date64ToLongEncoder(array.as_primitive::<Date64Type>()))
332 }
333 #[cfg(feature = "avro_custom_types")]
334 DataType::Date64 => {
335 Encoder::Date64(LongEncoder(array.as_primitive::<Date64Type>()))
336 }
337 #[cfg(feature = "avro_custom_types")]
338 DataType::Time32(TimeUnit::Second) => {
339 Encoder::Time32Secs(IntEncoder(array.as_primitive::<Time32SecondType>()))
340 }
341 DataType::Time32(TimeUnit::Millisecond) => {
342 Encoder::Time32Millis(IntEncoder(array.as_primitive::<Time32MillisecondType>()))
343 }
344 DataType::Time32(TimeUnit::Microsecond) => {
345 return Err(AvroError::InvalidArgument(
346 "Arrow Time32 only supports Second or Millisecond. Use Time64 for microseconds."
347 .into(),
348 ));
349 }
350 DataType::Time32(TimeUnit::Nanosecond) => {
351 return Err(AvroError::InvalidArgument(
352 "Arrow Time32 only supports Second or Millisecond. Use Time64 for nanoseconds."
353 .into(),
354 ));
355 }
356 DataType::Time64(TimeUnit::Microsecond) => Encoder::Time64Micros(LongEncoder(
357 array.as_primitive::<Time64MicrosecondType>(),
358 )),
359 #[cfg(not(feature = "avro_custom_types"))]
360 DataType::Time64(TimeUnit::Nanosecond) => {
361 Encoder::Time64NanosToMicros(Time64NanosToMicrosEncoder(
363 array.as_primitive::<Time64NanosecondType>(),
364 ))
365 }
366 #[cfg(feature = "avro_custom_types")]
367 DataType::Time64(TimeUnit::Nanosecond) => {
368 Encoder::Time64Nanos(LongEncoder(array.as_primitive::<Time64NanosecondType>()))
369 }
370 DataType::Time64(TimeUnit::Millisecond) => {
371 return Err(AvroError::InvalidArgument(
372 "Arrow Time64 with millisecond unit is not a valid Arrow type (use Time32 for millis)."
373 .into(),
374 ));
375 }
376 DataType::Time64(TimeUnit::Second) => {
377 return Err(AvroError::InvalidArgument(
378 "Arrow Time64 with second unit is not a valid Arrow type (use Time32 for seconds)."
379 .into(),
380 ));
381 }
382 DataType::Float32 => {
383 Encoder::Float32(F32Encoder(array.as_primitive::<Float32Type>()))
384 }
385 DataType::Float64 => {
386 Encoder::Float64(F64Encoder(array.as_primitive::<Float64Type>()))
387 }
388 DataType::Binary => Encoder::Binary(BinaryEncoder(array.as_binary::<i32>())),
389 DataType::LargeBinary => {
390 Encoder::LargeBinary(BinaryEncoder(array.as_binary::<i64>()))
391 }
392 DataType::FixedSizeBinary(_len) => {
393 let arr = array
394 .as_any()
395 .downcast_ref::<FixedSizeBinaryArray>()
396 .ok_or_else(|| {
397 AvroError::SchemaError("Expected FixedSizeBinaryArray".into())
398 })?;
399 Encoder::Fixed(FixedEncoder(arr))
400 }
401 DataType::Timestamp(unit, _) => match unit {
402 TimeUnit::Second => {
403 #[cfg(not(feature = "avro_custom_types"))]
404 {
405 Encoder::TimestampSecsToMillis(TimestampSecondsToMillisEncoder(
406 array.as_primitive::<TimestampSecondType>(),
407 ))
408 }
409 #[cfg(feature = "avro_custom_types")]
410 {
411 Encoder::TimestampSecs(LongEncoder(
412 array.as_primitive::<TimestampSecondType>(),
413 ))
414 }
415 }
416 TimeUnit::Millisecond => Encoder::TimestampMillis(LongEncoder(
417 array.as_primitive::<TimestampMillisecondType>(),
418 )),
419 TimeUnit::Microsecond => Encoder::TimestampMicros(LongEncoder(
420 array.as_primitive::<TimestampMicrosecondType>(),
421 )),
422 TimeUnit::Nanosecond => Encoder::TimestampNanos(LongEncoder(
423 array.as_primitive::<TimestampNanosecondType>(),
424 )),
425 },
426 #[cfg(feature = "avro_custom_types")]
427 DataType::Interval(unit) => match unit {
428 IntervalUnit::MonthDayNano => {
429 Encoder::IntervalMonthDayNanoFixed(IntervalMonthDayNanoFixedEncoder(
430 array.as_primitive::<IntervalMonthDayNanoType>(),
431 ))
432 }
433 IntervalUnit::YearMonth => {
434 Encoder::IntervalYearMonthFixed(IntervalYearMonthFixedEncoder(
435 array.as_primitive::<IntervalYearMonthType>(),
436 ))
437 }
438 IntervalUnit::DayTime => Encoder::IntervalDayTimeFixed(
439 IntervalDayTimeFixedEncoder(array.as_primitive::<IntervalDayTimeType>()),
440 ),
441 },
442 DataType::Duration(tu) => match tu {
443 TimeUnit::Second => Encoder::DurationSeconds(LongEncoder(
444 array.as_primitive::<DurationSecondType>(),
445 )),
446 TimeUnit::Millisecond => Encoder::DurationMillis(LongEncoder(
447 array.as_primitive::<DurationMillisecondType>(),
448 )),
449 TimeUnit::Microsecond => Encoder::DurationMicros(LongEncoder(
450 array.as_primitive::<DurationMicrosecondType>(),
451 )),
452 TimeUnit::Nanosecond => Encoder::DurationNanos(LongEncoder(
453 array.as_primitive::<DurationNanosecondType>(),
454 )),
455 },
456 other => {
457 return Err(AvroError::NYI(format!(
458 "Avro scalar type not yet supported: {other:?}"
459 )));
460 }
461 },
462 FieldPlan::Struct { bindings } => {
463 let arr = array
464 .as_any()
465 .downcast_ref::<StructArray>()
466 .ok_or_else(|| AvroError::SchemaError("Expected StructArray".into()))?;
467 Encoder::Struct(Box::new(StructEncoder::try_new(arr, bindings)?))
468 }
469 FieldPlan::List {
470 items_nullability,
471 item_plan,
472 } => match array.data_type() {
473 DataType::List(_) => {
474 let arr = array
475 .as_any()
476 .downcast_ref::<ListArray>()
477 .ok_or_else(|| AvroError::SchemaError("Expected ListArray".into()))?;
478 Encoder::List(Box::new(ListEncoder32::try_new(
479 arr,
480 *items_nullability,
481 item_plan.as_ref(),
482 )?))
483 }
484 DataType::LargeList(_) => {
485 let arr = array
486 .as_any()
487 .downcast_ref::<LargeListArray>()
488 .ok_or_else(|| AvroError::SchemaError("Expected LargeListArray".into()))?;
489 Encoder::LargeList(Box::new(ListEncoder64::try_new(
490 arr,
491 *items_nullability,
492 item_plan.as_ref(),
493 )?))
494 }
495 DataType::ListView(_) => {
496 let arr = array
497 .as_any()
498 .downcast_ref::<ListViewArray>()
499 .ok_or_else(|| AvroError::SchemaError("Expected ListViewArray".into()))?;
500 Encoder::ListView(Box::new(ListViewEncoder32::try_new(
501 arr,
502 *items_nullability,
503 item_plan.as_ref(),
504 )?))
505 }
506 DataType::LargeListView(_) => {
507 let arr = array
508 .as_any()
509 .downcast_ref::<LargeListViewArray>()
510 .ok_or_else(|| {
511 AvroError::SchemaError("Expected LargeListViewArray".into())
512 })?;
513 Encoder::LargeListView(Box::new(ListViewEncoder64::try_new(
514 arr,
515 *items_nullability,
516 item_plan.as_ref(),
517 )?))
518 }
519 DataType::FixedSizeList(_, _) => {
520 let arr = array
521 .as_any()
522 .downcast_ref::<FixedSizeListArray>()
523 .ok_or_else(|| {
524 AvroError::SchemaError("Expected FixedSizeListArray".into())
525 })?;
526 Encoder::FixedSizeList(Box::new(FixedSizeListEncoder::try_new(
527 arr,
528 *items_nullability,
529 item_plan.as_ref(),
530 )?))
531 }
532 other => {
533 return Err(AvroError::SchemaError(format!(
534 "Avro array site requires Arrow List/LargeList/ListView/LargeListView/FixedSizeList, found: {other:?}"
535 )));
536 }
537 },
538 FieldPlan::Decimal { size } => match array.data_type() {
539 #[cfg(feature = "small_decimals")]
540 DataType::Decimal32(_, _) => {
541 let arr = array
542 .as_any()
543 .downcast_ref::<Decimal32Array>()
544 .ok_or_else(|| AvroError::SchemaError("Expected Decimal32Array".into()))?;
545 Encoder::Decimal32(DecimalEncoder::<4, Decimal32Array>::new(arr, *size))
546 }
547 #[cfg(feature = "small_decimals")]
548 DataType::Decimal64(_, _) => {
549 let arr = array
550 .as_any()
551 .downcast_ref::<Decimal64Array>()
552 .ok_or_else(|| AvroError::SchemaError("Expected Decimal64Array".into()))?;
553 Encoder::Decimal64(DecimalEncoder::<8, Decimal64Array>::new(arr, *size))
554 }
555 DataType::Decimal128(_, _) => {
556 let arr = array
557 .as_any()
558 .downcast_ref::<Decimal128Array>()
559 .ok_or_else(|| AvroError::SchemaError("Expected Decimal128Array".into()))?;
560 Encoder::Decimal128(DecimalEncoder::<16, Decimal128Array>::new(arr, *size))
561 }
562 DataType::Decimal256(_, _) => {
563 let arr = array
564 .as_any()
565 .downcast_ref::<Decimal256Array>()
566 .ok_or_else(|| AvroError::SchemaError("Expected Decimal256Array".into()))?;
567 Encoder::Decimal256(DecimalEncoder::<32, Decimal256Array>::new(arr, *size))
568 }
569 other => {
570 return Err(AvroError::SchemaError(format!(
571 "Avro decimal site requires Arrow Decimal 32, 64, 128, or 256, found: {other:?}"
572 )));
573 }
574 },
575 FieldPlan::Uuid => {
576 let arr = array
577 .as_any()
578 .downcast_ref::<FixedSizeBinaryArray>()
579 .ok_or_else(|| {
580 AvroError::SchemaError("Expected FixedSizeBinaryArray".into())
581 })?;
582 Encoder::Uuid(UuidEncoder(arr))
583 }
584 FieldPlan::Map {
585 values_nullability,
586 value_plan,
587 } => {
588 let arr = array
589 .as_any()
590 .downcast_ref::<MapArray>()
591 .ok_or_else(|| AvroError::SchemaError("Expected MapArray".into()))?;
592 Encoder::Map(Box::new(MapEncoder::try_new(
593 arr,
594 *values_nullability,
595 value_plan.as_ref(),
596 )?))
597 }
598 FieldPlan::Enum { symbols } => match array.data_type() {
599 DataType::Dictionary(key_dt, value_dt) => {
600 if **key_dt != DataType::Int32 || **value_dt != DataType::Utf8 {
601 return Err(AvroError::SchemaError(
602 "Avro enum requires Dictionary<Int32, Utf8>".into(),
603 ));
604 }
605 let dict = array
606 .as_any()
607 .downcast_ref::<DictionaryArray<Int32Type>>()
608 .ok_or_else(|| {
609 AvroError::SchemaError("Expected DictionaryArray<Int32>".into())
610 })?;
611 let values = dict
612 .values()
613 .as_any()
614 .downcast_ref::<StringArray>()
615 .ok_or_else(|| {
616 AvroError::SchemaError("Dictionary values must be Utf8".into())
617 })?;
618 if values.len() != symbols.len() {
619 return Err(AvroError::SchemaError(format!(
620 "Enum symbol length {} != dictionary size {}",
621 symbols.len(),
622 values.len()
623 )));
624 }
625 for i in 0..values.len() {
626 if values.value(i) != symbols[i].as_str() {
627 return Err(AvroError::SchemaError(format!(
628 "Enum symbol mismatch at {i}: schema='{}' dict='{}'",
629 symbols[i],
630 values.value(i)
631 )));
632 }
633 }
634 let keys = dict.keys();
635 Encoder::Enum(EnumEncoder { keys })
636 }
637 other => {
638 return Err(AvroError::SchemaError(format!(
639 "Avro enum site requires DataType::Dictionary, found: {other:?}"
640 )));
641 }
642 },
643 FieldPlan::Union { bindings } => {
644 let arr = array
645 .as_any()
646 .downcast_ref::<UnionArray>()
647 .ok_or_else(|| AvroError::SchemaError("Expected UnionArray".into()))?;
648
649 Encoder::Union(Box::new(UnionEncoder::try_new(arr, bindings)?))
650 }
651 FieldPlan::RunEndEncoded {
652 values_nullability,
653 value_plan,
654 } => {
655 let build = |run_arr_any: &'a dyn Array| -> Result<Encoder<'a>, AvroError> {
657 if let Some(arr) = run_arr_any.as_any().downcast_ref::<RunArray<Int16Type>>() {
658 return Ok(Encoder::RunEncoded16(Box::new(RunEncodedEncoder::<
659 Int16Type,
660 >::new(
661 arr,
662 FieldEncoder::make_encoder(
663 arr.values().as_ref(),
664 value_plan.as_ref(),
665 *values_nullability,
666 )?,
667 ))));
668 }
669 if let Some(arr) = run_arr_any.as_any().downcast_ref::<RunArray<Int32Type>>() {
670 return Ok(Encoder::RunEncoded32(Box::new(RunEncodedEncoder::<
671 Int32Type,
672 >::new(
673 arr,
674 FieldEncoder::make_encoder(
675 arr.values().as_ref(),
676 value_plan.as_ref(),
677 *values_nullability,
678 )?,
679 ))));
680 }
681 if let Some(arr) = run_arr_any.as_any().downcast_ref::<RunArray<Int64Type>>() {
682 return Ok(Encoder::RunEncoded64(Box::new(RunEncodedEncoder::<
683 Int64Type,
684 >::new(
685 arr,
686 FieldEncoder::make_encoder(
687 arr.values().as_ref(),
688 value_plan.as_ref(),
689 *values_nullability,
690 )?,
691 ))));
692 }
693 Err(AvroError::SchemaError(
694 "Unsupported run-ends index type for RunEndEncoded; expected Int16/Int32/Int64"
695 .into(),
696 ))
697 };
698 build(array)?
699 }
700 FieldPlan::Duration => match array.data_type() {
701 DataType::Interval(IntervalUnit::MonthDayNano) => {
702 Encoder::IntervalMonthDayNanoDuration(DurationEncoder(
703 array.as_primitive::<IntervalMonthDayNanoType>(),
704 ))
705 }
706 DataType::Interval(IntervalUnit::YearMonth) => Encoder::IntervalYearMonthDuration(
707 DurationEncoder(array.as_primitive::<IntervalYearMonthType>()),
708 ),
709 DataType::Interval(IntervalUnit::DayTime) => Encoder::IntervalDayTimeDuration(
710 DurationEncoder(array.as_primitive::<IntervalDayTimeType>()),
711 ),
712 other => {
713 return Err(AvroError::SchemaError(format!(
714 "Avro duration requires Arrow Interval type, found: {other:?}"
715 )));
716 }
717 },
718 FieldPlan::TimeMillisFromSecs => match array.data_type() {
719 DataType::Time32(TimeUnit::Second) => Encoder::Time32SecsToMillis(
720 Time32SecondsToMillisEncoder(array.as_primitive::<Time32SecondType>()),
721 ),
722 other => {
723 return Err(AvroError::SchemaError(format!(
724 "Avro time-millis-from-seconds requires Arrow Time32(Second), found: {other:?}"
725 )));
726 }
727 },
728 };
729 let null_state = match nullability {
731 None => NullState::NonNullable,
732 Some(null_order) => match array.nulls() {
733 Some(nulls) if array.null_count() > 0 => NullState::Nullable { nulls, null_order },
734 _ => NullState::NullableNoNulls {
735 union_value_byte: union_value_branch_byte(null_order, false),
737 },
738 },
739 };
740 Ok(Self {
741 encoder,
742 null_state,
743 })
744 }
745
746 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
747 match &self.null_state {
748 NullState::NonNullable => {}
749 NullState::NullableNoNulls { union_value_byte } => out
750 .write_all(&[*union_value_byte])
751 .map_err(|e| AvroError::IoError(format!("write union value branch: {e}"), e))?,
752 NullState::Nullable { nulls, null_order } if nulls.is_null(idx) => {
753 return write_optional_index(out, true, *null_order); }
755 NullState::Nullable { null_order, .. } => {
756 write_optional_index(out, false, *null_order)?;
757 }
758 }
759 self.encoder.encode(out, idx)
760 }
761}
762
763fn union_value_branch_byte(null_order: Nullability, is_null: bool) -> u8 {
764 let nulls_first = null_order == Nullability::default();
765 if nulls_first == is_null { 0x00 } else { 0x02 }
766}
767
768#[derive(Debug, Clone)]
771enum FieldPlan {
772 Scalar,
774 Struct { bindings: Vec<FieldBinding> },
776 List {
778 items_nullability: Option<Nullability>,
779 item_plan: Box<FieldPlan>,
780 },
781 Decimal { size: Option<usize> },
783 Uuid,
785 Map {
787 values_nullability: Option<Nullability>,
788 value_plan: Box<FieldPlan>,
789 },
790 Enum { symbols: Arc<[String]> },
793 Union { bindings: Vec<FieldBinding> },
795 RunEndEncoded {
798 values_nullability: Option<Nullability>,
799 value_plan: Box<FieldPlan>,
800 },
801 Duration,
805 TimeMillisFromSecs,
809}
810
811#[derive(Debug, Clone)]
812struct FieldBinding {
813 arrow_index: usize,
815 nullability: Option<Nullability>,
817 plan: FieldPlan,
819}
820
821#[derive(Debug)]
823pub(crate) struct RecordEncoderBuilder<'a> {
824 avro_root: &'a AvroField,
825 arrow_schema: &'a ArrowSchema,
826 fingerprint: Option<Fingerprint>,
827}
828
829impl<'a> RecordEncoderBuilder<'a> {
830 pub(crate) fn new(avro_root: &'a AvroField, arrow_schema: &'a ArrowSchema) -> Self {
832 Self {
833 avro_root,
834 arrow_schema,
835 fingerprint: None,
836 }
837 }
838
839 pub(crate) fn with_fingerprint(mut self, fingerprint: Option<Fingerprint>) -> Self {
840 self.fingerprint = fingerprint;
841 self
842 }
843
844 pub(crate) fn build(self) -> Result<RecordEncoder, AvroError> {
847 let avro_root_dt = self.avro_root.data_type();
848 let Codec::Struct(root_fields) = avro_root_dt.codec() else {
849 return Err(AvroError::SchemaError(
850 "Top-level Avro schema must be a record/struct".into(),
851 ));
852 };
853 let mut columns = Vec::with_capacity(root_fields.len());
854 for root_field in root_fields.as_ref() {
855 let name = root_field.name();
856 let arrow_index = self.arrow_schema.index_of(name).map_err(|e| {
857 AvroError::SchemaError(format!("Schema mismatch for field '{name}': {e}"))
858 })?;
859 columns.push(FieldBinding {
860 arrow_index,
861 nullability: root_field.data_type().nullability(),
862 plan: FieldPlan::build(
863 root_field.data_type(),
864 self.arrow_schema.field(arrow_index),
865 )?,
866 });
867 }
868 Ok(RecordEncoder {
869 columns,
870 prefix: self.fingerprint.map(|fp| fp.make_prefix()),
871 })
872 }
873}
874
875#[derive(Debug, Clone)]
881pub(crate) struct RecordEncoder {
882 columns: Vec<FieldBinding>,
883 prefix: Option<Prefix>,
885}
886
887impl RecordEncoder {
888 fn prepare_for_batch<'a>(
889 &'a self,
890 batch: &'a RecordBatch,
891 ) -> Result<Vec<FieldEncoder<'a>>, AvroError> {
892 let arrays = batch.columns();
893 let mut out = Vec::with_capacity(self.columns.len());
894 for col_plan in &self.columns {
895 let arrow_index = col_plan.arrow_index;
896 let array = arrays.get(arrow_index).ok_or_else(|| {
897 AvroError::SchemaError(format!("Column index {arrow_index} out of range"))
898 })?;
899 #[cfg(not(feature = "avro_custom_types"))]
900 let site_nullability = match &col_plan.plan {
901 FieldPlan::RunEndEncoded { .. } => None,
902 _ => col_plan.nullability,
903 };
904 #[cfg(feature = "avro_custom_types")]
905 let site_nullability = col_plan.nullability;
906 out.push(FieldEncoder::make_encoder(
907 array.as_ref(),
908 &col_plan.plan,
909 site_nullability,
910 )?);
911 }
912 Ok(out)
913 }
914
915 pub(crate) fn encode<W: Write>(
919 &self,
920 out: &mut W,
921 batch: &RecordBatch,
922 ) -> Result<(), AvroError> {
923 let mut column_encoders = self.prepare_for_batch(batch)?;
924 let n = batch.num_rows();
925 let prefix = self.prefix.as_ref().map(|p| p.as_slice());
926 for_rows_with_prefix!(n, prefix, out, |row| {
927 for enc in &mut column_encoders {
928 enc.encode(out, row)?;
929 }
930 });
931 Ok(())
932 }
933
934 pub(crate) fn encode_rows(
943 &self,
944 batch: &RecordBatch,
945 row_capacity: usize,
946 out: &mut BytesMut,
947 offsets: &mut Vec<usize>,
948 ) -> Result<(), AvroError> {
949 let out_len = out.len();
950 if offsets.first() != Some(&0) || offsets.last() != Some(&out_len) {
951 return Err(AvroError::General(
952 "encode_rows requires offsets to start with 0 and end at out.len()".to_string(),
953 ));
954 }
955 let n = batch.num_rows();
956 if n == 0 {
957 return Ok(());
958 }
959 if offsets.len().checked_add(n).is_none() {
960 return Err(AvroError::General(
961 "encode_rows cannot append offsets: too many rows".to_string(),
962 ));
963 }
964 let mut column_encoders = self.prepare_for_batch(batch)?;
965 offsets.reserve(n);
966 let prefix_bytes = self.prefix.as_ref().map(|p| p.as_slice());
967 let prefix_len = prefix_bytes.map_or(0, |p| p.len());
968 let per_row_hint = row_capacity.max(prefix_len);
969 if let Some(additional) = n
970 .checked_mul(per_row_hint)
971 .filter(|&a| out_len.checked_add(a).is_some())
972 {
973 out.reserve(additional);
974 }
975 let start_out_len = out.len();
976 let start_offsets_len = offsets.len();
977 let res = (|| -> Result<(), AvroError> {
978 let mut w = out.writer();
979 if let [enc0] = column_encoders.as_mut_slice() {
980 for_rows_with_prefix!(n, prefix_bytes, w, |row| {
981 enc0.encode(&mut w, row)?;
982 offsets.push(w.get_ref().len());
983 });
984 } else {
985 for_rows_with_prefix!(n, prefix_bytes, w, |row| {
986 for enc in &mut column_encoders {
987 enc.encode(&mut w, row)?;
988 }
989 offsets.push(w.get_ref().len());
990 });
991 }
992 Ok(())
993 })();
994 if res.is_err() {
995 out.truncate(start_out_len);
996 offsets.truncate(start_offsets_len);
997 } else {
998 debug_assert_eq!(
999 *offsets.last().unwrap(),
1000 out.len(),
1001 "encode_rows: offsets/out length mismatch after successful encode"
1002 );
1003 }
1004 res
1005 }
1006}
1007
1008fn find_struct_child_index(fields: &arrow_schema::Fields, name: &str) -> Option<usize> {
1009 fields.iter().position(|f| f.name() == name)
1010}
1011
1012fn find_map_value_field_index(fields: &arrow_schema::Fields) -> Option<usize> {
1013 find_struct_child_index(fields, Field::MAP_VALUE_FIELD_DEFAULT_NAME)
1015 .or_else(|| find_struct_child_index(fields, "values"))
1016 .or_else(|| if fields.len() == 2 { Some(1) } else { None })
1017}
1018
1019impl FieldPlan {
1020 fn build(avro_dt: &AvroDataType, arrow_field: &Field) -> Result<Self, AvroError> {
1021 #[cfg(not(feature = "avro_custom_types"))]
1022 if let DataType::RunEndEncoded(_re_field, values_field) = arrow_field.data_type() {
1023 let values_nullability = avro_dt.nullability();
1024 let value_site_dt: &AvroDataType = match avro_dt.codec() {
1025 Codec::Union(branches, _, _) => branches
1026 .iter()
1027 .find(|b| !matches!(b.codec(), Codec::Null))
1028 .ok_or_else(|| {
1029 AvroError::SchemaError(
1030 "Avro union at RunEndEncoded site has no non-null branch".into(),
1031 )
1032 })?,
1033 _ => avro_dt,
1034 };
1035 return Ok(FieldPlan::RunEndEncoded {
1036 values_nullability,
1037 value_plan: Box::new(FieldPlan::build(value_site_dt, values_field.as_ref())?),
1038 });
1039 }
1040 if let DataType::FixedSizeBinary(len) = arrow_field.data_type() {
1041 let ext_is_uuid = {
1043 #[cfg(feature = "canonical_extension_types")]
1044 {
1045 matches!(
1046 arrow_field.extension_type_name(),
1047 Some("arrow.uuid" | "uuid")
1048 )
1049 }
1050 #[cfg(not(feature = "canonical_extension_types"))]
1051 {
1052 false
1053 }
1054 };
1055 let md_is_uuid = arrow_field
1056 .metadata()
1057 .get("logicalType")
1058 .map(|s| s.as_str())
1059 == Some("uuid");
1060 if ext_is_uuid || md_is_uuid {
1061 if *len != 16 {
1062 return Err(AvroError::InvalidArgument(
1063 "logicalType=uuid requires FixedSizeBinary(16)".into(),
1064 ));
1065 }
1066 return Ok(FieldPlan::Uuid);
1067 }
1068 }
1069 match avro_dt.codec() {
1070 Codec::Struct(avro_fields) => {
1071 let fields = match arrow_field.data_type() {
1072 DataType::Struct(struct_fields) => struct_fields,
1073 other => {
1074 return Err(AvroError::SchemaError(format!(
1075 "Avro struct maps to Arrow Struct, found: {other:?}"
1076 )));
1077 }
1078 };
1079 let mut bindings = Vec::with_capacity(avro_fields.len());
1080 for avro_field in avro_fields.iter() {
1081 let name = avro_field.name().to_string();
1082 let idx = find_struct_child_index(fields, &name).ok_or_else(|| {
1083 AvroError::SchemaError(format!(
1084 "Struct field '{name}' not present in Arrow field '{}'",
1085 arrow_field.name()
1086 ))
1087 })?;
1088 bindings.push(FieldBinding {
1089 arrow_index: idx,
1090 nullability: avro_field.data_type().nullability(),
1091 plan: FieldPlan::build(avro_field.data_type(), fields[idx].as_ref())?,
1092 });
1093 }
1094 Ok(FieldPlan::Struct { bindings })
1095 }
1096 Codec::List(items_dt) => match arrow_field.data_type() {
1097 DataType::List(field_ref)
1098 | DataType::LargeList(field_ref)
1099 | DataType::ListView(field_ref)
1100 | DataType::LargeListView(field_ref) => Ok(FieldPlan::List {
1101 items_nullability: items_dt.nullability(),
1102 item_plan: Box::new(FieldPlan::build(items_dt.as_ref(), field_ref.as_ref())?),
1103 }),
1104 DataType::FixedSizeList(field_ref, _len) => Ok(FieldPlan::List {
1105 items_nullability: items_dt.nullability(),
1106 item_plan: Box::new(FieldPlan::build(items_dt.as_ref(), field_ref.as_ref())?),
1107 }),
1108 other => Err(AvroError::SchemaError(format!(
1109 "Avro array maps to Arrow List/LargeList/ListView/LargeListView/FixedSizeList, found: {other:?}"
1110 ))),
1111 },
1112 Codec::Map(values_dt) => {
1113 let entries_field = match arrow_field.data_type() {
1114 DataType::Map(entries, _sorted) => entries.as_ref(),
1115 other => {
1116 return Err(AvroError::SchemaError(format!(
1117 "Avro map maps to Arrow DataType::Map, found: {other:?}"
1118 )));
1119 }
1120 };
1121 let entries_struct_fields = match entries_field.data_type() {
1122 DataType::Struct(fs) => fs,
1123 other => {
1124 return Err(AvroError::SchemaError(format!(
1125 "Arrow Map entries must be Struct, found: {other:?}"
1126 )));
1127 }
1128 };
1129 let value_idx =
1130 find_map_value_field_index(entries_struct_fields).ok_or_else(|| {
1131 AvroError::SchemaError("Map entries struct missing value field".into())
1132 })?;
1133 let value_field = entries_struct_fields[value_idx].as_ref();
1134 let value_plan = FieldPlan::build(values_dt.as_ref(), value_field)?;
1135 Ok(FieldPlan::Map {
1136 values_nullability: values_dt.nullability(),
1137 value_plan: Box::new(value_plan),
1138 })
1139 }
1140 Codec::Enum(symbols) => match arrow_field.data_type() {
1141 DataType::Dictionary(key_dt, value_dt) => {
1142 if **key_dt != DataType::Int32 {
1143 return Err(AvroError::SchemaError(
1144 "Avro enum requires Dictionary<Int32, Utf8>".into(),
1145 ));
1146 }
1147 if **value_dt != DataType::Utf8 {
1148 return Err(AvroError::SchemaError(
1149 "Avro enum requires Dictionary<Int32, Utf8>".into(),
1150 ));
1151 }
1152 Ok(FieldPlan::Enum {
1153 symbols: symbols.clone(),
1154 })
1155 }
1156 other => Err(AvroError::SchemaError(format!(
1157 "Avro enum maps to Arrow Dictionary<Int32, Utf8>, found: {other:?}"
1158 ))),
1159 },
1160 Codec::Decimal(precision, scale_opt, fixed_size_opt) => {
1162 let (ap, as_) = match arrow_field.data_type() {
1163 #[cfg(feature = "small_decimals")]
1164 DataType::Decimal32(p, s) => (*p as usize, *s as i32),
1165 #[cfg(feature = "small_decimals")]
1166 DataType::Decimal64(p, s) => (*p as usize, *s as i32),
1167 DataType::Decimal128(p, s) => (*p as usize, *s as i32),
1168 DataType::Decimal256(p, s) => (*p as usize, *s as i32),
1169 other => {
1170 return Err(AvroError::SchemaError(format!(
1171 "Avro decimal requires Arrow decimal, got {other:?} for field '{}'",
1172 arrow_field.name()
1173 )));
1174 }
1175 };
1176 let sc = scale_opt.unwrap_or(0) as i32; if ap != *precision || as_ != sc {
1178 return Err(AvroError::SchemaError(format!(
1179 "Decimal precision/scale mismatch for field '{}': Avro({precision},{sc}) vs Arrow({ap},{as_})",
1180 arrow_field.name()
1181 )));
1182 }
1183 Ok(FieldPlan::Decimal {
1184 size: *fixed_size_opt,
1185 })
1186 }
1187 Codec::Interval => match arrow_field.data_type() {
1188 DataType::Interval(
1189 IntervalUnit::MonthDayNano | IntervalUnit::YearMonth | IntervalUnit::DayTime,
1190 ) => Ok(FieldPlan::Duration),
1191 other => Err(AvroError::SchemaError(format!(
1192 "Avro 'duration' logical type requires an Arrow Interval (MonthDayNano, YearMonth, or DayTime), found: {other:?}"
1193 ))),
1194 },
1195 Codec::Union(avro_branches, _, UnionMode::Dense) => {
1196 let arrow_union_fields = match arrow_field.data_type() {
1197 DataType::Union(fields, UnionMode::Dense) => fields,
1198 DataType::Union(_, UnionMode::Sparse) => {
1199 return Err(AvroError::NYI(
1200 "Sparse Arrow unions are not yet supported".to_string(),
1201 ));
1202 }
1203 other => {
1204 return Err(AvroError::SchemaError(format!(
1205 "Avro union maps to Arrow Union, found: {other:?}"
1206 )));
1207 }
1208 };
1209 if avro_branches.len() != arrow_union_fields.len() {
1210 return Err(AvroError::SchemaError(format!(
1211 "Mismatched number of branches between Avro union ({}) and Arrow union ({}) for field '{}'",
1212 avro_branches.len(),
1213 arrow_union_fields.len(),
1214 arrow_field.name()
1215 )));
1216 }
1217 let bindings = avro_branches
1218 .iter()
1219 .zip(arrow_union_fields.iter())
1220 .enumerate()
1221 .map(|(i, (avro_branch, (_, arrow_child_field)))| {
1222 Ok(FieldBinding {
1223 arrow_index: i,
1224 nullability: avro_branch.nullability(),
1225 plan: FieldPlan::build(avro_branch, arrow_child_field)?,
1226 })
1227 })
1228 .collect::<Result<Vec<_>, AvroError>>()?;
1229 Ok(FieldPlan::Union { bindings })
1230 }
1231 Codec::Union(_, _, UnionMode::Sparse) => Err(AvroError::NYI(
1232 "Sparse Arrow unions are not yet supported".to_string(),
1233 )),
1234 #[cfg(feature = "avro_custom_types")]
1235 Codec::RunEndEncoded(values_dt, _width_code) => {
1236 let values_field = match arrow_field.data_type() {
1237 DataType::RunEndEncoded(_run_ends_field, values_field) => values_field.as_ref(),
1238 other => {
1239 return Err(AvroError::SchemaError(format!(
1240 "Avro RunEndEncoded maps to Arrow DataType::RunEndEncoded, found: {other:?}"
1241 )));
1242 }
1243 };
1244 Ok(FieldPlan::RunEndEncoded {
1245 values_nullability: values_dt.nullability(),
1246 value_plan: Box::new(FieldPlan::build(values_dt.as_ref(), values_field)?),
1247 })
1248 }
1249 Codec::TimeMillis => match arrow_field.data_type() {
1250 DataType::Time32(TimeUnit::Second) => Ok(FieldPlan::TimeMillisFromSecs),
1251 _ => Ok(FieldPlan::Scalar),
1252 },
1253 _ => Ok(FieldPlan::Scalar),
1254 }
1255 }
1256}
1257
1258enum Encoder<'a> {
1259 Boolean(BooleanEncoder<'a>),
1260 Int(IntEncoder<'a, Int32Type>),
1261 Long(LongEncoder<'a, Int64Type>),
1262 TimestampMicros(LongEncoder<'a, TimestampMicrosecondType>),
1263 TimestampMillis(LongEncoder<'a, TimestampMillisecondType>),
1264 TimestampNanos(LongEncoder<'a, TimestampNanosecondType>),
1265 #[cfg(not(feature = "avro_custom_types"))]
1266 TimestampSecsToMillis(TimestampSecondsToMillisEncoder<'a>),
1267 Date32(IntEncoder<'a, Date32Type>),
1268 Time32SecsToMillis(Time32SecondsToMillisEncoder<'a>),
1269 Time32Millis(IntEncoder<'a, Time32MillisecondType>),
1270 Time64Micros(LongEncoder<'a, Time64MicrosecondType>),
1271 DurationSeconds(LongEncoder<'a, DurationSecondType>),
1272 DurationMillis(LongEncoder<'a, DurationMillisecondType>),
1273 DurationMicros(LongEncoder<'a, DurationMicrosecondType>),
1274 DurationNanos(LongEncoder<'a, DurationNanosecondType>),
1275 Float32(F32Encoder<'a>),
1276 Float64(F64Encoder<'a>),
1277 Binary(BinaryEncoder<'a, i32>),
1278 LargeBinary(BinaryEncoder<'a, i64>),
1279 Utf8(Utf8Encoder<'a>),
1280 Utf8Large(Utf8LargeEncoder<'a>),
1281 Utf8View(Utf8ViewEncoder<'a>),
1282 BinaryView(BinaryViewEncoder<'a>),
1283 List(Box<ListEncoder32<'a>>),
1284 LargeList(Box<ListEncoder64<'a>>),
1285 ListView(Box<ListViewEncoder32<'a>>),
1286 LargeListView(Box<ListViewEncoder64<'a>>),
1287 FixedSizeList(Box<FixedSizeListEncoder<'a>>),
1288 Struct(Box<StructEncoder<'a>>),
1289 Fixed(FixedEncoder<'a>),
1291 Uuid(UuidEncoder<'a>),
1293 IntervalMonthDayNanoDuration(DurationEncoder<'a, IntervalMonthDayNanoType>),
1295 #[cfg(feature = "avro_custom_types")]
1297 IntervalMonthDayNanoFixed(IntervalMonthDayNanoFixedEncoder<'a>),
1298 IntervalYearMonthDuration(DurationEncoder<'a, IntervalYearMonthType>),
1300 IntervalDayTimeDuration(DurationEncoder<'a, IntervalDayTimeType>),
1302 #[cfg(feature = "small_decimals")]
1303 Decimal32(Decimal32Encoder<'a>),
1304 #[cfg(feature = "small_decimals")]
1305 Decimal64(Decimal64Encoder<'a>),
1306 Decimal128(Decimal128Encoder<'a>),
1307 Decimal256(Decimal256Encoder<'a>),
1308 Enum(EnumEncoder<'a>),
1310 Map(Box<MapEncoder<'a>>),
1311 Union(Box<UnionEncoder<'a>>),
1312 RunEncoded16(Box<RunEncodedEncoder16<'a>>),
1314 RunEncoded32(Box<RunEncodedEncoder32<'a>>),
1315 RunEncoded64(Box<RunEncodedEncoder64<'a>>),
1316 Null,
1317 #[cfg(feature = "avro_custom_types")]
1318 Int8(Int8Encoder<'a>),
1319 #[cfg(feature = "avro_custom_types")]
1320 Int16(Int16Encoder<'a>),
1321 #[cfg(feature = "avro_custom_types")]
1322 UInt8(UInt8Encoder<'a>),
1323 #[cfg(feature = "avro_custom_types")]
1324 UInt16(UInt16Encoder<'a>),
1325 #[cfg(feature = "avro_custom_types")]
1326 UInt32(UInt32Encoder<'a>),
1327 #[cfg(feature = "avro_custom_types")]
1328 UInt64Fixed(UInt64FixedEncoder<'a>),
1329 #[cfg(feature = "avro_custom_types")]
1330 Float16Fixed(Float16FixedEncoder<'a>),
1331 #[cfg(feature = "avro_custom_types")]
1332 Date64(LongEncoder<'a, Date64Type>),
1333 #[cfg(feature = "avro_custom_types")]
1334 Time64Nanos(LongEncoder<'a, Time64NanosecondType>),
1335 #[cfg(feature = "avro_custom_types")]
1336 Time32Secs(IntEncoder<'a, Time32SecondType>),
1337 #[cfg(feature = "avro_custom_types")]
1338 TimestampSecs(LongEncoder<'a, TimestampSecondType>),
1339 #[cfg(feature = "avro_custom_types")]
1340 IntervalYearMonthFixed(IntervalYearMonthFixedEncoder<'a>),
1341 #[cfg(feature = "avro_custom_types")]
1342 IntervalDayTimeFixed(IntervalDayTimeFixedEncoder<'a>),
1343 #[cfg(not(feature = "avro_custom_types"))]
1344 Int8ToInt(Int8ToIntEncoder<'a>),
1345 #[cfg(not(feature = "avro_custom_types"))]
1346 Int16ToInt(Int16ToIntEncoder<'a>),
1347 #[cfg(not(feature = "avro_custom_types"))]
1348 UInt8ToInt(UInt8ToIntEncoder<'a>),
1349 #[cfg(not(feature = "avro_custom_types"))]
1350 UInt16ToInt(UInt16ToIntEncoder<'a>),
1351 #[cfg(not(feature = "avro_custom_types"))]
1352 UInt32ToLong(UInt32ToLongEncoder<'a>),
1353 #[cfg(not(feature = "avro_custom_types"))]
1354 UInt64ToLong(UInt64ToLongEncoder<'a>),
1355 #[cfg(not(feature = "avro_custom_types"))]
1356 Float16ToFloat(Float16ToFloatEncoder<'a>),
1357 #[cfg(not(feature = "avro_custom_types"))]
1358 Date64ToLong(Date64ToLongEncoder<'a>),
1359 #[cfg(not(feature = "avro_custom_types"))]
1360 Time64NanosToMicros(Time64NanosToMicrosEncoder<'a>),
1361}
1362
1363impl Encoder<'_> {
1364 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1366 match self {
1367 Encoder::Boolean(e) => e.encode(out, idx),
1368 Encoder::Int(e) => e.encode(out, idx),
1369 Encoder::Long(e) => e.encode(out, idx),
1370 Encoder::TimestampMicros(e) => e.encode(out, idx),
1371 Encoder::TimestampMillis(e) => e.encode(out, idx),
1372 Encoder::TimestampNanos(e) => e.encode(out, idx),
1373 #[cfg(not(feature = "avro_custom_types"))]
1374 Encoder::TimestampSecsToMillis(e) => e.encode(out, idx),
1375 Encoder::Date32(e) => e.encode(out, idx),
1376 Encoder::Time32SecsToMillis(e) => e.encode(out, idx),
1377 Encoder::Time32Millis(e) => e.encode(out, idx),
1378 Encoder::Time64Micros(e) => e.encode(out, idx),
1379 Encoder::DurationSeconds(e) => e.encode(out, idx),
1380 Encoder::DurationMicros(e) => e.encode(out, idx),
1381 Encoder::DurationMillis(e) => e.encode(out, idx),
1382 Encoder::DurationNanos(e) => e.encode(out, idx),
1383 Encoder::Float32(e) => e.encode(out, idx),
1384 Encoder::Float64(e) => e.encode(out, idx),
1385 Encoder::Binary(e) => e.encode(out, idx),
1386 Encoder::LargeBinary(e) => e.encode(out, idx),
1387 Encoder::Utf8(e) => e.encode(out, idx),
1388 Encoder::Utf8Large(e) => e.encode(out, idx),
1389 Encoder::Utf8View(e) => e.encode(out, idx),
1390 Encoder::BinaryView(e) => e.encode(out, idx),
1391 Encoder::List(e) => e.encode(out, idx),
1392 Encoder::LargeList(e) => e.encode(out, idx),
1393 Encoder::ListView(e) => e.encode(out, idx),
1394 Encoder::LargeListView(e) => e.encode(out, idx),
1395 Encoder::FixedSizeList(e) => e.encode(out, idx),
1396 Encoder::Struct(e) => e.encode(out, idx),
1397 Encoder::Fixed(e) => (e).encode(out, idx),
1398 Encoder::Uuid(e) => (e).encode(out, idx),
1399 Encoder::IntervalMonthDayNanoDuration(e) => e.encode(out, idx),
1400 #[cfg(feature = "avro_custom_types")]
1401 Encoder::IntervalMonthDayNanoFixed(e) => e.encode(out, idx),
1402 Encoder::IntervalYearMonthDuration(e) => e.encode(out, idx),
1403 Encoder::IntervalDayTimeDuration(e) => e.encode(out, idx),
1404 #[cfg(feature = "small_decimals")]
1405 Encoder::Decimal32(e) => (e).encode(out, idx),
1406 #[cfg(feature = "small_decimals")]
1407 Encoder::Decimal64(e) => (e).encode(out, idx),
1408 Encoder::Decimal128(e) => (e).encode(out, idx),
1409 Encoder::Decimal256(e) => (e).encode(out, idx),
1410 Encoder::Map(e) => (e).encode(out, idx),
1411 Encoder::Enum(e) => (e).encode(out, idx),
1412 Encoder::Union(e) => (e).encode(out, idx),
1413 Encoder::RunEncoded16(e) => (e).encode(out, idx),
1414 Encoder::RunEncoded32(e) => (e).encode(out, idx),
1415 Encoder::RunEncoded64(e) => (e).encode(out, idx),
1416 Encoder::Null => Ok(()),
1417 #[cfg(feature = "avro_custom_types")]
1418 Encoder::Int8(e) => e.encode(out, idx),
1419 #[cfg(feature = "avro_custom_types")]
1420 Encoder::Int16(e) => e.encode(out, idx),
1421 #[cfg(feature = "avro_custom_types")]
1422 Encoder::UInt8(e) => e.encode(out, idx),
1423 #[cfg(feature = "avro_custom_types")]
1424 Encoder::UInt16(e) => e.encode(out, idx),
1425 #[cfg(feature = "avro_custom_types")]
1426 Encoder::UInt32(e) => e.encode(out, idx),
1427 #[cfg(feature = "avro_custom_types")]
1428 Encoder::UInt64Fixed(e) => e.encode(out, idx),
1429 #[cfg(feature = "avro_custom_types")]
1430 Encoder::Float16Fixed(e) => e.encode(out, idx),
1431 #[cfg(feature = "avro_custom_types")]
1432 Encoder::Date64(e) => e.encode(out, idx),
1433 #[cfg(feature = "avro_custom_types")]
1434 Encoder::Time64Nanos(e) => e.encode(out, idx),
1435 #[cfg(feature = "avro_custom_types")]
1436 Encoder::Time32Secs(e) => e.encode(out, idx),
1437 #[cfg(feature = "avro_custom_types")]
1438 Encoder::TimestampSecs(e) => e.encode(out, idx),
1439 #[cfg(feature = "avro_custom_types")]
1440 Encoder::IntervalYearMonthFixed(e) => e.encode(out, idx),
1441 #[cfg(feature = "avro_custom_types")]
1442 Encoder::IntervalDayTimeFixed(e) => e.encode(out, idx),
1443 #[cfg(not(feature = "avro_custom_types"))]
1444 Encoder::Int8ToInt(e) => e.encode(out, idx),
1445 #[cfg(not(feature = "avro_custom_types"))]
1446 Encoder::Int16ToInt(e) => e.encode(out, idx),
1447 #[cfg(not(feature = "avro_custom_types"))]
1448 Encoder::UInt8ToInt(e) => e.encode(out, idx),
1449 #[cfg(not(feature = "avro_custom_types"))]
1450 Encoder::UInt16ToInt(e) => e.encode(out, idx),
1451 #[cfg(not(feature = "avro_custom_types"))]
1452 Encoder::UInt32ToLong(e) => e.encode(out, idx),
1453 #[cfg(not(feature = "avro_custom_types"))]
1454 Encoder::UInt64ToLong(e) => e.encode(out, idx),
1455 #[cfg(not(feature = "avro_custom_types"))]
1456 Encoder::Float16ToFloat(e) => e.encode(out, idx),
1457 #[cfg(not(feature = "avro_custom_types"))]
1458 Encoder::Date64ToLong(e) => e.encode(out, idx),
1459 #[cfg(not(feature = "avro_custom_types"))]
1460 Encoder::Time64NanosToMicros(e) => e.encode(out, idx),
1461 }
1462 }
1463}
1464
1465struct BooleanEncoder<'a>(&'a arrow_array::BooleanArray);
1466impl BooleanEncoder<'_> {
1467 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1468 write_bool(out, self.0.value(idx))
1469 }
1470}
1471
1472struct IntEncoder<'a, P: ArrowPrimitiveType<Native = i32>>(&'a PrimitiveArray<P>);
1474impl<P: ArrowPrimitiveType<Native = i32>> IntEncoder<'_, P> {
1475 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1476 write_int(out, self.0.value(idx))
1477 }
1478}
1479
1480struct LongEncoder<'a, P: ArrowPrimitiveType<Native = i64>>(&'a PrimitiveArray<P>);
1482impl<P: ArrowPrimitiveType<Native = i64>> LongEncoder<'_, P> {
1483 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1484 write_long(out, self.0.value(idx))
1485 }
1486}
1487
1488struct Time32SecondsToMillisEncoder<'a>(&'a PrimitiveArray<Time32SecondType>);
1490impl Time32SecondsToMillisEncoder<'_> {
1491 #[inline]
1492 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1493 let secs = self.0.value(idx);
1494 let millis = secs
1495 .checked_mul(1000)
1496 .ok_or_else(|| AvroError::InvalidArgument("time32(secs) * 1000 overflowed".into()))?;
1497 write_int(out, millis)
1498 }
1499}
1500
1501#[cfg(not(feature = "avro_custom_types"))]
1503struct TimestampSecondsToMillisEncoder<'a>(&'a PrimitiveArray<TimestampSecondType>);
1504#[cfg(not(feature = "avro_custom_types"))]
1505impl TimestampSecondsToMillisEncoder<'_> {
1506 #[inline]
1507 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1508 let secs = self.0.value(idx);
1509 let millis = secs.checked_mul(1000).ok_or_else(|| {
1510 AvroError::InvalidArgument("timestamp(secs) * 1000 overflowed".into())
1511 })?;
1512 write_long(out, millis)
1513 }
1514}
1515
1516#[cfg(feature = "avro_custom_types")]
1518struct Int8Encoder<'a>(&'a PrimitiveArray<Int8Type>);
1519#[cfg(feature = "avro_custom_types")]
1520impl Int8Encoder<'_> {
1521 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1522 write_int(out, self.0.value(idx) as i32)
1523 }
1524}
1525
1526#[cfg(feature = "avro_custom_types")]
1528struct Int16Encoder<'a>(&'a PrimitiveArray<Int16Type>);
1529#[cfg(feature = "avro_custom_types")]
1530impl Int16Encoder<'_> {
1531 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1532 write_int(out, self.0.value(idx) as i32)
1533 }
1534}
1535
1536#[cfg(feature = "avro_custom_types")]
1538struct UInt8Encoder<'a>(&'a PrimitiveArray<UInt8Type>);
1539#[cfg(feature = "avro_custom_types")]
1540impl UInt8Encoder<'_> {
1541 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1542 write_int(out, self.0.value(idx) as i32)
1543 }
1544}
1545
1546#[cfg(feature = "avro_custom_types")]
1548struct UInt16Encoder<'a>(&'a PrimitiveArray<UInt16Type>);
1549#[cfg(feature = "avro_custom_types")]
1550impl UInt16Encoder<'_> {
1551 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1552 write_int(out, self.0.value(idx) as i32)
1553 }
1554}
1555
1556#[cfg(feature = "avro_custom_types")]
1558struct UInt32Encoder<'a>(&'a PrimitiveArray<UInt32Type>);
1559#[cfg(feature = "avro_custom_types")]
1560impl UInt32Encoder<'_> {
1561 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1562 write_long(out, self.0.value(idx) as i64)
1563 }
1564}
1565
1566#[cfg(feature = "avro_custom_types")]
1568struct UInt64FixedEncoder<'a>(&'a PrimitiveArray<UInt64Type>);
1569#[cfg(feature = "avro_custom_types")]
1570impl UInt64FixedEncoder<'_> {
1571 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1572 let v = self.0.value(idx);
1573 out.write_all(&v.to_le_bytes())?;
1574 Ok(())
1575 }
1576}
1577
1578#[cfg(feature = "avro_custom_types")]
1580struct Float16FixedEncoder<'a>(&'a Float16Array);
1581#[cfg(feature = "avro_custom_types")]
1582impl Float16FixedEncoder<'_> {
1583 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1584 let v = self.0.value(idx);
1585 out.write_all(&v.to_le_bytes())?;
1586 Ok(())
1587 }
1588}
1589
1590#[cfg(feature = "avro_custom_types")]
1598struct IntervalMonthDayNanoFixedEncoder<'a>(&'a PrimitiveArray<IntervalMonthDayNanoType>);
1599#[cfg(feature = "avro_custom_types")]
1600impl IntervalMonthDayNanoFixedEncoder<'_> {
1601 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1602 let v = self.0.value(idx);
1603 let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(v);
1604 out.write_all(&months.to_le_bytes())?;
1605 out.write_all(&days.to_le_bytes())?;
1606 out.write_all(&nanos.to_le_bytes())?;
1607 Ok(())
1608 }
1609}
1610
1611#[cfg(feature = "avro_custom_types")]
1613struct IntervalYearMonthFixedEncoder<'a>(&'a PrimitiveArray<IntervalYearMonthType>);
1614#[cfg(feature = "avro_custom_types")]
1615impl IntervalYearMonthFixedEncoder<'_> {
1616 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1617 let months = self.0.value(idx);
1618 out.write_all(&months.to_le_bytes())?;
1619 Ok(())
1620 }
1621}
1622
1623#[cfg(feature = "avro_custom_types")]
1625struct IntervalDayTimeFixedEncoder<'a>(&'a PrimitiveArray<IntervalDayTimeType>);
1626#[cfg(feature = "avro_custom_types")]
1627impl IntervalDayTimeFixedEncoder<'_> {
1628 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1629 let dt = self.0.value(idx);
1630 out.write_all(&dt.days.to_le_bytes())?;
1631 out.write_all(&dt.milliseconds.to_le_bytes())?;
1632 Ok(())
1633 }
1634}
1635
1636#[cfg(not(feature = "avro_custom_types"))]
1638struct Int8ToIntEncoder<'a>(&'a PrimitiveArray<Int8Type>);
1639#[cfg(not(feature = "avro_custom_types"))]
1640impl Int8ToIntEncoder<'_> {
1641 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1642 write_int(out, self.0.value(idx) as i32)
1643 }
1644}
1645
1646#[cfg(not(feature = "avro_custom_types"))]
1648struct Int16ToIntEncoder<'a>(&'a PrimitiveArray<Int16Type>);
1649#[cfg(not(feature = "avro_custom_types"))]
1650impl Int16ToIntEncoder<'_> {
1651 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1652 write_int(out, self.0.value(idx) as i32)
1653 }
1654}
1655
1656#[cfg(not(feature = "avro_custom_types"))]
1658struct UInt8ToIntEncoder<'a>(&'a PrimitiveArray<UInt8Type>);
1659#[cfg(not(feature = "avro_custom_types"))]
1660impl UInt8ToIntEncoder<'_> {
1661 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1662 write_int(out, self.0.value(idx) as i32)
1663 }
1664}
1665
1666#[cfg(not(feature = "avro_custom_types"))]
1668struct UInt16ToIntEncoder<'a>(&'a PrimitiveArray<UInt16Type>);
1669#[cfg(not(feature = "avro_custom_types"))]
1670impl UInt16ToIntEncoder<'_> {
1671 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1672 write_int(out, self.0.value(idx) as i32)
1673 }
1674}
1675
1676#[cfg(not(feature = "avro_custom_types"))]
1678struct UInt32ToLongEncoder<'a>(&'a PrimitiveArray<UInt32Type>);
1679#[cfg(not(feature = "avro_custom_types"))]
1680impl UInt32ToLongEncoder<'_> {
1681 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1682 write_long(out, self.0.value(idx) as i64)
1683 }
1684}
1685
1686#[cfg(not(feature = "avro_custom_types"))]
1688struct UInt64ToLongEncoder<'a>(&'a PrimitiveArray<UInt64Type>);
1689#[cfg(not(feature = "avro_custom_types"))]
1690impl UInt64ToLongEncoder<'_> {
1691 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1692 let v = self.0.value(idx);
1693 if v > i64::MAX as u64 {
1694 return Err(AvroError::InvalidArgument(format!(
1695 "UInt64 value {v} exceeds i64::MAX; enable avro_custom_types feature for full UInt64 support",
1696 )));
1697 }
1698 write_long(out, v as i64)
1699 }
1700}
1701
1702#[cfg(not(feature = "avro_custom_types"))]
1704struct Float16ToFloatEncoder<'a>(&'a Float16Array);
1705#[cfg(not(feature = "avro_custom_types"))]
1706impl Float16ToFloatEncoder<'_> {
1707 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1708 out.write_all(&self.0.value(idx).to_f32().to_bits().to_le_bytes())?;
1709 Ok(())
1710 }
1711}
1712
1713#[cfg(not(feature = "avro_custom_types"))]
1715struct Date64ToLongEncoder<'a>(&'a PrimitiveArray<Date64Type>);
1716#[cfg(not(feature = "avro_custom_types"))]
1717impl Date64ToLongEncoder<'_> {
1718 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1719 write_long(out, self.0.value(idx))
1720 }
1721}
1722
1723#[cfg(not(feature = "avro_custom_types"))]
1725struct Time64NanosToMicrosEncoder<'a>(&'a PrimitiveArray<Time64NanosecondType>);
1726#[cfg(not(feature = "avro_custom_types"))]
1727impl Time64NanosToMicrosEncoder<'_> {
1728 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1729 let nanos = self.0.value(idx);
1730 let micros = nanos / 1000;
1731 write_long(out, micros)
1732 }
1733}
1734
1735struct BinaryEncoder<'a, O: OffsetSizeTrait>(&'a GenericBinaryArray<O>);
1737impl<O: OffsetSizeTrait> BinaryEncoder<'_, O> {
1738 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1739 write_len_prefixed(out, self.0.value(idx))
1740 }
1741}
1742
1743struct BinaryViewEncoder<'a>(&'a BinaryViewArray);
1745impl BinaryViewEncoder<'_> {
1746 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1747 write_len_prefixed(out, self.0.value(idx))
1748 }
1749}
1750
1751struct Utf8ViewEncoder<'a>(&'a StringViewArray);
1753impl Utf8ViewEncoder<'_> {
1754 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1755 write_len_prefixed(out, self.0.value(idx).as_bytes())
1756 }
1757}
1758
1759struct F32Encoder<'a>(&'a arrow_array::Float32Array);
1760impl F32Encoder<'_> {
1761 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1762 out.write_all(&self.0.value(idx).to_bits().to_le_bytes())?;
1764 Ok(())
1765 }
1766}
1767
1768struct F64Encoder<'a>(&'a arrow_array::Float64Array);
1769impl F64Encoder<'_> {
1770 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1771 out.write_all(&self.0.value(idx).to_bits().to_le_bytes())
1773 .map_err(Into::into)
1774 }
1775}
1776
1777struct Utf8GenericEncoder<'a, O: OffsetSizeTrait>(&'a GenericStringArray<O>);
1778
1779impl<O: OffsetSizeTrait> Utf8GenericEncoder<'_, O> {
1780 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1781 write_len_prefixed(out, self.0.value(idx).as_bytes())
1782 }
1783}
1784
1785type Utf8Encoder<'a> = Utf8GenericEncoder<'a, i32>;
1786type Utf8LargeEncoder<'a> = Utf8GenericEncoder<'a, i64>;
1787
1788enum KeyKind<'a> {
1790 Utf8(&'a GenericStringArray<i32>),
1791 LargeUtf8(&'a GenericStringArray<i64>),
1792}
1793struct MapEncoder<'a> {
1794 map: &'a MapArray,
1795 keys: KeyKind<'a>,
1796 values: FieldEncoder<'a>,
1797 keys_offset: usize,
1798 values_offset: usize,
1799}
1800
1801impl<'a> MapEncoder<'a> {
1802 fn try_new(
1803 map: &'a MapArray,
1804 values_nullability: Option<Nullability>,
1805 value_plan: &FieldPlan,
1806 ) -> Result<Self, AvroError> {
1807 let keys_arr = map.keys();
1808 let keys_kind = match keys_arr.data_type() {
1809 DataType::Utf8 => KeyKind::Utf8(keys_arr.as_string::<i32>()),
1810 DataType::LargeUtf8 => KeyKind::LargeUtf8(keys_arr.as_string::<i64>()),
1811 other => {
1812 return Err(AvroError::SchemaError(format!(
1813 "Avro map requires string keys; Arrow key type must be Utf8/LargeUtf8, found: {other:?}"
1814 )));
1815 }
1816 };
1817 Ok(Self {
1818 map,
1819 keys: keys_kind,
1820 values: FieldEncoder::make_encoder(
1821 map.values().as_ref(),
1822 value_plan,
1823 values_nullability,
1824 )?,
1825 keys_offset: keys_arr.offset(),
1826 values_offset: map.values().offset(),
1827 })
1828 }
1829
1830 fn encode_map_entries<W, O>(
1831 out: &mut W,
1832 keys: &GenericStringArray<O>,
1833 keys_offset: usize,
1834 start: usize,
1835 end: usize,
1836 mut write_item: impl FnMut(&mut W, usize) -> Result<(), AvroError>,
1837 ) -> Result<(), AvroError>
1838 where
1839 W: Write + ?Sized,
1840 O: OffsetSizeTrait,
1841 {
1842 encode_blocked_range(out, start, end, |out, j| {
1843 let j_key = j.saturating_sub(keys_offset);
1844 write_len_prefixed(out, keys.value(j_key).as_bytes())?;
1845 write_item(out, j)
1846 })
1847 }
1848
1849 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1850 let offsets = self.map.offsets();
1851 let start = offsets[idx] as usize;
1852 let end = offsets[idx + 1] as usize;
1853 let write_item = |out: &mut W, j: usize| {
1854 let j_val = j.saturating_sub(self.values_offset);
1855 self.values.encode(out, j_val)
1856 };
1857 match self.keys {
1858 KeyKind::Utf8(arr) => MapEncoder::<'a>::encode_map_entries(
1859 out,
1860 arr,
1861 self.keys_offset,
1862 start,
1863 end,
1864 write_item,
1865 ),
1866 KeyKind::LargeUtf8(arr) => MapEncoder::<'a>::encode_map_entries(
1867 out,
1868 arr,
1869 self.keys_offset,
1870 start,
1871 end,
1872 write_item,
1873 ),
1874 }
1875 }
1876}
1877
1878struct EnumEncoder<'a> {
1885 keys: &'a PrimitiveArray<Int32Type>,
1886}
1887impl EnumEncoder<'_> {
1888 fn encode<W: Write + ?Sized>(&self, out: &mut W, row: usize) -> Result<(), AvroError> {
1889 write_int(out, self.keys.value(row))
1890 }
1891}
1892
1893struct UnionEncoder<'a> {
1894 encoders: Vec<FieldEncoder<'a>>,
1895 array: &'a UnionArray,
1896 type_id_to_encoder_index: Vec<Option<usize>>,
1897}
1898
1899impl<'a> UnionEncoder<'a> {
1900 fn try_new(array: &'a UnionArray, field_bindings: &[FieldBinding]) -> Result<Self, AvroError> {
1901 let DataType::Union(fields, UnionMode::Dense) = array.data_type() else {
1902 return Err(AvroError::SchemaError("Expected Dense UnionArray".into()));
1903 };
1904 if fields.len() != field_bindings.len() {
1905 return Err(AvroError::SchemaError(format!(
1906 "Mismatched number of union branches between Arrow array ({}) and encoding plan ({})",
1907 fields.len(),
1908 field_bindings.len()
1909 )));
1910 }
1911 let max_type_id = fields.iter().map(|(tid, _)| tid).max().unwrap_or(0);
1912 let mut type_id_to_encoder_index: Vec<Option<usize>> =
1913 vec![None; (max_type_id + 1) as usize];
1914 let mut encoders = Vec::with_capacity(fields.len());
1915 for (i, (type_id, _)) in fields.iter().enumerate() {
1916 let binding = field_bindings
1917 .get(i)
1918 .ok_or_else(|| AvroError::SchemaError("Binding and field mismatch".to_string()))?;
1919 encoders.push(FieldEncoder::make_encoder(
1920 array.child(type_id).as_ref(),
1921 &binding.plan,
1922 binding.nullability,
1923 )?);
1924 type_id_to_encoder_index[type_id as usize] = Some(i);
1925 }
1926 Ok(Self {
1927 encoders,
1928 array,
1929 type_id_to_encoder_index,
1930 })
1931 }
1932
1933 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1934 let type_id = self.array.type_ids()[idx];
1939 let encoder_index = self
1940 .type_id_to_encoder_index
1941 .get(type_id as usize)
1942 .and_then(|opt| *opt)
1943 .ok_or_else(|| AvroError::SchemaError(format!("Invalid type_id {type_id}")))?;
1944 write_int(out, encoder_index as i32)?;
1945 let encoder = self.encoders.get_mut(encoder_index).ok_or_else(|| {
1946 AvroError::SchemaError(format!("Invalid encoder index {encoder_index}"))
1947 })?;
1948 encoder.encode(out, self.array.value_offset(idx))
1949 }
1950}
1951
1952struct StructEncoder<'a> {
1953 encoders: Vec<FieldEncoder<'a>>,
1954}
1955
1956impl<'a> StructEncoder<'a> {
1957 fn try_new(array: &'a StructArray, field_bindings: &[FieldBinding]) -> Result<Self, AvroError> {
1958 let mut encoders = Vec::with_capacity(field_bindings.len());
1959 for field_binding in field_bindings {
1960 let idx = field_binding.arrow_index;
1961 let column = array.columns().get(idx).ok_or_else(|| {
1962 AvroError::SchemaError(format!("Struct child index {idx} out of range"))
1963 })?;
1964 let encoder = FieldEncoder::make_encoder(
1965 column.as_ref(),
1966 &field_binding.plan,
1967 field_binding.nullability,
1968 )?;
1969 encoders.push(encoder);
1970 }
1971 Ok(Self { encoders })
1972 }
1973
1974 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
1975 for encoder in &mut self.encoders {
1976 encoder.encode(out, idx)?;
1977 }
1978 Ok(())
1979 }
1980}
1981
1982fn encode_blocked_range<W: Write + ?Sized, F>(
1986 out: &mut W,
1987 start: usize,
1988 end: usize,
1989 mut write_item: F,
1990) -> Result<(), AvroError>
1991where
1992 F: FnMut(&mut W, usize) -> Result<(), AvroError>,
1993{
1994 let len = end.saturating_sub(start);
1995 if len == 0 {
1996 write_long(out, 0)?;
1998 return Ok(());
1999 }
2000 write_long(out, len as i64)?;
2002 for row in start..end {
2003 write_item(out, row)?;
2004 }
2005 write_long(out, 0)?;
2006 Ok(())
2007}
2008
2009struct ListEncoder<'a, O: OffsetSizeTrait> {
2010 list: &'a GenericListArray<O>,
2011 values: FieldEncoder<'a>,
2012 values_offset: usize,
2013}
2014
2015type ListEncoder32<'a> = ListEncoder<'a, i32>;
2016type ListEncoder64<'a> = ListEncoder<'a, i64>;
2017
2018impl<'a, O: OffsetSizeTrait> ListEncoder<'a, O> {
2019 fn try_new(
2020 list: &'a GenericListArray<O>,
2021 items_nullability: Option<Nullability>,
2022 item_plan: &FieldPlan,
2023 ) -> Result<Self, AvroError> {
2024 Ok(Self {
2025 list,
2026 values: FieldEncoder::make_encoder(
2027 list.values().as_ref(),
2028 item_plan,
2029 items_nullability,
2030 )?,
2031 values_offset: list.values().offset(),
2032 })
2033 }
2034
2035 fn encode_list_range<W: Write + ?Sized>(
2036 &mut self,
2037 out: &mut W,
2038 start: usize,
2039 end: usize,
2040 ) -> Result<(), AvroError> {
2041 encode_blocked_range(out, start, end, |out, row| {
2042 self.values
2043 .encode(out, row.saturating_sub(self.values_offset))
2044 })
2045 }
2046
2047 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2048 let offsets = self.list.offsets();
2049 let start = offsets[idx].to_usize().ok_or_else(|| {
2050 AvroError::InvalidArgument(format!("Error converting offset[{idx}] to usize"))
2051 })?;
2052 let end = offsets[idx + 1].to_usize().ok_or_else(|| {
2053 AvroError::InvalidArgument(format!("Error converting offset[{}] to usize", idx + 1))
2054 })?;
2055 self.encode_list_range(out, start, end)
2056 }
2057}
2058
2059struct ListViewEncoder<'a, O: OffsetSizeTrait> {
2061 list: &'a GenericListViewArray<O>,
2062 values: FieldEncoder<'a>,
2063 values_offset: usize,
2064}
2065type ListViewEncoder32<'a> = ListViewEncoder<'a, i32>;
2066type ListViewEncoder64<'a> = ListViewEncoder<'a, i64>;
2067
2068impl<'a, O: OffsetSizeTrait> ListViewEncoder<'a, O> {
2069 fn try_new(
2070 list: &'a GenericListViewArray<O>,
2071 items_nullability: Option<Nullability>,
2072 item_plan: &FieldPlan,
2073 ) -> Result<Self, AvroError> {
2074 Ok(Self {
2075 list,
2076 values: FieldEncoder::make_encoder(
2077 list.values().as_ref(),
2078 item_plan,
2079 items_nullability,
2080 )?,
2081 values_offset: list.values().offset(),
2082 })
2083 }
2084
2085 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2086 let start = self.list.value_offset(idx).to_usize().ok_or_else(|| {
2087 AvroError::InvalidArgument(format!("Error converting value_offset[{idx}] to usize"))
2088 })?;
2089 let len = self.list.value_size(idx).to_usize().ok_or_else(|| {
2090 AvroError::InvalidArgument(format!("Error converting value_size[{idx}] to usize"))
2091 })?;
2092 let start = start + self.values_offset;
2093 let end = start + len;
2094 encode_blocked_range(out, start, end, |out, row| {
2095 self.values
2096 .encode(out, row.saturating_sub(self.values_offset))
2097 })
2098 }
2099}
2100
2101struct FixedSizeListEncoder<'a> {
2103 values: FieldEncoder<'a>,
2104 values_offset: usize,
2105 elem_len: usize,
2106}
2107
2108impl<'a> FixedSizeListEncoder<'a> {
2109 fn try_new(
2110 list: &'a FixedSizeListArray,
2111 items_nullability: Option<Nullability>,
2112 item_plan: &FieldPlan,
2113 ) -> Result<Self, AvroError> {
2114 Ok(Self {
2115 values: FieldEncoder::make_encoder(
2116 list.values().as_ref(),
2117 item_plan,
2118 items_nullability,
2119 )?,
2120 values_offset: list.values().offset(),
2121 elem_len: list.value_length() as usize,
2122 })
2123 }
2124
2125 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2126 let rel = idx * self.elem_len;
2128 let start = self.values_offset + rel;
2129 let end = start + self.elem_len;
2130 encode_blocked_range(out, start, end, |out, row| {
2131 self.values
2132 .encode(out, row.saturating_sub(self.values_offset))
2133 })
2134 }
2135}
2136
2137struct FixedEncoder<'a>(&'a FixedSizeBinaryArray);
2140impl FixedEncoder<'_> {
2141 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2142 let v = self.0.value(idx); out.write_all(v)?;
2144 Ok(())
2145 }
2146}
2147
2148struct UuidEncoder<'a>(&'a FixedSizeBinaryArray);
2151impl UuidEncoder<'_> {
2152 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2153 let mut buf = [0u8; 1 + uuid::fmt::Hyphenated::LENGTH];
2154 buf[0] = 0x48;
2155 let v = self.0.value(idx);
2156 let u = Uuid::from_slice(v)
2157 .map_err(|e| AvroError::InvalidArgument(format!("Invalid UUID bytes: {e}")))?;
2158 let _ = u.hyphenated().encode_lower(&mut buf[1..]);
2159 out.write_all(&buf)?;
2160 Ok(())
2161 }
2162}
2163
2164#[derive(Copy, Clone)]
2165struct DurationParts {
2166 months: u32,
2167 days: u32,
2168 millis: u32,
2169}
2170trait IntervalToDurationParts: ArrowPrimitiveType {
2172 fn duration_parts(native: Self::Native) -> Result<DurationParts, AvroError>;
2173}
2174impl IntervalToDurationParts for IntervalMonthDayNanoType {
2175 fn duration_parts(native: Self::Native) -> Result<DurationParts, AvroError> {
2176 let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(native);
2177 if months < 0 || days < 0 || nanos < 0 {
2178 return Err(AvroError::InvalidArgument(
2179 "Avro 'duration' cannot encode negative months/days/nanoseconds; enable `avro_custom_types` to round-trip signed Arrow intervals".into(),
2180 ));
2181 }
2182 if nanos % 1_000_000 != 0 {
2183 return Err(AvroError::InvalidArgument(
2184 "Avro 'duration' requires whole milliseconds; nanoseconds must be divisible by 1_000_000 (enable `avro_custom_types` to preserve nanosecond intervals)"
2185 .into(),
2186 ));
2187 }
2188 let millis = nanos / 1_000_000;
2189 if millis > u32::MAX as i64 {
2190 return Err(AvroError::InvalidArgument(
2191 "Avro 'duration' milliseconds exceed u32::MAX; enable `avro_custom_types` to preserve full Arrow Interval(MonthDayNano) range".into(),
2192 ));
2193 }
2194 Ok(DurationParts {
2195 months: months as u32,
2196 days: days as u32,
2197 millis: millis as u32,
2198 })
2199 }
2200}
2201impl IntervalToDurationParts for IntervalYearMonthType {
2202 fn duration_parts(native: Self::Native) -> Result<DurationParts, AvroError> {
2203 if native < 0 {
2204 return Err(AvroError::InvalidArgument(
2205 "Avro 'duration' cannot encode negative months; enable `avro_custom_types` to round-trip signed Arrow Interval(YearMonth)".into(),
2206 ));
2207 }
2208 Ok(DurationParts {
2209 months: native as u32,
2210 days: 0,
2211 millis: 0,
2212 })
2213 }
2214}
2215impl IntervalToDurationParts for IntervalDayTimeType {
2216 fn duration_parts(native: Self::Native) -> Result<DurationParts, AvroError> {
2217 let (days, millis) = IntervalDayTimeType::to_parts(native);
2218 if days < 0 || millis < 0 {
2219 return Err(AvroError::InvalidArgument(
2220 "Avro 'duration' cannot encode negative days or milliseconds; enable `avro_custom_types` to round-trip signed Arrow Interval(DayTime)".into(),
2221 ));
2222 }
2223 Ok(DurationParts {
2224 months: 0,
2225 days: days as u32,
2226 millis: millis as u32,
2227 })
2228 }
2229}
2230
2231struct DurationEncoder<'a, P: ArrowPrimitiveType + IntervalToDurationParts>(&'a PrimitiveArray<P>);
2234impl<P: ArrowPrimitiveType + IntervalToDurationParts> DurationEncoder<'_, P> {
2235 #[inline(always)]
2236 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2237 let parts = P::duration_parts(self.0.value(idx))?;
2238 let months = parts.months.to_le_bytes();
2239 let days = parts.days.to_le_bytes();
2240 let ms = parts.millis.to_le_bytes();
2241 let buf = [
2257 months[0], months[1], months[2], months[3], days[0], days[1], days[2], days[3], ms[0],
2258 ms[1], ms[2], ms[3],
2259 ];
2260 out.write_all(&buf)?;
2261 Ok(())
2262 }
2263}
2264
2265trait DecimalBeBytes<const N: usize> {
2268 fn value_be_bytes(&self, idx: usize) -> [u8; N];
2269}
2270#[cfg(feature = "small_decimals")]
2271impl DecimalBeBytes<4> for Decimal32Array {
2272 fn value_be_bytes(&self, idx: usize) -> [u8; 4] {
2273 self.value(idx).to_be_bytes()
2274 }
2275}
2276#[cfg(feature = "small_decimals")]
2277impl DecimalBeBytes<8> for Decimal64Array {
2278 fn value_be_bytes(&self, idx: usize) -> [u8; 8] {
2279 self.value(idx).to_be_bytes()
2280 }
2281}
2282impl DecimalBeBytes<16> for Decimal128Array {
2283 fn value_be_bytes(&self, idx: usize) -> [u8; 16] {
2284 self.value(idx).to_be_bytes()
2285 }
2286}
2287impl DecimalBeBytes<32> for Decimal256Array {
2288 fn value_be_bytes(&self, idx: usize) -> [u8; 32] {
2289 self.value(idx).to_be_bytes()
2291 }
2292}
2293
2294struct DecimalEncoder<'a, const N: usize, A: DecimalBeBytes<N>> {
2300 arr: &'a A,
2301 fixed_size: Option<usize>,
2302}
2303
2304impl<'a, const N: usize, A: DecimalBeBytes<N>> DecimalEncoder<'a, N, A> {
2305 fn new(arr: &'a A, fixed_size: Option<usize>) -> Self {
2306 Self { arr, fixed_size }
2307 }
2308
2309 fn encode<W: Write + ?Sized>(&self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2310 let be = self.arr.value_be_bytes(idx);
2311 match self.fixed_size {
2312 Some(n) => write_sign_extended(out, &be, n),
2313 None => write_len_prefixed(out, minimal_twos_complement(&be)),
2314 }
2315 }
2316}
2317
2318#[cfg(feature = "small_decimals")]
2319type Decimal32Encoder<'a> = DecimalEncoder<'a, 4, Decimal32Array>;
2320#[cfg(feature = "small_decimals")]
2321type Decimal64Encoder<'a> = DecimalEncoder<'a, 8, Decimal64Array>;
2322type Decimal128Encoder<'a> = DecimalEncoder<'a, 16, Decimal128Array>;
2323type Decimal256Encoder<'a> = DecimalEncoder<'a, 32, Decimal256Array>;
2324
2325struct RunEncodedEncoder<'a, R: RunEndIndexType> {
2329 ends_slice: &'a [<R as ArrowPrimitiveType>::Native],
2330 base: usize,
2331 len: usize,
2332 values: FieldEncoder<'a>,
2333 cur_run: usize,
2335 cur_end: usize,
2337}
2338
2339type RunEncodedEncoder16<'a> = RunEncodedEncoder<'a, Int16Type>;
2340type RunEncodedEncoder32<'a> = RunEncodedEncoder<'a, Int32Type>;
2341type RunEncodedEncoder64<'a> = RunEncodedEncoder<'a, Int64Type>;
2342
2343impl<'a, R: RunEndIndexType> RunEncodedEncoder<'a, R> {
2344 fn new(arr: &'a RunArray<R>, values: FieldEncoder<'a>) -> Self {
2345 let ends = arr.run_ends();
2346 let base = ends.get_start_physical_index();
2347 let slice = ends.values();
2348 let len = ends.len();
2349 let cur_end = if len == 0 { 0 } else { slice[base].as_usize() };
2350 Self {
2351 ends_slice: slice,
2352 base,
2353 len,
2354 values,
2355 cur_run: 0,
2356 cur_end,
2357 }
2358 }
2359
2360 #[inline(always)]
2363 fn advance_to_row(&mut self, idx: usize) -> Result<(), AvroError> {
2364 if idx < self.cur_end {
2365 return Ok(());
2366 }
2367 while self.cur_run + 1 < self.len && idx >= self.cur_end {
2369 self.cur_run += 1;
2370 self.cur_end = self.ends_slice[self.base + self.cur_run].as_usize();
2371 }
2372 if idx < self.cur_end {
2373 Ok(())
2374 } else {
2375 Err(AvroError::InvalidArgument(format!(
2376 "row index {idx} out of bounds for run-ends ({} runs)",
2377 self.len
2378 )))
2379 }
2380 }
2381
2382 #[inline(always)]
2383 fn encode<W: Write + ?Sized>(&mut self, out: &mut W, idx: usize) -> Result<(), AvroError> {
2384 self.advance_to_row(idx)?;
2385 self.values.encode(out, self.cur_run)
2388 }
2389}
2390
2391#[cfg(test)]
2392mod tests {
2393 use super::*;
2394 use arrow_array::types::Int32Type;
2395 use arrow_array::{
2396 Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array,
2397 Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, NullArray,
2398 StringArray,
2399 };
2400 use arrow_buffer::Buffer;
2401 use arrow_schema::{DataType, Field, Fields, UnionFields};
2402
2403 fn zigzag_i64(v: i64) -> u64 {
2404 ((v << 1) ^ (v >> 63)) as u64
2405 }
2406
2407 fn varint(mut x: u64) -> Vec<u8> {
2408 let mut out = Vec::new();
2409 while (x & !0x7f) != 0 {
2410 out.push(((x & 0x7f) as u8) | 0x80);
2411 x >>= 7;
2412 }
2413 out.push((x & 0x7f) as u8);
2414 out
2415 }
2416
2417 fn avro_long_bytes(v: i64) -> Vec<u8> {
2418 varint(zigzag_i64(v))
2419 }
2420
2421 fn avro_len_prefixed_bytes(payload: &[u8]) -> Vec<u8> {
2422 let mut out = avro_long_bytes(payload.len() as i64);
2423 out.extend_from_slice(payload);
2424 out
2425 }
2426
2427 fn duration_fixed12(months: u32, days: u32, millis: u32) -> [u8; 12] {
2428 let m = months.to_le_bytes();
2429 let d = days.to_le_bytes();
2430 let ms = millis.to_le_bytes();
2431 [
2432 m[0], m[1], m[2], m[3], d[0], d[1], d[2], d[3], ms[0], ms[1], ms[2], ms[3],
2433 ]
2434 }
2435
2436 #[cfg(feature = "avro_custom_types")]
2437 fn interval_mdn_fixed16(months: i32, days: i32, nanos: i64) -> [u8; 16] {
2438 let m = months.to_le_bytes();
2439 let d = days.to_le_bytes();
2440 let n = nanos.to_le_bytes();
2441 [
2442 m[0], m[1], m[2], m[3], d[0], d[1], d[2], d[3], n[0], n[1], n[2], n[3], n[4], n[5],
2443 n[6], n[7],
2444 ]
2445 }
2446
2447 fn encode_all(
2448 array: &dyn Array,
2449 plan: &FieldPlan,
2450 nullability: Option<Nullability>,
2451 ) -> Vec<u8> {
2452 let mut enc = FieldEncoder::make_encoder(array, plan, nullability).unwrap();
2453 let mut out = Vec::new();
2454 for i in 0..array.len() {
2455 enc.encode(&mut out, i).unwrap();
2456 }
2457 out
2458 }
2459
2460 fn assert_bytes_eq(actual: &[u8], expected: &[u8]) {
2461 if actual != expected {
2462 let to_hex = |b: &[u8]| {
2463 b.iter()
2464 .map(|x| format!("{x:02X}"))
2465 .collect::<Vec<_>>()
2466 .join(" ")
2467 };
2468 panic!(
2469 "mismatch\n expected: [{}]\n actual: [{}]",
2470 to_hex(expected),
2471 to_hex(actual)
2472 );
2473 }
2474 }
2475
2476 fn row_slice<'a>(buf: &'a [u8], offsets: &[usize], row: usize) -> &'a [u8] {
2477 let start = offsets[row];
2478 let end = offsets[row + 1];
2479 &buf[start..end]
2480 }
2481
2482 #[test]
2483 fn binary_encoder() {
2484 let values: Vec<&[u8]> = vec![b"", b"ab", b"\x00\xFF"];
2485 let arr = BinaryArray::from_vec(values);
2486 let mut expected = Vec::new();
2487 for payload in [b"" as &[u8], b"ab", b"\x00\xFF"] {
2488 expected.extend(avro_len_prefixed_bytes(payload));
2489 }
2490 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2491 assert_bytes_eq(&got, &expected);
2492 }
2493
2494 #[test]
2495 fn large_binary_encoder() {
2496 let values: Vec<&[u8]> = vec![b"xyz", b""];
2497 let arr = LargeBinaryArray::from_vec(values);
2498 let mut expected = Vec::new();
2499 for payload in [b"xyz" as &[u8], b""] {
2500 expected.extend(avro_len_prefixed_bytes(payload));
2501 }
2502 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2503 assert_bytes_eq(&got, &expected);
2504 }
2505
2506 #[test]
2507 fn utf8_encoder() {
2508 let arr = StringArray::from(vec!["", "A", "BC"]);
2509 let mut expected = Vec::new();
2510 for s in ["", "A", "BC"] {
2511 expected.extend(avro_len_prefixed_bytes(s.as_bytes()));
2512 }
2513 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2514 assert_bytes_eq(&got, &expected);
2515 }
2516
2517 #[test]
2518 fn large_utf8_encoder() {
2519 let arr = LargeStringArray::from(vec!["hello", ""]);
2520 let mut expected = Vec::new();
2521 for s in ["hello", ""] {
2522 expected.extend(avro_len_prefixed_bytes(s.as_bytes()));
2523 }
2524 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2525 assert_bytes_eq(&got, &expected);
2526 }
2527
2528 #[test]
2529 fn list_encoder_int32() {
2530 let values = Int32Array::from(vec![1, 2, 3]);
2532 let offsets = vec![0, 2, 2, 3];
2533 let list = ListArray::new(
2534 Field::new("item", DataType::Int32, true).into(),
2535 arrow_buffer::OffsetBuffer::new(offsets.into()),
2536 Arc::new(values) as ArrayRef,
2537 None,
2538 );
2539 let mut expected = Vec::new();
2541 expected.extend(avro_long_bytes(2));
2543 expected.extend(avro_long_bytes(1));
2544 expected.extend(avro_long_bytes(2));
2545 expected.extend(avro_long_bytes(0));
2546 expected.extend(avro_long_bytes(0));
2548 expected.extend(avro_long_bytes(1));
2550 expected.extend(avro_long_bytes(3));
2551 expected.extend(avro_long_bytes(0));
2552
2553 let plan = FieldPlan::List {
2554 items_nullability: None,
2555 item_plan: Box::new(FieldPlan::Scalar),
2556 };
2557 let got = encode_all(&list, &plan, None);
2558 assert_bytes_eq(&got, &expected);
2559 }
2560
2561 #[test]
2562 fn fixed_size_list_encoder_int32_sliced() {
2563 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2565 let list = FixedSizeListArray::new(
2566 Field::new("item", DataType::Int32, true).into(),
2567 2,
2568 Arc::new(values) as ArrayRef,
2569 None,
2570 )
2571 .slice(1, 2);
2572 let mut expected = Vec::new();
2573 for row in [[3, 4], [5, 6]] {
2574 expected.extend(avro_long_bytes(2));
2575 for value in row {
2576 expected.extend(avro_long_bytes(value));
2577 }
2578 expected.extend(avro_long_bytes(0));
2579 }
2580
2581 let plan = FieldPlan::List {
2582 items_nullability: None,
2583 item_plan: Box::new(FieldPlan::Scalar),
2584 };
2585 let got = encode_all(&list, &plan, None);
2586 assert_bytes_eq(&got, &expected);
2587 }
2588
2589 #[test]
2590 fn struct_encoder_two_fields() {
2591 let a = Int32Array::from(vec![1, 2]);
2593 let b = StringArray::from(vec!["x", "y"]);
2594 let fields = Fields::from(vec![
2595 Field::new("a", DataType::Int32, true),
2596 Field::new("b", DataType::Utf8, true),
2597 ]);
2598 let struct_arr = StructArray::new(
2599 fields.clone(),
2600 vec![Arc::new(a) as ArrayRef, Arc::new(b) as ArrayRef],
2601 None,
2602 );
2603 let plan = FieldPlan::Struct {
2604 bindings: vec![
2605 FieldBinding {
2606 arrow_index: 0,
2607 nullability: None,
2608 plan: FieldPlan::Scalar,
2609 },
2610 FieldBinding {
2611 arrow_index: 1,
2612 nullability: None,
2613 plan: FieldPlan::Scalar,
2614 },
2615 ],
2616 };
2617 let got = encode_all(&struct_arr, &plan, None);
2618 let mut expected = Vec::new();
2620 expected.extend(avro_long_bytes(1)); expected.extend(avro_len_prefixed_bytes(b"x")); expected.extend(avro_long_bytes(2)); expected.extend(avro_len_prefixed_bytes(b"y")); assert_bytes_eq(&got, &expected);
2625 }
2626
2627 #[test]
2628 fn enum_encoder_dictionary() {
2629 let dict_values = StringArray::from(vec!["A", "B", "C"]);
2631 let keys = Int32Array::from(vec![2, 0, 1]);
2632 let dict =
2633 DictionaryArray::<Int32Type>::try_new(keys, Arc::new(dict_values) as ArrayRef).unwrap();
2634 let symbols = Arc::<[String]>::from(
2635 vec!["A".to_string(), "B".to_string(), "C".to_string()].into_boxed_slice(),
2636 );
2637 let plan = FieldPlan::Enum { symbols };
2638 let got = encode_all(&dict, &plan, None);
2639 let mut expected = Vec::new();
2640 expected.extend(avro_long_bytes(2));
2641 expected.extend(avro_long_bytes(0));
2642 expected.extend(avro_long_bytes(1));
2643 assert_bytes_eq(&got, &expected);
2644 }
2645
2646 #[test]
2647 fn decimal_bytes_and_fixed() {
2648 let dec = Decimal128Array::from(vec![1i128, -1i128, 0i128])
2650 .with_precision_and_scale(20, 0)
2651 .unwrap();
2652 let plan_bytes = FieldPlan::Decimal { size: None };
2654 let got_bytes = encode_all(&dec, &plan_bytes, None);
2655 let mut expected_bytes = Vec::new();
2657 expected_bytes.extend(avro_len_prefixed_bytes(&[0x01]));
2658 expected_bytes.extend(avro_len_prefixed_bytes(&[0xFF]));
2659 expected_bytes.extend(avro_len_prefixed_bytes(&[0x00]));
2660 assert_bytes_eq(&got_bytes, &expected_bytes);
2661
2662 let plan_fixed = FieldPlan::Decimal { size: Some(16) };
2663 let got_fixed = encode_all(&dec, &plan_fixed, None);
2664 let mut expected_fixed = Vec::new();
2665 expected_fixed.extend_from_slice(&1i128.to_be_bytes());
2666 expected_fixed.extend_from_slice(&(-1i128).to_be_bytes());
2667 expected_fixed.extend_from_slice(&0i128.to_be_bytes());
2668 assert_bytes_eq(&got_fixed, &expected_fixed);
2669 }
2670
2671 #[test]
2672 fn decimal_bytes_256() {
2673 use arrow_buffer::i256;
2674 let dec = Decimal256Array::from(vec![
2676 i256::from_i128(1),
2677 i256::from_i128(-1),
2678 i256::from_i128(0),
2679 ])
2680 .with_precision_and_scale(76, 0)
2681 .unwrap();
2682 let plan_bytes = FieldPlan::Decimal { size: None };
2684 let got_bytes = encode_all(&dec, &plan_bytes, None);
2685 let mut expected_bytes = Vec::new();
2687 expected_bytes.extend(avro_len_prefixed_bytes(&[0x01]));
2688 expected_bytes.extend(avro_len_prefixed_bytes(&[0xFF]));
2689 expected_bytes.extend(avro_len_prefixed_bytes(&[0x00]));
2690 assert_bytes_eq(&got_bytes, &expected_bytes);
2691
2692 let plan_fixed = FieldPlan::Decimal { size: Some(32) };
2694 let got_fixed = encode_all(&dec, &plan_fixed, None);
2695 let mut expected_fixed = Vec::new();
2696 expected_fixed.extend_from_slice(&i256::from_i128(1).to_be_bytes());
2697 expected_fixed.extend_from_slice(&i256::from_i128(-1).to_be_bytes());
2698 expected_fixed.extend_from_slice(&i256::from_i128(0).to_be_bytes());
2699 assert_bytes_eq(&got_fixed, &expected_fixed);
2700 }
2701
2702 #[cfg(feature = "small_decimals")]
2703 #[test]
2704 fn decimal_bytes_and_fixed_32() {
2705 let dec = Decimal32Array::from(vec![1i32, -1i32, 0i32])
2707 .with_precision_and_scale(9, 0)
2708 .unwrap();
2709 let plan_bytes = FieldPlan::Decimal { size: None };
2711 let got_bytes = encode_all(&dec, &plan_bytes, None);
2712 let mut expected_bytes = Vec::new();
2713 expected_bytes.extend(avro_len_prefixed_bytes(&[0x01]));
2714 expected_bytes.extend(avro_len_prefixed_bytes(&[0xFF]));
2715 expected_bytes.extend(avro_len_prefixed_bytes(&[0x00]));
2716 assert_bytes_eq(&got_bytes, &expected_bytes);
2717 let plan_fixed = FieldPlan::Decimal { size: Some(4) };
2719 let got_fixed = encode_all(&dec, &plan_fixed, None);
2720 let mut expected_fixed = Vec::new();
2721 expected_fixed.extend_from_slice(&1i32.to_be_bytes());
2722 expected_fixed.extend_from_slice(&(-1i32).to_be_bytes());
2723 expected_fixed.extend_from_slice(&0i32.to_be_bytes());
2724 assert_bytes_eq(&got_fixed, &expected_fixed);
2725 }
2726
2727 #[cfg(feature = "small_decimals")]
2728 #[test]
2729 fn decimal_bytes_and_fixed_64() {
2730 let dec = Decimal64Array::from(vec![1i64, -1i64, 0i64])
2732 .with_precision_and_scale(18, 0)
2733 .unwrap();
2734 let plan_bytes = FieldPlan::Decimal { size: None };
2736 let got_bytes = encode_all(&dec, &plan_bytes, None);
2737 let mut expected_bytes = Vec::new();
2738 expected_bytes.extend(avro_len_prefixed_bytes(&[0x01]));
2739 expected_bytes.extend(avro_len_prefixed_bytes(&[0xFF]));
2740 expected_bytes.extend(avro_len_prefixed_bytes(&[0x00]));
2741 assert_bytes_eq(&got_bytes, &expected_bytes);
2742 let plan_fixed = FieldPlan::Decimal { size: Some(8) };
2744 let got_fixed = encode_all(&dec, &plan_fixed, None);
2745 let mut expected_fixed = Vec::new();
2746 expected_fixed.extend_from_slice(&1i64.to_be_bytes());
2747 expected_fixed.extend_from_slice(&(-1i64).to_be_bytes());
2748 expected_fixed.extend_from_slice(&0i64.to_be_bytes());
2749 assert_bytes_eq(&got_fixed, &expected_fixed);
2750 }
2751
2752 #[test]
2753 fn float32_and_float64_encoders() {
2754 let f32a = Float32Array::from(vec![0.0f32, -1.5f32, f32::from_bits(0x7fc00000)]); let f64a = Float64Array::from(vec![0.0f64, -2.25f64]);
2756 let mut expected32 = Vec::new();
2758 for v in [0.0f32, -1.5f32, f32::from_bits(0x7fc00000)] {
2759 expected32.extend_from_slice(&v.to_bits().to_le_bytes());
2760 }
2761 let got32 = encode_all(&f32a, &FieldPlan::Scalar, None);
2762 assert_bytes_eq(&got32, &expected32);
2763 let mut expected64 = Vec::new();
2765 for v in [0.0f64, -2.25f64] {
2766 expected64.extend_from_slice(&v.to_bits().to_le_bytes());
2767 }
2768 let got64 = encode_all(&f64a, &FieldPlan::Scalar, None);
2769 assert_bytes_eq(&got64, &expected64);
2770 }
2771
2772 #[test]
2773 fn long_encoder_int64() {
2774 let arr = Int64Array::from(vec![0i64, 1i64, -1i64, 2i64, -2i64, i64::MIN + 1]);
2775 let mut expected = Vec::new();
2776 for v in [0, 1, -1, 2, -2, i64::MIN + 1] {
2777 expected.extend(avro_long_bytes(v));
2778 }
2779 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2780 assert_bytes_eq(&got, &expected);
2781 }
2782
2783 #[test]
2784 fn fixed_encoder_plain() {
2785 let data = [[0xDE, 0xAD, 0xBE, 0xEF], [0x00, 0x01, 0x02, 0x03]];
2787 let values: Vec<Vec<u8>> = data.iter().map(|x| x.to_vec()).collect();
2788 let arr = FixedSizeBinaryArray::try_from_iter(values.into_iter()).unwrap();
2789 let got = encode_all(&arr, &FieldPlan::Scalar, None);
2790 let mut expected = Vec::new();
2791 expected.extend_from_slice(&data[0]);
2792 expected.extend_from_slice(&data[1]);
2793 assert_bytes_eq(&got, &expected);
2794 }
2795
2796 #[test]
2797 fn uuid_encoder_test() {
2798 let u = Uuid::parse_str("00112233-4455-6677-8899-aabbccddeeff").unwrap();
2800 let bytes = *u.as_bytes();
2801 let arr_ok = FixedSizeBinaryArray::try_from_iter(vec![bytes.to_vec()].into_iter()).unwrap();
2802 let mut expected = Vec::new();
2804 expected.push(0x48);
2805 expected.extend_from_slice(u.hyphenated().to_string().as_bytes());
2806 let got = encode_all(&arr_ok, &FieldPlan::Uuid, None);
2807 assert_bytes_eq(&got, &expected);
2808 }
2809
2810 #[test]
2811 fn uuid_encoder_error() {
2812 let arr =
2814 FixedSizeBinaryArray::try_new(10, arrow_buffer::Buffer::from(vec![0u8; 10]), None)
2815 .unwrap();
2816 let plan = FieldPlan::Uuid;
2817 let mut enc = FieldEncoder::make_encoder(&arr, &plan, None).unwrap();
2818 let mut out = Vec::new();
2819 let err = enc.encode(&mut out, 0).unwrap_err();
2820 match err {
2821 AvroError::InvalidArgument(msg) => {
2822 assert!(msg.contains("Invalid UUID bytes"))
2823 }
2824 other => panic!("expected InvalidArgument, got {other:?}"),
2825 }
2826 }
2827
2828 fn test_scalar_primitive_encoding<T>(
2829 non_nullable_data: &[T::Native],
2830 nullable_data: &[Option<T::Native>],
2831 ) where
2832 T: ArrowPrimitiveType,
2833 T::Native: Into<i64> + Copy,
2834 PrimitiveArray<T>: From<Vec<<T as ArrowPrimitiveType>::Native>>,
2835 {
2836 let plan = FieldPlan::Scalar;
2837
2838 let array = PrimitiveArray::<T>::from(non_nullable_data.to_vec());
2839 let got = encode_all(&array, &plan, None);
2840
2841 let mut expected = Vec::new();
2842 for &value in non_nullable_data {
2843 expected.extend(avro_long_bytes(value.into()));
2844 }
2845 assert_bytes_eq(&got, &expected);
2846
2847 let array_nullable: PrimitiveArray<T> = nullable_data.iter().copied().collect();
2848 let got_nullable = encode_all(&array_nullable, &plan, Some(Nullability::NullFirst));
2849
2850 let mut expected_nullable = Vec::new();
2851 for &opt_value in nullable_data {
2852 match opt_value {
2853 Some(value) => {
2854 expected_nullable.extend(avro_long_bytes(1));
2856 expected_nullable.extend(avro_long_bytes(value.into()));
2857 }
2858 None => {
2859 expected_nullable.extend(avro_long_bytes(0));
2861 }
2862 }
2863 }
2864 assert_bytes_eq(&got_nullable, &expected_nullable);
2865 }
2866
2867 #[test]
2868 fn date32_encoder() {
2869 test_scalar_primitive_encoding::<Date32Type>(
2870 &[
2871 19345, 0, -1, ],
2875 &[Some(19345), None],
2876 );
2877 }
2878
2879 #[test]
2880 fn time32_millis_encoder() {
2881 test_scalar_primitive_encoding::<Time32MillisecondType>(
2882 &[
2883 0, 49530123, 86399999, ],
2887 &[None, Some(49530123)],
2888 );
2889 }
2890
2891 #[test]
2892 fn time64_micros_encoder() {
2893 test_scalar_primitive_encoding::<Time64MicrosecondType>(
2894 &[
2895 0, 86399999999, ],
2898 &[Some(86399999999), None],
2899 );
2900 }
2901
2902 #[test]
2903 fn timestamp_millis_encoder() {
2904 test_scalar_primitive_encoding::<TimestampMillisecondType>(
2905 &[
2906 1704067200000, 0, -123456789, ],
2910 &[None, Some(1704067200000)],
2911 );
2912 }
2913
2914 #[test]
2915 fn map_encoder_string_keys_int_values() {
2916 let keys = StringArray::from(vec!["k1", "k2"]);
2920 let values = Int32Array::from(vec![1, 2]);
2921 let entries_fields = Fields::from(vec![
2922 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2923 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
2924 ]);
2925 let entries = StructArray::new(
2926 entries_fields,
2927 vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef],
2928 None,
2929 );
2930 let offsets = arrow_buffer::OffsetBuffer::new(vec![0i32, 2, 2].into());
2931 let map = MapArray::new(
2932 Field::new(
2933 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2934 entries.data_type().clone(),
2935 false,
2936 )
2937 .into(),
2938 offsets,
2939 entries,
2940 None,
2941 false,
2942 );
2943 let plan = FieldPlan::Map {
2944 values_nullability: None,
2945 value_plan: Box::new(FieldPlan::Scalar),
2946 };
2947 let got = encode_all(&map, &plan, None);
2948 let mut expected = Vec::new();
2949 expected.extend(avro_long_bytes(2));
2951 expected.extend(avro_len_prefixed_bytes(b"k1"));
2952 expected.extend(avro_long_bytes(1));
2953 expected.extend(avro_len_prefixed_bytes(b"k2"));
2954 expected.extend(avro_long_bytes(2));
2955 expected.extend(avro_long_bytes(0));
2956 expected.extend(avro_long_bytes(0));
2958 assert_bytes_eq(&got, &expected);
2959 }
2960
2961 #[test]
2962 fn union_encoder_string_int() {
2963 let strings = StringArray::from(vec!["hello", "world"]);
2964 let ints = Int32Array::from(vec![10, 20, 30]);
2965
2966 let union_fields = UnionFields::try_new(
2967 vec![0, 1],
2968 vec![
2969 Field::new("v_str", DataType::Utf8, true),
2970 Field::new("v_int", DataType::Int32, true),
2971 ],
2972 )
2973 .unwrap();
2974
2975 let type_ids = Buffer::from_slice_ref([0_i8, 1, 1, 0, 1]);
2976 let offsets = Buffer::from_slice_ref([0_i32, 0, 1, 1, 2]);
2977
2978 let union_array = UnionArray::try_new(
2979 union_fields,
2980 type_ids.into(),
2981 Some(offsets.into()),
2982 vec![Arc::new(strings), Arc::new(ints)],
2983 )
2984 .unwrap();
2985
2986 let plan = FieldPlan::Union {
2987 bindings: vec![
2988 FieldBinding {
2989 arrow_index: 0,
2990 nullability: None,
2991 plan: FieldPlan::Scalar,
2992 },
2993 FieldBinding {
2994 arrow_index: 1,
2995 nullability: None,
2996 plan: FieldPlan::Scalar,
2997 },
2998 ],
2999 };
3000
3001 let got = encode_all(&union_array, &plan, None);
3002
3003 let mut expected = Vec::new();
3004 expected.extend(avro_long_bytes(0));
3005 expected.extend(avro_len_prefixed_bytes(b"hello"));
3006 expected.extend(avro_long_bytes(1));
3007 expected.extend(avro_long_bytes(10));
3008 expected.extend(avro_long_bytes(1));
3009 expected.extend(avro_long_bytes(20));
3010 expected.extend(avro_long_bytes(0));
3011 expected.extend(avro_len_prefixed_bytes(b"world"));
3012 expected.extend(avro_long_bytes(1));
3013 expected.extend(avro_long_bytes(30));
3014
3015 assert_bytes_eq(&got, &expected);
3016 }
3017
3018 #[test]
3019 fn union_encoder_null_string_int() {
3020 let nulls = NullArray::new(1);
3021 let strings = StringArray::from(vec!["hello"]);
3022 let ints = Int32Array::from(vec![10]);
3023
3024 let union_fields = UnionFields::try_new(
3025 vec![0, 1, 2],
3026 vec![
3027 Field::new("v_null", DataType::Null, true),
3028 Field::new("v_str", DataType::Utf8, true),
3029 Field::new("v_int", DataType::Int32, true),
3030 ],
3031 )
3032 .unwrap();
3033
3034 let type_ids = Buffer::from_slice_ref([0_i8, 1, 2]);
3035 let offsets = Buffer::from_slice_ref([0_i32, 0, 0]);
3039
3040 let union_array = UnionArray::try_new(
3041 union_fields,
3042 type_ids.into(),
3043 Some(offsets.into()),
3044 vec![Arc::new(nulls), Arc::new(strings), Arc::new(ints)],
3045 )
3046 .unwrap();
3047
3048 let plan = FieldPlan::Union {
3049 bindings: vec![
3050 FieldBinding {
3051 arrow_index: 0,
3052 nullability: None,
3053 plan: FieldPlan::Scalar,
3054 },
3055 FieldBinding {
3056 arrow_index: 1,
3057 nullability: None,
3058 plan: FieldPlan::Scalar,
3059 },
3060 FieldBinding {
3061 arrow_index: 2,
3062 nullability: None,
3063 plan: FieldPlan::Scalar,
3064 },
3065 ],
3066 };
3067
3068 let got = encode_all(&union_array, &plan, None);
3069
3070 let mut expected = Vec::new();
3071 expected.extend(avro_long_bytes(0));
3072 expected.extend(avro_long_bytes(1));
3073 expected.extend(avro_len_prefixed_bytes(b"hello"));
3074 expected.extend(avro_long_bytes(2));
3075 expected.extend(avro_long_bytes(10));
3076
3077 assert_bytes_eq(&got, &expected);
3078 }
3079
3080 #[test]
3081 fn list64_encoder_int32() {
3082 let values = Int32Array::from(vec![1, 2, 3]);
3084 let offsets: Vec<i64> = vec![0, 3, 3];
3085 let list = LargeListArray::new(
3086 Field::new("item", DataType::Int32, true).into(),
3087 arrow_buffer::OffsetBuffer::new(offsets.into()),
3088 Arc::new(values) as ArrayRef,
3089 None,
3090 );
3091 let plan = FieldPlan::List {
3092 items_nullability: None,
3093 item_plan: Box::new(FieldPlan::Scalar),
3094 };
3095 let got = encode_all(&list, &plan, None);
3096 let mut expected = Vec::new();
3098 expected.extend(avro_long_bytes(3));
3099 expected.extend(avro_long_bytes(1));
3100 expected.extend(avro_long_bytes(2));
3101 expected.extend(avro_long_bytes(3));
3102 expected.extend(avro_long_bytes(0));
3103 expected.extend(avro_long_bytes(0));
3104 assert_bytes_eq(&got, &expected);
3105 }
3106
3107 #[test]
3108 fn int_encoder_test() {
3109 let ints = Int32Array::from(vec![0, -1, 2]);
3110 let mut expected_i = Vec::new();
3111 for v in [0i32, -1, 2] {
3112 expected_i.extend(avro_long_bytes(v as i64));
3113 }
3114 let got_i = encode_all(&ints, &FieldPlan::Scalar, None);
3115 assert_bytes_eq(&got_i, &expected_i);
3116 }
3117
3118 #[test]
3119 fn boolean_encoder_test() {
3120 let bools = BooleanArray::from(vec![true, false]);
3121 let mut expected_b = Vec::new();
3122 expected_b.extend_from_slice(&[1]);
3123 expected_b.extend_from_slice(&[0]);
3124 let got_b = encode_all(&bools, &FieldPlan::Scalar, None);
3125 assert_bytes_eq(&got_b, &expected_b);
3126 }
3127
3128 #[test]
3129 #[cfg(feature = "avro_custom_types")]
3130 fn duration_encoding_seconds() {
3131 let arr: PrimitiveArray<DurationSecondType> = vec![0i64, -1, 2].into();
3132 let mut expected = Vec::new();
3133 for v in [0i64, -1, 2] {
3134 expected.extend_from_slice(&avro_long_bytes(v));
3135 }
3136 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3137 assert_bytes_eq(&got, &expected);
3138 }
3139
3140 #[test]
3141 #[cfg(feature = "avro_custom_types")]
3142 fn duration_encoding_milliseconds() {
3143 let arr: PrimitiveArray<DurationMillisecondType> = vec![1i64, 0, -2].into();
3144 let mut expected = Vec::new();
3145 for v in [1i64, 0, -2] {
3146 expected.extend_from_slice(&avro_long_bytes(v));
3147 }
3148 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3149 assert_bytes_eq(&got, &expected);
3150 }
3151
3152 #[test]
3153 #[cfg(feature = "avro_custom_types")]
3154 fn duration_encoding_microseconds() {
3155 let arr: PrimitiveArray<DurationMicrosecondType> = vec![5i64, -6, 7].into();
3156 let mut expected = Vec::new();
3157 for v in [5i64, -6, 7] {
3158 expected.extend_from_slice(&avro_long_bytes(v));
3159 }
3160 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3161 assert_bytes_eq(&got, &expected);
3162 }
3163
3164 #[test]
3165 #[cfg(feature = "avro_custom_types")]
3166 fn duration_encoding_nanoseconds() {
3167 let arr: PrimitiveArray<DurationNanosecondType> = vec![8i64, 9, -10].into();
3168 let mut expected = Vec::new();
3169 for v in [8i64, 9, -10] {
3170 expected.extend_from_slice(&avro_long_bytes(v));
3171 }
3172 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3173 assert_bytes_eq(&got, &expected);
3174 }
3175
3176 #[test]
3177 fn duration_encoder_year_month_happy_path() {
3178 let arr: PrimitiveArray<IntervalYearMonthType> = vec![0i32, 1i32, 25i32].into();
3179 let mut expected = Vec::new();
3180 for m in [0u32, 1u32, 25u32] {
3181 expected.extend_from_slice(&duration_fixed12(m, 0, 0));
3182 }
3183 let got = encode_all(&arr, &FieldPlan::Duration, None);
3184 assert_bytes_eq(&got, &expected);
3185 }
3186
3187 #[test]
3188 fn duration_encoder_year_month_rejects_negative() {
3189 let arr: PrimitiveArray<IntervalYearMonthType> = vec![-1i32].into();
3190 let mut enc = FieldEncoder::make_encoder(&arr, &FieldPlan::Duration, None).unwrap();
3191 let mut out = Vec::new();
3192 let err = enc.encode(&mut out, 0).unwrap_err();
3193 match err {
3194 AvroError::InvalidArgument(msg) => {
3195 assert!(msg.contains("cannot encode negative months"))
3196 }
3197 other => panic!("expected InvalidArgument, got {other:?}"),
3198 }
3199 }
3200
3201 #[test]
3202 fn duration_encoder_day_time_happy_path() {
3203 let v0 = IntervalDayTimeType::make_value(2, 500); let v1 = IntervalDayTimeType::make_value(0, 0);
3205 let arr: PrimitiveArray<IntervalDayTimeType> = vec![v0, v1].into();
3206 let mut expected = Vec::new();
3207 expected.extend_from_slice(&duration_fixed12(0, 2, 500));
3208 expected.extend_from_slice(&duration_fixed12(0, 0, 0));
3209 let got = encode_all(&arr, &FieldPlan::Duration, None);
3210 assert_bytes_eq(&got, &expected);
3211 }
3212
3213 #[test]
3214 fn duration_encoder_day_time_rejects_negative() {
3215 let bad = IntervalDayTimeType::make_value(-1, 0);
3216 let arr: PrimitiveArray<IntervalDayTimeType> = vec![bad].into();
3217 let mut enc = FieldEncoder::make_encoder(&arr, &FieldPlan::Duration, None).unwrap();
3218 let mut out = Vec::new();
3219 let err = enc.encode(&mut out, 0).unwrap_err();
3220 match err {
3221 AvroError::InvalidArgument(msg) => {
3222 assert!(msg.contains("cannot encode negative days"))
3223 }
3224 other => panic!("expected InvalidArgument, got {other:?}"),
3225 }
3226 }
3227
3228 #[cfg(feature = "avro_custom_types")]
3229 #[test]
3230 fn interval_month_day_nano_fixed_encoder_happy_path() {
3231 let v0 = IntervalMonthDayNanoType::make_value(1, 2, 3); let v1 = IntervalMonthDayNanoType::make_value(-4, -5, -6);
3234 let arr: PrimitiveArray<IntervalMonthDayNanoType> = vec![v0, v1].into();
3235
3236 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3237 let mut expected = Vec::new();
3238 expected.extend_from_slice(&interval_mdn_fixed16(1, 2, 3));
3239 expected.extend_from_slice(&interval_mdn_fixed16(-4, -5, -6));
3240 assert_bytes_eq(&got, &expected);
3241 }
3242
3243 #[test]
3244 fn duration_encoder_month_day_nano_happy_path() {
3245 let v0 = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000); let v1 = IntervalMonthDayNanoType::make_value(0, 0, 0);
3247 let arr: PrimitiveArray<IntervalMonthDayNanoType> = vec![v0, v1].into();
3248 let mut expected = Vec::new();
3249 expected.extend_from_slice(&duration_fixed12(1, 2, 3));
3250 expected.extend_from_slice(&duration_fixed12(0, 0, 0));
3251 let got = encode_all(&arr, &FieldPlan::Duration, None);
3252 assert_bytes_eq(&got, &expected);
3253 }
3254
3255 #[test]
3256 fn duration_encoder_month_day_nano_rejects_non_ms_multiple() {
3257 let bad = IntervalMonthDayNanoType::make_value(0, 0, 1);
3258 let arr: PrimitiveArray<IntervalMonthDayNanoType> = vec![bad].into();
3259 let mut enc = FieldEncoder::make_encoder(&arr, &FieldPlan::Duration, None).unwrap();
3260 let mut out = Vec::new();
3261 let err = enc.encode(&mut out, 0).unwrap_err();
3262 match err {
3263 AvroError::InvalidArgument(msg) => {
3264 assert!(msg.contains("requires whole milliseconds") || msg.contains("divisible"))
3265 }
3266 other => panic!("expected InvalidArgument, got {other:?}"),
3267 }
3268 }
3269
3270 #[test]
3271 fn minimal_twos_complement_test() {
3272 let pos = [0x00, 0x00, 0x01];
3273 assert_eq!(minimal_twos_complement(&pos), &pos[2..]);
3274 let neg = [0xFF, 0xFF, 0x80]; assert_eq!(minimal_twos_complement(&neg), &neg[2..]);
3276 let zero = [0x00, 0x00, 0x00];
3277 assert_eq!(minimal_twos_complement(&zero), &zero[2..]);
3278 }
3279
3280 #[test]
3281 fn write_sign_extend_test() {
3282 let mut out = Vec::new();
3283 write_sign_extended(&mut out, &[0x01], 4).unwrap();
3284 assert_eq!(out, vec![0x00, 0x00, 0x00, 0x01]);
3285 out.clear();
3286 write_sign_extended(&mut out, &[0xFF], 4).unwrap();
3287 assert_eq!(out, vec![0xFF, 0xFF, 0xFF, 0xFF]);
3288 out.clear();
3289 write_sign_extended(&mut out, &[0xFF, 0xFF, 0x80], 2).unwrap();
3291 assert_eq!(out, vec![0xFF, 0x80]);
3292 out.clear();
3293 let err = write_sign_extended(&mut out, &[0x01, 0x00], 1).unwrap_err();
3295 match err {
3296 AvroError::InvalidArgument(_) => {}
3297 _ => panic!("expected InvalidArgument"),
3298 }
3299 }
3300
3301 #[test]
3302 fn duration_month_day_nano_overflow_millis() {
3303 let nanos = ((u64::from(u32::MAX) + 1) * 1_000_000) as i64;
3305 let v = IntervalMonthDayNanoType::make_value(0, 0, nanos);
3306 let arr: PrimitiveArray<IntervalMonthDayNanoType> = vec![v].into();
3307 let mut enc = FieldEncoder::make_encoder(&arr, &FieldPlan::Duration, None).unwrap();
3308 let mut out = Vec::new();
3309 let err = enc.encode(&mut out, 0).unwrap_err();
3310 match err {
3311 AvroError::InvalidArgument(msg) => assert!(msg.contains("exceed u32::MAX")),
3312 _ => panic!("expected InvalidArgument"),
3313 }
3314 }
3315
3316 #[test]
3317 fn fieldplan_decimal_precision_scale_mismatch_errors() {
3318 use crate::codec::Codec;
3320 use std::collections::HashMap;
3321 let arrow_field = Field::new("d", DataType::Decimal128(12, 2), true);
3322 let avro_dt = AvroDataType::new(Codec::Decimal(10, Some(2), None), HashMap::new(), None);
3323 let err = FieldPlan::build(&avro_dt, &arrow_field).unwrap_err();
3324 match err {
3325 AvroError::SchemaError(msg) => {
3326 assert!(msg.contains("Decimal precision/scale mismatch"))
3327 }
3328 _ => panic!("expected SchemaError"),
3329 }
3330 }
3331
3332 #[test]
3333 fn timestamp_micros_encoder() {
3334 test_scalar_primitive_encoding::<TimestampMicrosecondType>(
3336 &[
3337 1_704_067_200_000_000, 0, -123_456_789, ],
3341 &[None, Some(1_704_067_200_000_000)],
3342 );
3343 }
3344
3345 #[test]
3346 fn list_encoder_nullable_items_null_first() {
3347 let values = Int32Array::from(vec![Some(1), None, Some(2)]);
3349 let offsets = arrow_buffer::OffsetBuffer::new(vec![0i32, 3].into());
3350 let list = ListArray::new(
3351 Field::new("item", DataType::Int32, true).into(),
3352 offsets,
3353 Arc::new(values) as ArrayRef,
3354 None,
3355 );
3356
3357 let plan = FieldPlan::List {
3358 items_nullability: Some(Nullability::NullFirst),
3359 item_plan: Box::new(FieldPlan::Scalar),
3360 };
3361
3362 let mut expected = Vec::new();
3365 expected.extend(avro_long_bytes(3)); expected.extend(avro_long_bytes(1)); expected.extend(avro_long_bytes(1)); expected.extend(avro_long_bytes(0)); expected.extend(avro_long_bytes(1)); expected.extend(avro_long_bytes(2)); expected.extend(avro_long_bytes(0)); let got = encode_all(&list, &plan, None);
3374 assert_bytes_eq(&got, &expected);
3375 }
3376
3377 #[test]
3378 fn large_list_encoder_nullable_items_null_first() {
3379 let values = Int32Array::from(vec![Some(10), None]);
3381 let offsets = arrow_buffer::OffsetBuffer::new(vec![0i64, 2].into());
3382 let list = LargeListArray::new(
3383 Field::new("item", DataType::Int32, true).into(),
3384 offsets,
3385 Arc::new(values) as ArrayRef,
3386 None,
3387 );
3388
3389 let plan = FieldPlan::List {
3390 items_nullability: Some(Nullability::NullFirst),
3391 item_plan: Box::new(FieldPlan::Scalar),
3392 };
3393
3394 let mut expected = Vec::new();
3395 expected.extend(avro_long_bytes(2)); expected.extend(avro_long_bytes(1)); expected.extend(avro_long_bytes(10)); expected.extend(avro_long_bytes(0)); expected.extend(avro_long_bytes(0)); let got = encode_all(&list, &plan, None);
3402 assert_bytes_eq(&got, &expected);
3403 }
3404
3405 #[test]
3406 fn map_encoder_string_keys_nullable_int_values_null_first() {
3407 let keys = StringArray::from(vec!["k1", "k2"]);
3409 let values = Int32Array::from(vec![Some(7), None]);
3410
3411 let entries_fields = Fields::from(vec![
3412 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3413 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
3414 ]);
3415 let entries = StructArray::new(
3416 entries_fields,
3417 vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef],
3418 None,
3419 );
3420
3421 let offsets = arrow_buffer::OffsetBuffer::new(vec![0i32, 2].into());
3423 let map = MapArray::new(
3424 Field::new(
3425 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3426 entries.data_type().clone(),
3427 false,
3428 )
3429 .into(),
3430 offsets,
3431 entries,
3432 None,
3433 false,
3434 );
3435
3436 let plan = FieldPlan::Map {
3437 values_nullability: Some(Nullability::NullFirst),
3438 value_plan: Box::new(FieldPlan::Scalar),
3439 };
3440
3441 let mut expected = Vec::new();
3447 expected.extend(avro_long_bytes(2)); expected.extend(avro_len_prefixed_bytes(b"k1")); expected.extend(avro_long_bytes(1)); expected.extend(avro_long_bytes(7)); expected.extend(avro_len_prefixed_bytes(b"k2")); expected.extend(avro_long_bytes(0)); expected.extend(avro_long_bytes(0)); let got = encode_all(&map, &plan, None);
3456 assert_bytes_eq(&got, &expected);
3457 }
3458
3459 #[test]
3460 fn time32_seconds_to_millis_encoder() {
3461 let arr: arrow_array::PrimitiveArray<arrow_array::types::Time32SecondType> =
3463 vec![0i32, 1, -2, 12_345].into();
3464 let got = encode_all(&arr, &FieldPlan::TimeMillisFromSecs, None);
3465 let mut expected = Vec::new();
3466 for secs in [0i32, 1, -2, 12_345] {
3467 let millis = (secs as i64) * 1000;
3468 expected.extend_from_slice(&avro_long_bytes(millis));
3469 }
3470 assert_bytes_eq(&got, &expected);
3471 }
3472
3473 #[test]
3474 fn time32_seconds_to_millis_overflow() {
3475 let overflow_secs: i32 = i32::MAX / 1000 + 1;
3477 let arr: PrimitiveArray<Time32SecondType> = vec![overflow_secs].into();
3478 let mut enc =
3479 FieldEncoder::make_encoder(&arr, &FieldPlan::TimeMillisFromSecs, None).unwrap();
3480 let mut out = Vec::new();
3481 let err = enc.encode(&mut out, 0).unwrap_err();
3482 match err {
3483 AvroError::InvalidArgument(msg) => {
3484 assert!(
3485 msg.contains("overflowed") || msg.contains("overflow"),
3486 "unexpected message: {msg}"
3487 )
3488 }
3489 other => panic!("expected InvalidArgument, got {other:?}"),
3490 }
3491 }
3492
3493 #[test]
3494 fn time32_seconds_to_millis_type_mismatch_returns_schema_error() {
3495 let arr = Int32Array::from(vec![1, 2, 3]);
3496 match FieldEncoder::make_encoder(&arr, &FieldPlan::TimeMillisFromSecs, None) {
3497 Err(AvroError::SchemaError(msg)) => {
3498 assert!(msg.contains("Time32(Second)"), "unexpected message: {msg}");
3499 assert!(msg.contains("Int32"), "unexpected message: {msg}");
3500 }
3501 Ok(_) => panic!("expected SchemaError"),
3502 Err(other) => panic!("expected SchemaError, got {other:?}"),
3503 }
3504 }
3505
3506 #[test]
3507 fn encode_rows_time32_seconds_plan_rejects_millisecond_column() {
3508 let schema = ArrowSchema::new(vec![Field::new(
3509 "t",
3510 DataType::Time32(TimeUnit::Millisecond),
3511 false,
3512 )]);
3513 let arr: PrimitiveArray<Time32MillisecondType> = vec![1i32].into();
3514 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arr)]).unwrap();
3515 let encoder = RecordEncoder {
3516 columns: vec![FieldBinding {
3517 arrow_index: 0,
3518 nullability: None,
3519 plan: FieldPlan::TimeMillisFromSecs,
3520 }],
3521 prefix: None,
3522 };
3523
3524 let mut out = BytesMut::new();
3525 let mut offsets = vec![0usize];
3526 let err = encoder
3527 .encode_rows(&batch, 16, &mut out, &mut offsets)
3528 .unwrap_err();
3529 match err {
3530 AvroError::SchemaError(msg) => {
3531 assert!(msg.contains("Time32(Second)"), "unexpected message: {msg}");
3532 assert!(
3533 msg.contains("Time32(Millisecond)"),
3534 "unexpected message: {msg}"
3535 );
3536 }
3537 other => panic!("expected SchemaError, got {other:?}"),
3538 }
3539 }
3540
3541 #[cfg(not(feature = "avro_custom_types"))]
3542 #[test]
3543 fn timestamp_seconds_to_millis_encoder() {
3544 let arr: PrimitiveArray<TimestampSecondType> = vec![0i64, 1, -1, 1_234_567_890].into();
3546 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3547 let mut expected = Vec::new();
3548 for secs in [0i64, 1, -1, 1_234_567_890] {
3549 let millis = secs * 1000;
3550 expected.extend_from_slice(&avro_long_bytes(millis));
3551 }
3552 assert_bytes_eq(&got, &expected);
3553 }
3554
3555 #[cfg(not(feature = "avro_custom_types"))]
3556 #[test]
3557 fn timestamp_seconds_to_millis_overflow() {
3558 let overflow_secs: i64 = i64::MAX / 1000 + 1;
3560 let arr: PrimitiveArray<TimestampSecondType> = vec![overflow_secs].into();
3561 let mut enc = FieldEncoder::make_encoder(&arr, &FieldPlan::Scalar, None).unwrap();
3562 let mut out = Vec::new();
3563 let err = enc.encode(&mut out, 0).unwrap_err();
3564 match err {
3565 AvroError::InvalidArgument(msg) => {
3566 assert!(
3567 msg.contains("overflowed") || msg.contains("overflow"),
3568 "unexpected message: {msg}"
3569 )
3570 }
3571 other => panic!("expected InvalidArgument, got {other:?}"),
3572 }
3573 }
3574
3575 #[test]
3576 fn timestamp_nanos_encoder() {
3577 let arr: PrimitiveArray<TimestampNanosecondType> = vec![0i64, 1, -1, 123].into();
3578 let got = encode_all(&arr, &FieldPlan::Scalar, None);
3579 let mut expected = Vec::new();
3580 for ns in [0i64, 1, -1, 123] {
3581 expected.extend_from_slice(&avro_long_bytes(ns));
3582 }
3583 assert_bytes_eq(&got, &expected);
3584 }
3585
3586 #[test]
3587 fn union_encoder_string_int_nonzero_type_ids() {
3588 let strings = StringArray::from(vec!["hello", "world"]);
3589 let ints = Int32Array::from(vec![10, 20, 30]);
3590 let union_fields = UnionFields::try_new(
3591 vec![2, 5],
3592 vec![
3593 Field::new("v_str", DataType::Utf8, true),
3594 Field::new("v_int", DataType::Int32, true),
3595 ],
3596 )
3597 .unwrap();
3598 let type_ids = Buffer::from_slice_ref([2_i8, 5, 5, 2, 5]);
3599 let offsets = Buffer::from_slice_ref([0_i32, 0, 1, 1, 2]);
3600 let union_array = UnionArray::try_new(
3601 union_fields,
3602 type_ids.into(),
3603 Some(offsets.into()),
3604 vec![Arc::new(strings), Arc::new(ints)],
3605 )
3606 .unwrap();
3607 let plan = FieldPlan::Union {
3608 bindings: vec![
3609 FieldBinding {
3610 arrow_index: 0,
3611 nullability: None,
3612 plan: FieldPlan::Scalar,
3613 },
3614 FieldBinding {
3615 arrow_index: 1,
3616 nullability: None,
3617 plan: FieldPlan::Scalar,
3618 },
3619 ],
3620 };
3621 let got = encode_all(&union_array, &plan, None);
3622 let mut expected = Vec::new();
3623 expected.extend(avro_long_bytes(0));
3624 expected.extend(avro_len_prefixed_bytes(b"hello"));
3625 expected.extend(avro_long_bytes(1));
3626 expected.extend(avro_long_bytes(10));
3627 expected.extend(avro_long_bytes(1));
3628 expected.extend(avro_long_bytes(20));
3629 expected.extend(avro_long_bytes(0));
3630 expected.extend(avro_len_prefixed_bytes(b"world"));
3631 expected.extend(avro_long_bytes(1));
3632 expected.extend(avro_long_bytes(30));
3633 assert_bytes_eq(&got, &expected);
3634 }
3635
3636 #[test]
3637 fn nullable_state_with_null_buffer_and_zero_nulls() {
3638 let values = vec![1i32, 2, 3];
3639 let arr = Int32Array::from_iter_values_with_nulls(values, Some(NullBuffer::new_valid(3)));
3640 assert_eq!(arr.null_count(), 0);
3641 assert!(arr.nulls().is_some());
3642 let plan = FieldPlan::Scalar;
3643 let enc = FieldEncoder::make_encoder(&arr, &plan, Some(Nullability::NullFirst)).unwrap();
3644 match enc.null_state {
3645 NullState::NullableNoNulls { union_value_byte } => {
3646 assert_eq!(
3647 union_value_byte,
3648 union_value_branch_byte(Nullability::NullFirst, false)
3649 );
3650 }
3651 other => panic!("expected NullableNoNulls, got {other:?}"),
3652 }
3653 }
3654
3655 #[test]
3656 fn encode_rows_single_column_int32() {
3657 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
3658 let arr = Int32Array::from(vec![1, 2, 3]);
3659 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3660 let encoder = RecordEncoder {
3661 columns: vec![FieldBinding {
3662 arrow_index: 0,
3663 nullability: None,
3664 plan: FieldPlan::Scalar,
3665 }],
3666 prefix: None,
3667 };
3668 let mut out = BytesMut::new();
3669 let mut offsets: Vec<usize> = vec![0];
3670 encoder
3671 .encode_rows(&batch, 16, &mut out, &mut offsets)
3672 .unwrap();
3673 assert_eq!(offsets.len(), 4);
3674 assert_eq!(*offsets.last().unwrap(), out.len());
3675 assert_bytes_eq(row_slice(&out, &offsets, 0), &avro_long_bytes(1));
3676 assert_bytes_eq(row_slice(&out, &offsets, 1), &avro_long_bytes(2));
3677 assert_bytes_eq(row_slice(&out, &offsets, 2), &avro_long_bytes(3));
3678 }
3679
3680 #[test]
3681 fn encode_rows_multiple_columns() {
3682 let schema = ArrowSchema::new(vec![
3683 Field::new("a", DataType::Int32, false),
3684 Field::new("b", DataType::Utf8, false),
3685 ]);
3686 let int_arr = Int32Array::from(vec![10, 20]);
3687 let str_arr = StringArray::from(vec!["hello", "world"]);
3688 let batch = RecordBatch::try_new(
3689 Arc::new(schema.clone()),
3690 vec![Arc::new(int_arr), Arc::new(str_arr)],
3691 )
3692 .unwrap();
3693 let encoder = RecordEncoder {
3694 columns: vec![
3695 FieldBinding {
3696 arrow_index: 0,
3697 nullability: None,
3698 plan: FieldPlan::Scalar,
3699 },
3700 FieldBinding {
3701 arrow_index: 1,
3702 nullability: None,
3703 plan: FieldPlan::Scalar,
3704 },
3705 ],
3706 prefix: None,
3707 };
3708 let mut out = BytesMut::new();
3709 let mut offsets: Vec<usize> = vec![0];
3710 encoder
3711 .encode_rows(&batch, 32, &mut out, &mut offsets)
3712 .unwrap();
3713 assert_eq!(offsets.len(), 3);
3714 assert_eq!(*offsets.last().unwrap(), out.len());
3715 let mut expected_row0 = Vec::new();
3716 expected_row0.extend(avro_long_bytes(10));
3717 expected_row0.extend(avro_len_prefixed_bytes(b"hello"));
3718 assert_bytes_eq(row_slice(&out, &offsets, 0), &expected_row0);
3719 let mut expected_row1 = Vec::new();
3720 expected_row1.extend(avro_long_bytes(20));
3721 expected_row1.extend(avro_len_prefixed_bytes(b"world"));
3722 assert_bytes_eq(row_slice(&out, &offsets, 1), &expected_row1);
3723 }
3724
3725 #[test]
3726 fn encode_rows_with_prefix() {
3727 use crate::codec::AvroFieldBuilder;
3728 use crate::schema::AvroSchema;
3729 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
3730 let arr = Int32Array::from(vec![42]);
3731 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3732 let avro_schema = AvroSchema::try_from(&schema).unwrap();
3733 let fingerprint = avro_schema
3734 .fingerprint(crate::schema::FingerprintAlgorithm::Rabin)
3735 .unwrap();
3736 let avro_root = AvroFieldBuilder::new(&avro_schema.schema().unwrap())
3737 .build()
3738 .unwrap();
3739 let encoder = RecordEncoderBuilder::new(&avro_root, &schema)
3740 .with_fingerprint(Some(fingerprint))
3741 .build()
3742 .unwrap();
3743 let mut out = BytesMut::new();
3744 let mut offsets: Vec<usize> = vec![0];
3745 encoder
3746 .encode_rows(&batch, 32, &mut out, &mut offsets)
3747 .unwrap();
3748 assert_eq!(offsets.len(), 2);
3749 let row0 = row_slice(&out, &offsets, 0);
3750 assert!(row0.len() > 10, "Row should contain prefix + encoded value");
3751 assert_eq!(row0[0], 0xC3);
3752 assert_eq!(row0[1], 0x01);
3753 }
3754
3755 #[test]
3756 fn encode_rows_empty_batch() {
3757 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
3758 let arr = Int32Array::from(Vec::<i32>::new());
3759 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3760 let encoder = RecordEncoder {
3761 columns: vec![FieldBinding {
3762 arrow_index: 0,
3763 nullability: None,
3764 plan: FieldPlan::Scalar,
3765 }],
3766 prefix: None,
3767 };
3768 let mut out = BytesMut::new();
3769 let mut offsets: Vec<usize> = vec![0];
3770 encoder
3771 .encode_rows(&batch, 16, &mut out, &mut offsets)
3772 .unwrap();
3773 assert_eq!(offsets, vec![0]);
3774 assert!(out.is_empty());
3775 }
3776
3777 #[test]
3778 fn encode_rows_matches_encode_output() {
3779 let schema = ArrowSchema::new(vec![
3780 Field::new("a", DataType::Int64, false),
3781 Field::new("b", DataType::Float64, false),
3782 ]);
3783 let int_arr = Int64Array::from(vec![100i64, 200, 300]);
3784 let float_arr = Float64Array::from(vec![1.5, 2.5, 3.5]);
3785 let batch = RecordBatch::try_new(
3786 Arc::new(schema.clone()),
3787 vec![Arc::new(int_arr), Arc::new(float_arr)],
3788 )
3789 .unwrap();
3790 let encoder = RecordEncoder {
3791 columns: vec![
3792 FieldBinding {
3793 arrow_index: 0,
3794 nullability: None,
3795 plan: FieldPlan::Scalar,
3796 },
3797 FieldBinding {
3798 arrow_index: 1,
3799 nullability: None,
3800 plan: FieldPlan::Scalar,
3801 },
3802 ],
3803 prefix: None,
3804 };
3805 let mut stream_buf = Vec::new();
3806 encoder.encode(&mut stream_buf, &batch).unwrap();
3807 let mut out = BytesMut::new();
3808 let mut offsets: Vec<usize> = vec![0];
3809 encoder
3810 .encode_rows(&batch, 32, &mut out, &mut offsets)
3811 .unwrap();
3812 assert_eq!(offsets.len(), 1 + batch.num_rows());
3813 assert_bytes_eq(&out[..], &stream_buf);
3814 }
3815
3816 #[test]
3817 fn encode_rows_appends_to_existing_buffer() {
3818 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
3819 let arr = Int32Array::from(vec![5, 6]);
3820 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3821 let encoder = RecordEncoder {
3822 columns: vec![FieldBinding {
3823 arrow_index: 0,
3824 nullability: None,
3825 plan: FieldPlan::Scalar,
3826 }],
3827 prefix: None,
3828 };
3829 let mut out = BytesMut::new();
3830 out.extend_from_slice(&[0xAA, 0xBB]);
3831 let mut offsets: Vec<usize> = vec![0, out.len()];
3832 encoder
3833 .encode_rows(&batch, 16, &mut out, &mut offsets)
3834 .unwrap();
3835 assert_eq!(offsets.len(), 4);
3836 assert_eq!(*offsets.last().unwrap(), out.len());
3837 assert_bytes_eq(row_slice(&out, &offsets, 0), &[0xAA, 0xBB]);
3838 assert_bytes_eq(row_slice(&out, &offsets, 1), &avro_long_bytes(5));
3839 assert_bytes_eq(row_slice(&out, &offsets, 2), &avro_long_bytes(6));
3840 }
3841
3842 #[test]
3843 fn encode_rows_nullable_column() {
3844 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, true)]);
3845 let arr = Int32Array::from(vec![Some(1), None, Some(3)]);
3846 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3847 let encoder = RecordEncoder {
3848 columns: vec![FieldBinding {
3849 arrow_index: 0,
3850 nullability: Some(Nullability::NullFirst),
3851 plan: FieldPlan::Scalar,
3852 }],
3853 prefix: None,
3854 };
3855 let mut out = BytesMut::new();
3856 let mut offsets: Vec<usize> = vec![0];
3857 encoder
3858 .encode_rows(&batch, 16, &mut out, &mut offsets)
3859 .unwrap();
3860 assert_eq!(offsets.len(), 4);
3861 let mut expected_row0 = Vec::new();
3862 expected_row0.extend(avro_long_bytes(1)); expected_row0.extend(avro_long_bytes(1)); assert_bytes_eq(row_slice(&out, &offsets, 0), &expected_row0);
3865 let expected_row1 = avro_long_bytes(0); assert_bytes_eq(row_slice(&out, &offsets, 1), &expected_row1);
3867 let mut expected_row2 = Vec::new();
3868 expected_row2.extend(avro_long_bytes(1)); expected_row2.extend(avro_long_bytes(3)); assert_bytes_eq(row_slice(&out, &offsets, 2), &expected_row2);
3871 }
3872
3873 #[test]
3874 fn encode_prefix_write_error() {
3875 use crate::codec::AvroFieldBuilder;
3876 use crate::schema::{AvroSchema, FingerprintAlgorithm};
3877 use std::io;
3878
3879 struct FailWriter {
3880 failed: bool,
3881 }
3882
3883 impl io::Write for FailWriter {
3884 fn write(&mut self, _buf: &[u8]) -> io::Result<usize> {
3885 if !self.failed {
3886 self.failed = true;
3887 Err(io::Error::other("fail write"))
3888 } else {
3889 Ok(0)
3890 }
3891 }
3892
3893 fn flush(&mut self) -> io::Result<()> {
3894 Ok(())
3895 }
3896 }
3897
3898 let schema = ArrowSchema::new(vec![Field::new("x", DataType::Int32, false)]);
3899 let arr = Int32Array::from(vec![42]);
3900 let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(arr)]).unwrap();
3901 let avro_schema = AvroSchema::try_from(&schema).unwrap();
3902 let fingerprint = avro_schema
3903 .fingerprint(FingerprintAlgorithm::Rabin)
3904 .unwrap();
3905 let avro_root = AvroFieldBuilder::new(&avro_schema.schema().unwrap())
3906 .build()
3907 .unwrap();
3908 let encoder = RecordEncoderBuilder::new(&avro_root, &schema)
3909 .with_fingerprint(Some(fingerprint))
3910 .build()
3911 .unwrap();
3912
3913 let mut writer = FailWriter { failed: false };
3914 let err = encoder.encode(&mut writer, &batch).unwrap_err();
3915 let msg = format!("{err}");
3916 assert!(msg.contains("write prefix"), "unexpected error: {msg}");
3917 }
3918}