From 3ddeacaf0a20831c78cb14d5b8ed56825bbfbe39 Mon Sep 17 00:00:00 2001 From: Jefffrey Date: Mon, 23 Sep 2024 19:13:47 +1000 Subject: [PATCH] Refactor PrimitiveArrayDecoder to fully use PrimitiveValueDecoder::decode --- src/array_decoder/mod.rs | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/array_decoder/mod.rs b/src/array_decoder/mod.rs index 1d6c734f..517e5bcd 100644 --- a/src/array_decoder/mod.rs +++ b/src/array_decoder/mod.rs @@ -34,7 +34,8 @@ use crate::encoding::byte::ByteRleReader; use crate::encoding::float::FloatIter; use crate::encoding::{get_rle_reader, PrimitiveValueDecoder}; use crate::error::{ - self, ArrowSnafu, MismatchedSchemaSnafu, Result, UnexpectedSnafu, UnsupportedTypeVariantSnafu, + self, ArrowSnafu, MismatchedSchemaSnafu, OutOfSpecSnafu, Result, UnexpectedSnafu, + UnsupportedTypeVariantSnafu, }; use crate::proto::stream::Kind; use crate::schema::DataType; @@ -78,16 +79,22 @@ impl PrimitiveArrayDecoder { match present { Some(present) => { + let count = present.iter().filter(|&&p| p).count(); + let mut data = vec![T::Native::ZERO; count]; + if self.iter.decode(data.as_mut_slice())? != count { + // TODO: use a more specific error type, be more descriptive + return OutOfSpecSnafu { + msg: "Not enough data in array", + } + .fail(); + } + + let mut data = data.iter(); let mut builder = PrimitiveBuilder::::with_capacity(batch_size); for is_present in present { if is_present { - // TODO: return as error instead - let val = self - .iter - .next() - .transpose()? - .expect("array less than expected length"); - builder.append_value(val); + // Safe unwrap as we guarantee length above + builder.append_value(*data.next().unwrap()); } else { builder.append_null(); }