Fix arrow type inference error (fall back to string)

vega · Mar 21, 2024 · 2fe8155 · 2fe8155
1 parent 3615368
commit 2fe8155
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git a/python/vegafusion/tests/test_datasource.py b/python/vegafusion/tests/test_datasource.py
@@ -0,0 +1,25 @@
+import pandas as pd
+import pyarrow as pa
+from vegafusion.datasource import PandasDatasource
+
+def test_mixed_col_type_inference():
+    df = pd.DataFrame({
+        'a': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'],
+        'b': [28, 55, 43, 91, 81, 53, 19, 87, 52],
+        'c': [28, 55, 43, 91, 81, 53, 19, 87, 52],
+        'd': [28, 55, 43, 91, 81, 53, 19, 87, 52],
+        'e': [28, 55, 43, 91, 81, 53, 19, 87, 52],
+        'f': [28, 55, 43, 91, '81.5', 53, 19, 87, 52],
+    })
+    datasource = PandasDatasource(df)
+
+    expected_schema = pa.schema([
+        pa.field("a", pa.string()),
+        pa.field("b", pa.int64()),
+        pa.field("c", pa.int64()),
+        pa.field("d", pa.int64()),
+        pa.field("e", pa.int64()),
+        pa.field("f", pa.string()),
+    ])
+
+    assert datasource.schema() == expected_schema
diff --git a/python/vegafusion/vegafusion/datasource/pandas_datasource.py b/python/vegafusion/vegafusion/datasource/pandas_datasource.py
@@ -23,13 +23,18 @@ def __init__(self, df: pd.DataFrame, sample_size: int = 1000, batch_size: int =
             if not isinstance(col, str):
                 continue
             try:
-                # We will expand categoricals (not yet supported in VegaFusion)
-                if isinstance(pd_type, pd.CategoricalDtype):
-                    cat = df[col].cat
-                    field = pa.Schema.from_pandas(pd.DataFrame({col: cat.categories})).field(col)
-                else:
-                    candidate_schema = pa.Schema.from_pandas(df.iloc[::sample_stride][[col]])
-                    field = candidate_schema.field(col)
+                try:
+                    # We will expand categoricals (not yet supported in VegaFusion)
+                    if isinstance(pd_type, pd.CategoricalDtype):
+                        cat = df[col].cat
+                        field = pa.Schema.from_pandas(pd.DataFrame({col: cat.categories})).field(col)
+                    else:
+                        candidate_schema = pa.Schema.from_pandas(df.iloc[::sample_stride][[col]])
+                        field = candidate_schema.field(col)
+                except (pa.ArrowTypeError, pa.ArrowInvalid):
+                    # If arrow fails to infer the type, fall back to string
+                    casts[col] = "str"
+                    field = pa.field(col, pa.string())
 
                 # Convert Decimal columns to float
                 if isinstance(field.type, (pa.Decimal128Type, pa.Decimal256Type)):