This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 4d4c0c62b Implement arrow_json encoder for Decimal128 & Decimal256
(#6606)
4d4c0c62b is described below
commit 4d4c0c62b30f2112860f8074d0d64176b484cdac
Author: Phillip LeBlanc <[email protected]>
AuthorDate: Mon Oct 21 18:41:06 2024 +0900
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 dfe8edb350bd427b552ffbca1b853a7ab52baa0a.
* Use ArrayFormatter directly
---------
Co-authored-by: Raphael Taylor-Davies <[email protected]>
---
arrow-json/src/writer/encoder.rs | 14 ++++++++
arrow-json/src/writer/mod.rs | 78 +++++++++++++++++++++++++++++++++++++++-
2 files changed, 91 insertions(+), 1 deletion(-)
diff --git a/arrow-json/src/writer/encoder.rs b/arrow-json/src/writer/encoder.rs
index dfe628521..29a302f74 100644
--- a/arrow-json/src/writer/encoder.rs
+++ b/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 {
diff --git a/arrow-json/src/writer/mod.rs b/arrow-json/src/writer/mod.rs
index 9a93cf041..df18a9ad5 100644
--- a/arrow-json/src/writer/mod.rs
+++ b/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}
+"#,
+ );
+ }
}
