handle empty arrays when inferring schema (#297)

Co-authored-by: Thomas McKenna <[email protected]>
pola-rs · Nov 26, 2024 · 9ae34f8 · 9ae34f8
1 parent d806f3c
commit 9ae34f8
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 14 deletions.
diff --git a/__tests__/dataframe.test.ts b/__tests__/dataframe.test.ts
@@ -2008,6 +2008,16 @@ describe("create", () => {
       string: pl.String,
     });
   });
+  test("from row objects, inferred schema, empty array", () => {
+    const df = pl.readRecords([
+      { a: [], b: 0 },
+      { a: [""], b: 0 },
+    ]);
+    expect(df.schema).toStrictEqual({
+      a: pl.List(pl.String),
+      b: pl.Float64,
+    });
+  });
   test("from row objects, with schema", () => {
     const rows = [
       { num: 1, date: "foo", string: "foo1" },

diff --git a/src/dataframe.rs b/src/dataframe.rs
@@ -1657,21 +1657,26 @@ fn obj_to_pairs(rows: &Array, len: usize) -> impl '_ + Iterator<Item = Vec<(Stri
                                 if val.is_array().unwrap() {
                                     let arr: napi::JsObject = unsafe { val.cast() };
                                     let len = arr.get_array_length().unwrap();
-                                    // dont compare too many items, as it could be expensive
-                                    let max_take = std::cmp::min(len as usize, 10);
-                                    let mut dtypes: Vec<DataType> =
-                                        Vec::with_capacity(len as usize);
-
-                                    for idx in 0..max_take {
-                                        let item: napi::JsUnknown =
-                                            arr.get_element(idx as u32).unwrap();
-                                        let ty = item.get_type().unwrap();
-                                        let dt: Wrap<DataType> = ty.into();
-                                        dtypes.push(dt.0)
-                                    }
-                                    let dtype = coerce_data_type(&dtypes);
 
-                                    DataType::List(dtype.into())
+                                    if len == 0 {
+                                        DataType::List(DataType::Null.into())
+                                    } else {
+                                        // dont compare too many items, as it could be expensive
+                                        let max_take = std::cmp::min(len as usize, 10);
+                                        let mut dtypes: Vec<DataType> =
+                                            Vec::with_capacity(len as usize);
+
+                                        for idx in 0..max_take {
+                                            let item: napi::JsUnknown =
+                                                arr.get_element(idx as u32).unwrap();
+                                            let ty = item.get_type().unwrap();
+                                            let dt: Wrap<DataType> = ty.into();
+                                            dtypes.push(dt.0)
+                                        }
+                                        let dtype = coerce_data_type(&dtypes);
+
+                                        DataType::List(dtype.into())
+                                    }
                                 } else if val.is_date().unwrap() {
                                     DataType::Datetime(TimeUnit::Milliseconds, None)
                                 } else {