fix: serialization of decimal (#5801)

apache · May 29, 2024 · 09e58a4 · 09e58a4
1 parent 12f57af
commit 09e58a4
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 0 deletions.
diff --git 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
@@ -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);