fix: Fix SO in json inference (#18887)

pola-rs · Sep 24, 2024 · a716989 · a716989
1 parent f1133a4
commit a716989
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/crates/polars-json/src/json/infer_schema.rs b/crates/polars-json/src/json/infer_schema.rs
@@ -134,11 +134,15 @@ pub(crate) fn coerce_dtype<A: Borrow<ArrowDataType>>(datatypes: &[A]) -> ArrowDa
             true,
         )));
     } else if datatypes.len() > 2 {
-        return coerce_dtype(datatypes);
+        return datatypes
+            .iter()
+            .map(|t| t.borrow().clone())
+            .reduce(|a, b| coerce_dtype(&[a, b]))
+            .expect("not empty");
     }
     let (lhs, rhs) = (datatypes[0].borrow(), datatypes[1].borrow());
 
-    return match (lhs, rhs) {
+    match (lhs, rhs) {
         (lhs, rhs) if lhs == rhs => lhs.clone(),
         (LargeList(lhs), LargeList(rhs)) => {
             let inner = coerce_dtype(&[lhs.dtype(), rhs.dtype()]);
@@ -171,5 +175,5 @@ pub(crate) fn coerce_dtype<A: Borrow<ArrowDataType>>(datatypes: &[A]) -> ArrowDa
         (Null, rhs) => rhs.clone(),
         (lhs, Null) => lhs.clone(),
         (_, _) => LargeUtf8,
-    };
+    }
 }
diff --git a/py-polars/tests/unit/io/test_json.py b/py-polars/tests/unit/io/test_json.py
@@ -430,3 +430,11 @@ def test_empty_list_json() -> None:
     df = pl.read_json(b"[]")
     assert df.shape == (0, 0)
     assert isinstance(df, pl.DataFrame)
+
+
+def test_json_infer_3_dtypes() -> None:
+    # would SO before
+    df = pl.DataFrame({"a": ["{}", "1", "[1, 2]"]})
+    out = df.select(pl.col("a").str.json_decode())
+    assert out["a"].to_list() == [None, ["1"], ["1", "2"]]
+    assert out.dtypes[0] == pl.List(pl.String)