fix(pyspark): default to inferring the schema of CSV files and assumi…

…ng they have a header with `header=True`
ibis-project · Aug 11, 2023 · 0ffda75 · 0ffda75
1 parent 3323156
commit 0ffda75
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 1 deletion.
diff --git a/ibis/backends/pyspark/__init__.py b/ibis/backends/pyspark/__init__.py
@@ -671,8 +671,12 @@ def read_csv(
         ir.Table
             The just-registered table
         """
+        inferSchema = kwargs.pop("inferSchema", True)
+        header = kwargs.pop("header", True)
         source_list = normalize_filenames(source_list)
-        spark_df = self._session.read.csv(source_list, **kwargs)
+        spark_df = self._session.read.csv(
+            source_list, inferSchema=inferSchema, header=header, **kwargs
+        )
         table_name = table_name or util.gen_name("read_csv")
 
         spark_df.createOrReplaceTempView(table_name)

diff --git a/ibis/backends/tests/test_register.py b/ibis/backends/tests/test_register.py
@@ -10,6 +10,7 @@
 import pytest
 from pytest import param
 
+import ibis
 from ibis.backends.conftest import TEST_TABLES
 
 if TYPE_CHECKING:
@@ -444,6 +445,21 @@ def num_diamonds(data_dir):
         return sum(1 for _ in f) - 1
 
 
+DIAMONDS_COLUMN_TYPES = {
+    # snowflake's `INFER_SCHEMA` returns this for the diamonds CSV `price`
+    # column type
+    "snowflake": {
+        "carat": "decimal(3, 2)",
+        "depth": "decimal(3, 1)",
+        "table": "decimal(3, 1)",
+        "x": "decimal(4, 2)",
+        "y": "decimal(4, 2)",
+        "z": "decimal(4, 2)",
+    },
+    "pyspark": {"price": "int32"},
+}
+
+
 @pytest.mark.parametrize(
     ("in_table_name", "out_table_name"),
     [
@@ -474,4 +490,22 @@ def test_read_csv(con, data_dir, in_table_name, out_table_name, num_diamonds):
         table = con.read_csv(fname, table_name=in_table_name)
 
     assert any(out_table_name in t for t in con.list_tables())
+
+    special_types = DIAMONDS_COLUMN_TYPES.get(con.name, {})
+
+    assert table.schema() == ibis.schema(
+        {
+            "carat": "float64",
+            "cut": "string",
+            "color": "string",
+            "clarity": "string",
+            "depth": "float64",
+            "table": "float64",
+            "price": "int64",
+            "x": "float64",
+            "y": "float64",
+            "z": "float64",
+            **special_types,
+        }
+    )
     assert table.count().execute() == num_diamonds