alamb commented on code in PR #10670:
URL: https://github.com/apache/arrow-rs/pull/10670#discussion_r3875145595
##########
arrow-json/src/reader/mod.rs:
##########
@@ -704,11 +720,133 @@ impl Decoder {
}
}
-trait ArrayDecoder: Send {
+/// Decodes a column of JSON values from a [`Tape`] into an [`ArrayRef`]
+///
+/// Implement this together with [`DecoderFactory`] to override how a type is
+/// decoded, or to add support for a type the reader does not handle.
+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. 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 default decoder for a data type, or adds support for one the
reader
+/// 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,
+/// the inverse of the [`EncoderFactory`] example.
+///
+/// ```
+/// 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, 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_custom_decoder(
+/// &self,
+/// _ctx: &DecoderContext,
+/// data_type: &DataType,
+/// _is_nullable: bool,
+/// ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
+/// match 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 `data_type`, or `Ok(None)` to use the reader's
default.
+ ///
+ /// Use [`DecoderContext::make_decoder`] on `ctx` to build child decoders.
Calling
+ /// it for the `data_type` this was invoked with recurses back here and
loops.
+ ///
+ /// `is_nullable` folds in ancestor nullability, so it may be `true` even
where the
+ /// corresponding field is not.
+ fn make_custom_decoder(
+ &self,
+ _ctx: &DecoderContext,
+ _data_type: &DataType,
Review Comment:
Rather than pass in the fields of a DataType, how about just pass in the
`&Field` reference directly:
That would also allow the decoder to key off Metadata (where the extension
type information is stored)
```diff
@@ -861,17 +864,24 @@ pub trait ArrayDecoder: Send {
///
/// [`EncoderFactory`]: crate::EncoderFactory
pub trait DecoderFactory: std::fmt::Debug + Send + Sync {
- /// Make a decoder for `data_type`, or `Ok(None)` to use the reader's
default.
+ /// 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. Calling
- /// it for the `data_type` this was invoked with recurses back here and
loops.
+ /// it with the field this was invoked with recurses back here and
loops.
///
- /// `is_nullable` folds in ancestor nullability, so it may be `true`
even where the
- /// corresponding field is not.
+ /// `is_nullable` folds in ancestor nullability, so it may be `true`
even when
+ /// `field.is_nullable()` is not.
+ ///
+ /// [extension types]:
https://arrow.apache.org/docs/format/Columnar.html#extension-types
fn make_custom_decoder(
&self,
_ctx: &DecoderContext,
- _data_type: &DataType,
+ _field: &FieldRef,
_is_nullable: bool,
) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
Ok(None)
```
##########
arrow-json/src/reader/mod.rs:
##########
@@ -745,11 +885,18 @@ impl DecoderContext {
self.ignore_type_conflicts
}
+ /// Returns the optional hook for customizing decoding behavior
+ pub fn decoder_factory(&self) -> Option<&Arc<dyn DecoderFactory>> {
+ self.decoder_factory.as_ref()
+ }
+
/// Create a decoder for a type.
///
- /// This is the standard way to create child decoders from within a decoder
- /// implementation.
- fn make_decoder(
+ /// The standard way to create child decoders from within a decoder, and
how a
+ /// [`DecoderFactory`] delegates to the decoder the reader would otherwise
use.
+ /// The factory is consulted first, so calling this for the same data type
the
Review Comment:
this seems like a non trivial foot gun and would be nice to have a way to
get the standard decoder from within the factor (e.g. to delegate to the
default decoders)
Could we add some method like `make_default_decoder` that bypasses the
Decoder factor?
##########
arrow-json/src/reader/mod.rs:
##########
@@ -704,11 +720,133 @@ impl Decoder {
}
}
-trait ArrayDecoder: Send {
+/// Decodes a column of JSON values from a [`Tape`] into an [`ArrayRef`]
+///
+/// Implement this together with [`DecoderFactory`] to override how a type is
+/// decoded, or to add support for a type the reader does not handle.
+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. 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 default decoder for a data type, or adds support for one the
reader
+/// 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,
+/// the inverse of the [`EncoderFactory`] example.
+///
+/// ```
+/// 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, 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_custom_decoder(
+/// &self,
+/// _ctx: &DecoderContext,
+/// data_type: &DataType,
+/// _is_nullable: bool,
+/// ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
+/// match 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 `data_type`, or `Ok(None)` to use the reader's
default.
+ ///
+ /// Use [`DecoderContext::make_decoder`] on `ctx` to build child decoders.
Calling
+ /// it for the `data_type` this was invoked with recurses back here and
loops.
+ ///
+ /// `is_nullable` folds in ancestor nullability, so it may be `true` even
where the
+ /// corresponding field is not.
+ fn make_custom_decoder(
Review Comment:
It would be nice to make this name symmetric with the writer side -
https://github.com/apache/arrow-rs/blob/4cc296bca10303382989ec7cebc98522e3230ddb/arrow-json/src/writer/encoder.rs#L251
So something like
```rust
fn make_default_decoder(
```
--
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]
