hareshkh commented on code in PR #10670:
URL: https://github.com/apache/arrow-rs/pull/10670#discussion_r3882377405
##########
arrow-json/src/reader/mod.rs:
##########
@@ -735,11 +748,171 @@ impl Decoder {
}
}
-trait ArrayDecoder: Send {
+/// Decodes a column of JSON values from a [`Tape`] into an [`ArrayRef`]
+///
+/// Implement alongside [`DecoderFactory`] to override or add support for a
type.
+pub trait ArrayDecoder: Send {
/// Decode elements from `tape` starting at the indexes contained in `pos`
+ ///
+ /// `pos` contains one tape index per output row, so the returned array
must have
+ /// exactly `pos.len()` elements and the field's data type. A row's value
may be
+ /// [`TapeElement::Null`].
fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef,
ArrowError>;
}
+/// A trait to create custom decoders for specific data types.
+///
+/// Overrides the reader's decoder for a data type, or adds support for one it
does not
+/// handle. The reader-side counterpart of [`EncoderFactory`]; register an
+/// implementation with [`ReaderBuilder::with_decoder_factory`].
+///
+/// # Examples
+///
+/// Decodes `Binary` from a JSON array of integers rather than the default hex
string.
+///
+/// ```
+/// use std::sync::Arc;
+/// use arrow_array::{Array, ArrayRef, BinaryArray};
+/// use arrow_array::cast::AsArray;
+/// use arrow_json::reader::{ArrayDecoder, DecoderContext, DecoderFactory,
Tape, TapeElement};
+/// use arrow_json::ReaderBuilder;
+/// use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema};
+///
+/// /// Decodes `[104, 105]` into the bytes `b"hi"`
+/// struct IntArrayBinaryDecoder;
+///
+/// impl ArrayDecoder for IntArrayBinaryDecoder {
+/// fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef,
ArrowError> {
+/// let mut values: Vec<Option<Vec<u8>>> =
Vec::with_capacity(pos.len());
+/// for p in pos {
+/// match tape.get(*p) {
+/// TapeElement::Null => values.push(None),
+/// TapeElement::StartList(end) => {
+/// let mut bytes = Vec::new();
+/// let mut cur = p + 1;
+/// while cur < end {
+/// match tape.get(cur) {
+/// // JSON text yields `Number`; serde yields
`I32`
+/// TapeElement::Number(idx) => {
+/// let s = tape.get_string(idx);
+/// bytes.push(s.parse::<u8>().map_err(|e| {
+/// ArrowError::JsonError(format!("invalid
byte {s}: {e}"))
+/// })?);
+/// }
+/// TapeElement::I32(v) => bytes.push(v as u8),
+/// _ => return Err(tape.error(cur, "byte")),
+/// }
+/// cur = tape.next(cur, "byte")?;
+/// }
+/// values.push(Some(bytes));
+/// }
+/// _ => return Err(tape.error(*p, "list of bytes")),
+/// }
+/// }
+/// Ok(Arc::new(BinaryArray::from_iter(values.iter().map(|v|
v.as_deref()))))
+/// }
+/// }
+///
+/// #[derive(Debug)]
+/// struct IntArrayBinaryDecoderFactory;
+///
+/// impl DecoderFactory for IntArrayBinaryDecoderFactory {
+/// fn make_default_decoder(
+/// &self,
+/// _ctx: &DecoderContext,
+/// field: &FieldRef,
+/// _is_nullable: bool,
+/// ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
+/// // `field.metadata()` is also available, e.g. for extension types
+/// match field.data_type() {
+/// DataType::Binary => Ok(Some(Box::new(IntArrayBinaryDecoder))),
+/// // Returning `None` uses the reader's default decoder
+/// _ => Ok(None),
+/// }
+/// }
+/// }
+///
+/// let schema = Arc::new(Schema::new(vec![
+/// Field::new("bytes", DataType::Binary, true),
+/// Field::new("float", DataType::Float64, true),
+/// ]));
+///
+/// let json = r#"{"bytes": [104, 105], "float": 1.0}
+/// {"float": 2.3}
+/// {"bytes": [98]}
+/// "#;
+///
+/// let batch = ReaderBuilder::new(schema)
+/// .with_decoder_factory(Arc::new(IntArrayBinaryDecoderFactory))
+/// .build(json.as_bytes())
+/// .unwrap()
+/// .next()
+/// .unwrap()
+/// .unwrap();
+///
+/// let bytes = batch.column(0).as_binary::<i32>();
+/// assert_eq!(bytes.value(0), b"hi");
+/// assert!(bytes.is_null(1));
+/// assert_eq!(bytes.value(2), b"b");
+/// ```
+///
+/// [`EncoderFactory`]: crate::EncoderFactory
+pub trait DecoderFactory: std::fmt::Debug + Send + Sync {
+ /// Make a decoder for `field`, or `Ok(None)` to use the reader's default.
+ ///
+ /// Receives the [`FieldRef`] rather than just its [`DataType`] so decoder
+ /// selection can consider the field's metadata, e.g. to identify
[extension
+ /// types]. The root of a [`ReaderBuilder::new`] schema is presented as a
+ /// synthesized nameless `Struct` field.
+ ///
+ /// Use [`DecoderContext::make_decoder`] on `ctx` to build child decoders,
and
+ /// [`DecoderContext::make_builtin_decoder`] to build on the reader's own
decoder
+ /// for this field. Calling `make_decoder` with the field this was invoked
with
+ /// recurses back here and loops.
+ ///
+ /// `is_nullable` folds in ancestor nullability, so it may differ from
+ /// `field.is_nullable()` in either direction: a nullable struct widens its
+ /// children, while a run-end encoded array narrows its values.
+ ///
+ /// [extension types]:
https://arrow.apache.org/docs/format/Columnar.html#extension-types
+ fn make_default_decoder(
+ &self,
+ _ctx: &DecoderContext,
+ _field: &FieldRef,
+ _is_nullable: bool,
+ ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
+ Ok(None)
+ }
+}
+
+/// Validates the output of a [`DecoderFactory`] decoder before it reaches the
arrays
+/// built from it, some of which are constructed without further checks.
+struct CheckedDecoder {
+ inner: Box<dyn ArrayDecoder>,
+ data_type: DataType,
+}
+
+impl ArrayDecoder for CheckedDecoder {
+ fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef,
ArrowError> {
+ let array = self.inner.decode(tape, pos)?;
+ if array.data_type() != &self.data_type {
+ return Err(ArrowError::JsonError(format!(
+ "custom decoder returned {} for a field of type {}",
+ array.data_type(),
+ self.data_type
+ )));
+ }
+ if array.len() != pos.len() {
+ return Err(ArrowError::JsonError(format!(
+ "custom decoder returned {} values for {} rows",
+ array.len(),
+ pos.len()
+ )));
+ }
+ Ok(array)
+ }
+}
Review Comment:
@alamb: This is net new from the previous pass and adds validation to custom
decoders - what are your thoughts on this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
