From 09e58a435a37b169a11af6e194b1739c47df34ef Mon Sep 17 00:00:00 2001
From: Yijie Shen <henry.yijieshen@gmail.com>
Date: Wed, 29 May 2024 03:56:17 -0700
Subject: [PATCH] fix: serialization of decimal (#5801)

---
 arrow-json/src/reader/decimal_array.rs | 26 ++++++++++++++++++++++++++
 arrow-json/src/reader/mod.rs           | 24 ++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/arrow-json/src/reader/decimal_array.rs b/arrow-json/src/reader/decimal_array.rs
index fc3c9aaa6b43..c51c4e968930 100644
--- a/arrow-json/src/reader/decimal_array.rs
+++ b/arrow-json/src/reader/decimal_array.rs
@@ -64,6 +64,32 @@ where
                     let value = parse_decimal::<D>(s, self.precision, self.scale)?;
                     builder.append_value(value)
                 }
+                TapeElement::I64(high) => match tape.get(*p + 1) {
+                    TapeElement::I32(low) => {
+                        let val = ((high as i64) << 32 | (low as u32) as i64).to_string();
+                        let value = parse_decimal::<D>(&val, self.precision, self.scale)?;
+                        builder.append_value(value)
+                    }
+                    _ => unreachable!(),
+                },
+                TapeElement::I32(val) => {
+                    let s = val.to_string();
+                    let value = parse_decimal::<D>(&s, self.precision, self.scale)?;
+                    builder.append_value(value)
+                }
+                TapeElement::F64(high) => match tape.get(*p + 1) {
+                    TapeElement::F32(low) => {
+                        let val = f64::from_bits((high as u64) << 32 | low as u64).to_string();
+                        let value = parse_decimal::<D>(&val, self.precision, self.scale)?;
+                        builder.append_value(value)
+                    }
+                    _ => unreachable!(),
+                },
+                TapeElement::F32(val) => {
+                    let s = f32::from_bits(val).to_string();
+                    let value = parse_decimal::<D>(&s, self.precision, self.scale)?;
+                    builder.append_value(value)
+                }
                 _ => return Err(tape.error(*p, "decimal")),
             }
         }
diff --git a/arrow-json/src/reader/mod.rs b/arrow-json/src/reader/mod.rs
index 628e5c96693d..9a113ee2bd75 100644
--- a/arrow-json/src/reader/mod.rs
+++ b/arrow-json/src/reader/mod.rs
@@ -2212,6 +2212,30 @@ mod tests {
         assert_eq!(values.values(), &[1681319393, -7200]);
     }
 
+    #[test]
+    fn test_serialize_decimal() {
+        let json = vec![
+            json!({"decimal": 1.234}),
+            json!({"decimal": "1.234"}),
+            json!({"decimal": 1234}),
+            json!({"decimal": "1234"}),
+        ];
+        let schema = Schema::new(vec![Field::new(
+            "decimal",
+            DataType::Decimal128(10, 3),
+            true,
+        )]);
+        let mut decoder = ReaderBuilder::new(Arc::new(schema))
+            .build_decoder()
+            .unwrap();
+        decoder.serialize(&json).unwrap();
+        let batch = decoder.flush().unwrap().unwrap();
+        assert_eq!(batch.num_rows(), 4);
+        assert_eq!(batch.num_columns(), 1);
+        let values = batch.column(0).as_primitive::<Decimal128Type>();
+        assert_eq!(values.values(), &[1234, 1234, 1234000, 1234000]);
+    }
+
     #[test]
     fn test_serde_field() {
         let field = Field::new("int", DataType::Int32, true);