Updates

* Use infer_type over Schema.from_pandas for arrow type inference, as it can better handle extension types and pd.NA values
* Call __arrow_array__ directly if it is present to exit create_array early in _create_batch
* Add pandas version checks where required for tests
* Add tests covering pd.NA and BooleanDtype conversion

Author: moskvax

python/pyspark/sql/pandas/conversion.py

@@ -394,11 +394,11 @@ def _create_from_pandas_with_arrow(self, pdf, schema, timezone):
 
         # Create the Spark schema from list of names passed in with Arrow types
         if isinstance(schema, (list, tuple)):
-            # Arrow < 0.17.0 cannot handle ExtensionDtype columns when inferring the schema
-            arrow_schema = pa.Schema.from_pandas(pdf.astype('object'), preserve_index=False)
+            inferred_types = [pa.infer_type(s, mask=s.isna(), from_pandas=True)
+                              for s in (pdf[c] for c in pdf)]
             struct = StructType()
-            for name, field in zip(schema, arrow_schema):
-                struct.add(name, from_arrow_type(field.type), nullable=field.nullable)
+            for name, t in zip(schema, inferred_types):
+                struct.add(name, from_arrow_type(t), nullable=True)
             schema = struct
 
         # Determine arrow types to coerce data when creating batches

python/pyspark/sql/pandas/serializers.py

@@ -141,6 +141,7 @@ def _create_batch(self, series):
         :return: Arrow RecordBatch
         """
         import pandas as pd
+        from pandas.api.types import is_categorical_dtype
         import pyarrow as pa
         from pyspark.sql.pandas.types import _check_series_convert_timestamps_internal
         # Make input conform to [(series1, type1), (series2, type2), ...]
@@ -150,16 +151,22 @@ def _create_batch(self, series):
         series = ((s, None) if not isinstance(s, (list, tuple)) else s for s in series)
 
         def create_array(s, t):
-            # If the series implements __arrow_array__, conversion will fail if a mask is passed
-            mask = s.isnull() if not hasattr(s.values, "__arrow_array__") else None
+            # Create with __arrow_array__ if the series' backing array implements it
+            series_array = getattr(s, 'array', s._values)
+            if hasattr(series_array, "__arrow_array__"):
+                return series_array.__arrow_array__(type=t)
+
             # Ensure timestamp series are in expected form for Spark internal representation
             if t is not None and pa.types.is_timestamp(t):
                 s = _check_series_convert_timestamps_internal(s, self._timezone)
-            elif type(s.dtype) == pd.CategoricalDtype:
+            elif is_categorical_dtype(s.dtype):
                 # Note: This can be removed once minimum pyarrow version is >= 0.16.1
                 s = s.astype(s.dtypes.categories.dtype)
             try:
-                array = pa.Array.from_pandas(s, mask=mask, type=t, safe=self._safecheck)
+                mask = s.isnull()
+                # pass _ndarray_values to avoid potential failed type checks from pandas array types
+                array = pa.Array.from_pandas(s._ndarray_values, mask=mask, type=t,
+                                             safe=self._safecheck)
             except pa.ArrowException as e:
                 error_msg = "Exception thrown when converting pandas.Series (%s) to Arrow " + \
                             "Array (%s). It can be caused by overflows or other unsafe " + \

python/pyspark/sql/tests/test_arrow.py

@@ -34,6 +34,8 @@
 if have_pandas:
     import pandas as pd
     from pandas.util.testing import assert_frame_equal
+    from distutils.version import LooseVersion
+    pandas_version = LooseVersion(pd.__version__)
 
 if have_pyarrow:
     import pyarrow as pa
@@ -442,25 +444,47 @@ def test_createDataFrame_with_category_type(self):
         self.assertIsInstance(arrow_first_category_element, str)
         self.assertIsInstance(spark_first_category_element, str)
 
-    def test_createDataFrame_from_string_extension_dtype(self):
-        pdf = pd.DataFrame({u"A": [u"a", u"b", u"c", u"d"]})
-        pdf_ext_dtype = pd.DataFrame({u"A": [u"a", u"b", u"c", u"d"]}, dtype=pd.StringDtype())
+    def _assert_converted_dfs_equal(self, pdf1, pdf2):
+        df1 = self.spark.createDataFrame(pdf1)
+        df2 = self.spark.createDataFrame(pdf2)
+        self.assertEqual(df1.schema, df2.schema)
+        self.assertEqual(df1.collect(), df2.collect())
 
-        df = self.spark.createDataFrame(pdf)
-        df_ext_dtype = self.spark.createDataFrame(pdf_ext_dtype)
-
-        self.assertEqual(df_ext_dtype.schema, df.schema)
-        self.assertEqual(df_ext_dtype.collect(), df.collect())
-
-    def test_createDataFrame_from_integer_extension_dtype(self):
+    @unittest.skipIf(pandas_version < "0.24.0", "pandas < 0.24.0 missing Int64Dtype")
+    def test_createDataFrame_with_pandas_integer_dtype(self):
         pdf = pd.DataFrame({u"A": range(4)})
         pdf_ext_dtype = pd.DataFrame({u"A": range(4)}, dtype=pd.Int64Dtype())
+        self._assert_converted_dfs_equal(pdf, pdf_ext_dtype)
+
+    @unittest.skipIf(pandas_version < "1.0.0",
+                     "pandas < 1.0.0 missing StringDtype and BooleanDtype")
+    def test_createDataFrame_with_pandas_boolean_and_string_dtypes(self):
+        pdf = pd.DataFrame({
+            u"A": pd.Series([0, 1, 2, 3]),
+            u"B": pd.Series([u"a", u"b", u"c", u"d"]),
+            u"C": pd.Series([True, False, True, False]),
+        })
+        pdf_ext_dtype = pd.DataFrame({
+            u"A": pd.Series([0, 1, 2, 3], dtype=pd.Int64Dtype()),
+            u"B": pd.Series([u"a", u"b", u"c", u"d"], dtype=pd.StringDtype()),
+            u"C": pd.Series([True, False, True, False], dtype=pd.BooleanDtype()),
+        })
+        self._assert_converted_dfs_equal(pdf, pdf_ext_dtype)
+
+    @unittest.skipIf(pandas_version < "1.0.0", "pandas < 1.0.0 missing pd.NA")
+    def test_createDataFrame_with_pd_NA_values(self):
+        pdf = pd.DataFrame({
+            u"A": pd.Series([0, pd.NA, 2, 3]),
+            u"B": pd.Series([pd.NA, u"b", u"c", u"d"]),
+            u"C": pd.Series([True, False, pd.NA, False]),
+        })
+        pdf_ext_dtype = pd.DataFrame({
+            u"A": pd.Series([0, pd.NA, 2, 3], dtype=pd.Int64Dtype()),
+            u"B": pd.Series([pd.NA, u"b", u"c", u"d"], dtype=pd.StringDtype()),
+            u"C": pd.Series([True, False, pd.NA, False], dtype=pd.BooleanDtype()),
+        })
+        self._assert_converted_dfs_equal(pdf, pdf_ext_dtype)
 
-        df = self.spark.createDataFrame(pdf)
-        df_ext_dtype = self.spark.createDataFrame(pdf_ext_dtype)
-
-        self.assertEqual(df_ext_dtype.schema, df.schema)
-        self.assertEqual(df_ext_dtype.collect(), df.collect())
 
 
 @unittest.skipIf(