Skip to content

Commit d96fc62

Browse files
committed
Addressed PR Review comment and added a new encoder.rs unit test that locks in the desired zero-null check behavior.
Also implemented a small and non-breaking/non-public-facing micro-optimization that removed `cloned` from the updated zero nulls checking logic by adding the `'a` lifetime to `NullState`.
1 parent 8ffb25a commit d96fc62

File tree

1 file changed

+28
-10
lines changed

1 file changed

+28
-10
lines changed

arrow-avro/src/writer/encoder.rs

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -205,13 +205,13 @@ fn write_optional_index<W: Write + ?Sized>(
205205
}
206206

207207
#[derive(Debug, Clone)]
208-
enum NullState {
208+
enum NullState<'a> {
209209
NonNullable,
210210
NullableNoNulls {
211211
union_value_byte: u8,
212212
},
213213
Nullable {
214-
nulls: NullBuffer,
214+
nulls: &'a NullBuffer,
215215
null_order: Nullability,
216216
},
217217
}
@@ -221,7 +221,7 @@ enum NullState {
221221
/// - Carries the per-site nullability **state** as a single enum that enforces invariants
222222
pub(crate) struct FieldEncoder<'a> {
223223
encoder: Encoder<'a>,
224-
null_state: NullState,
224+
null_state: NullState<'a>,
225225
}
226226

227227
impl<'a> FieldEncoder<'a> {
@@ -617,16 +617,15 @@ impl<'a> FieldEncoder<'a> {
617617
// Compute the effective null state from writer-declared nullability and data nulls.
618618
let null_state = match nullability {
619619
None => NullState::NonNullable,
620-
Some(order) => {
621-
if let Some(nulls) = array.nulls().cloned() {
622-
NullState::Nullable {
623-
nulls,
624-
null_order: order,
625-
}
620+
Some(null_order) => {
621+
if array.null_count() > 0
622+
&& let Some(nulls) = array.nulls()
623+
{
624+
NullState::Nullable { nulls, null_order }
626625
} else {
627626
// Nullable site with no null buffer for this view
628627
NullState::NullableNoNulls {
629-
union_value_byte: union_value_branch_byte(order, false),
628+
union_value_byte: union_value_branch_byte(null_order, false),
630629
}
631630
}
632631
}
@@ -3025,4 +3024,23 @@ mod tests {
30253024
expected.extend(avro_long_bytes(30));
30263025
assert_bytes_eq(&got, &expected);
30273026
}
3027+
3028+
#[test]
3029+
fn nullable_state_with_null_buffer_and_zero_nulls() {
3030+
let values = vec![1i32, 2, 3];
3031+
let arr = Int32Array::from_iter_values_with_nulls(values, Some(NullBuffer::new_valid(3)));
3032+
assert_eq!(arr.null_count(), 0);
3033+
assert!(arr.nulls().is_some());
3034+
let plan = FieldPlan::Scalar;
3035+
let enc = FieldEncoder::make_encoder(&arr, &plan, Some(Nullability::NullFirst)).unwrap();
3036+
match enc.null_state {
3037+
NullState::NullableNoNulls { union_value_byte } => {
3038+
assert_eq!(
3039+
union_value_byte,
3040+
union_value_branch_byte(Nullability::NullFirst, false)
3041+
);
3042+
}
3043+
other => panic!("expected NullableNoNulls, got {other:?}"),
3044+
}
3045+
}
30283046
}

0 commit comments

Comments
 (0)