Implement arrow_json encoder for Decimal128 & Decimal256 (#6606)

* Implement arrow_json encoder for Decimal128 & Decimal256

* fix

* Use primitive helper for i128

* Revert "Use primitive helper for i128"

This reverts commit dfe8edb.

* Use ArrayFormatter directly

---------

Co-authored-by: Raphael Taylor-Davies <[email protected]>

arrow-json/src/writer/encoder.rs

@@ -134,6 +134,11 @@ fn make_encoder_impl<'a>(
             };
             (Box::new(encoder) as _, array.nulls().cloned())
         }
+        DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => {
+            let options = FormatOptions::new().with_display_error(true);
+            let formatter = ArrayFormatter::try_new(array, &options)?;
+            (Box::new(RawArrayFormatter(formatter)) as _, array.nulls().cloned())
+        }
         d => match d.is_temporal() {
             true => {
                 // Note: the implementation of Encoder for ArrayFormatter assumes it does not produce
@@ -434,6 +439,15 @@ impl<'a> Encoder for ArrayFormatter<'a> {
     }
 }
 
+/// A newtype wrapper around [`ArrayFormatter`] that skips surrounding the value with `"`
+struct RawArrayFormatter<'a>(ArrayFormatter<'a>);
+
+impl<'a> Encoder for RawArrayFormatter<'a> {
+    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
+        let _ = write!(out, "{}", self.0.value(idx));
+    }
+}
+
 struct NullEncoder;
 
 impl Encoder for NullEncoder {

arrow-json/src/writer/mod.rs

@@ -405,7 +405,7 @@ mod tests {
 
     use arrow_array::builder::*;
     use arrow_array::types::*;
-    use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer, ToByteSlice};
+    use arrow_buffer::{i256, Buffer, NullBuffer, OffsetBuffer, ToByteSlice};
     use arrow_data::ArrayData;
 
     use crate::reader::*;
@@ -1833,4 +1833,80 @@ mod tests {
             r#"[{"my_dict":"a"},{"my_dict":null},{"my_dict":null}]"#
         )
     }
+
+    #[test]
+    fn test_decimal128_encoder() {
+        let array = Decimal128Array::from_iter_values([1234, 5678, 9012])
+            .with_precision_and_scale(10, 2)
+            .unwrap();
+        let field = Arc::new(Field::new("decimal", array.data_type().clone(), true));
+        let schema = Schema::new(vec![field]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
+
+        let mut buf = Vec::new();
+        {
+            let mut writer = LineDelimitedWriter::new(&mut buf);
+            writer.write_batches(&[&batch]).unwrap();
+        }
+
+        assert_json_eq(
+            &buf,
+            r#"{"decimal":12.34}
+{"decimal":56.78}
+{"decimal":90.12}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_decimal256_encoder() {
+        let array = Decimal256Array::from_iter_values([
+            i256::from(123400),
+            i256::from(567800),
+            i256::from(901200),
+        ])
+        .with_precision_and_scale(10, 4)
+        .unwrap();
+        let field = Arc::new(Field::new("decimal", array.data_type().clone(), true));
+        let schema = Schema::new(vec![field]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
+
+        let mut buf = Vec::new();
+        {
+            let mut writer = LineDelimitedWriter::new(&mut buf);
+            writer.write_batches(&[&batch]).unwrap();
+        }
+
+        assert_json_eq(
+            &buf,
+            r#"{"decimal":12.3400}
+{"decimal":56.7800}
+{"decimal":90.1200}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_decimal_encoder_with_nulls() {
+        let array = Decimal128Array::from_iter([Some(1234), None, Some(5678)])
+            .with_precision_and_scale(10, 2)
+            .unwrap();
+        let field = Arc::new(Field::new("decimal", array.data_type().clone(), true));
+        let schema = Schema::new(vec![field]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
+
+        let mut buf = Vec::new();
+        {
+            let mut writer = LineDelimitedWriter::new(&mut buf);
+            writer.write_batches(&[&batch]).unwrap();
+        }
+
+        assert_json_eq(
+            &buf,
+            r#"{"decimal":12.34}
+{}
+{"decimal":56.78}
+"#,
+        );
+    }
 }