Address comments.

Author: viirya

python/pyspark/serializers.py

@@ -245,7 +245,7 @@ def __repr__(self):
         return "ArrowStreamSerializer"
 
 
-def _create_batch(series, timezone, runner_conf):
+def _create_batch(series, timezone, safecheck):
     """
     Create an Arrow record batch from the given pandas.Series or list of Series, with optional type.
 
@@ -285,10 +285,8 @@ def create_array(s, t):
             # TODO: see ARROW-1949. Remove when the minimum PyArrow version becomes 0.11.0.
             return pa.Array.from_pandas(s, mask=mask, type=t)
 
-        enabledArrowSafeTypeCheck = \
-            runner_conf.get("spark.sql.execution.pandas.arrowSafeTypeConversion", "true") == 'true'
         try:
-            array = pa.Array.from_pandas(s, mask=mask, type=t, safe=enabledArrowSafeTypeCheck)
+            array = pa.Array.from_pandas(s, mask=mask, type=t, safe=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 " + \
@@ -307,10 +305,10 @@ class ArrowStreamPandasSerializer(Serializer):
     Serializes Pandas.Series as Arrow data with Arrow streaming format.
     """
 
-    def __init__(self, timezone, runner_conf):
+    def __init__(self, timezone, safecheck):
         super(ArrowStreamPandasSerializer, self).__init__()
         self._timezone = timezone
-        self._runner_conf = runner_conf
+        self._safecheck = safecheck
 
     def arrow_to_pandas(self, arrow_column):
         from pyspark.sql.types import from_arrow_type, \
@@ -330,7 +328,7 @@ def dump_stream(self, iterator, stream):
         writer = None
         try:
             for series in iterator:
-                batch = _create_batch(series, self._timezone, self._runner_conf)
+                batch = _create_batch(series, self._timezone, self._safecheck)
                 if writer is None:
                     write_int(SpecialLengths.START_ARROW_STREAM, stream)
                     writer = pa.RecordBatchStreamWriter(stream, batch.schema)

python/pyspark/sql/session.py

@@ -556,10 +556,9 @@ def _create_from_pandas_with_arrow(self, pdf, schema, timezone):
         pdf_slices = (pdf[start:start + step] for start in xrange(0, len(pdf), step))
 
         # Create Arrow record batches
-        runner_conf = {"spark.sql.execution.pandas.arrowSafeTypeConversion":
-                       self._wrapped._conf.arrowSafeTypeConversion()}
+        safecheck = self._wrapped._conf.arrowSafeTypeConversion()
         batches = [_create_batch([(c, t) for (_, c), t in zip(pdf_slice.iteritems(), arrow_types)],
-                                 timezone, runner_conf)
+                                 timezone, safecheck)
                    for pdf_slice in pdf_slices]
 
         # Create the Spark schema from the first Arrow batch (always at least 1 batch after slicing)

python/pyspark/sql/tests/test_arrow.py

@@ -288,8 +288,10 @@ def test_createDataFrame_does_not_modify_input(self):
         # Integers with nulls will get NaNs filled with 0 and will be casted
         pdf.ix[1, '2_int_t'] = None
         pdf_copy = pdf.copy(deep=True)
-        self.spark.createDataFrame(pdf, schema=self.schema)
-        self.assertTrue(pdf.equals(pdf_copy))
+        with self.sql_conf({
+                "spark.sql.execution.pandas.arrowSafeTypeConversion": False}):
+            self.spark.createDataFrame(pdf, schema=self.schema)
+            self.assertTrue(pdf.equals(pdf_copy))
 
     def test_schema_conversion_roundtrip(self):
         from pyspark.sql.types import from_arrow_schema, to_arrow_schema

python/pyspark/sql/tests/test_pandas_udf.py

@@ -225,6 +225,38 @@ def udf(column):
                 "spark.sql.execution.pandas.arrowSafeTypeConversion": False}):
             df.select(['A']).withColumn('udf', udf('A')).collect()
 
+    def test_pandas_udf_arrow_overflow(self):
+        from distutils.version import LooseVersion
+        from pyspark.sql.functions import pandas_udf
+        import pandas as pd
+        import pyarrow as pa
+
+        df = self.spark.range(0, 1)
+
+        @pandas_udf(returnType="byte")
+        def udf(column):
+            return pd.Series([128])
+
+        # Arrow 0.11.0+ allows enabling or disabling safe type check.
+        if LooseVersion(pa.__version__) >= LooseVersion("0.11.0"):
+            # When enabling safe type check, Arrow 0.11.0+ disallows overflow cast.
+            with self.sql_conf({
+                    "spark.sql.execution.pandas.arrowSafeTypeConversion": True}):
+                with self.assertRaisesRegexp(Exception,
+                                             "Exception thrown when converting pandas.Series"):
+                    df.withColumn('udf', udf('id')).collect()
+
+            # Disabling safe type check, let Arrow do the cast anyway.
+            with self.sql_conf({"spark.sql.execution.pandas.arrowSafeTypeConversion": False}):
+                df.withColumn('udf', udf('id')).collect()
+        else:
+            # SQL config `arrowSafeTypeConversion` no matters for older Arrow.
+            # Overflow cast causes an error.
+            with self.sql_conf({"spark.sql.execution.pandas.arrowSafeTypeConversion": False}):
+                with self.assertRaisesRegexp(Exception,
+                                             "Integer value out of bounds"):
+                    df.withColumn('udf', udf('id')).collect()
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests.test_pandas_udf import *

python/pyspark/worker.py

@@ -253,7 +253,9 @@ def read_udfs(pickleSer, infile, eval_type):
 
         # NOTE: if timezone is set here, that implies respectSessionTimeZone is True
         timezone = runner_conf.get("spark.sql.session.timeZone", None)
-        ser = ArrowStreamPandasSerializer(timezone, runner_conf)
+        safecheck = runner_conf.get("spark.sql.execution.pandas.arrowSafeTypeConversion",
+                                    "true").lower() == 'true'
+        ser = ArrowStreamPandasSerializer(timezone, safecheck)
     else:
         ser = BatchedSerializer(PickleSerializer(), 100)
 

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1334,9 +1334,9 @@ object SQLConf {
   val PANDAS_ARROW_SAFE_TYPE_CONVERSION =
     buildConf("spark.sql.execution.pandas.arrowSafeTypeConversion")
       .internal()
-      .doc("When true, enabling Arrow do safe type conversion check when converting" +
-        "Pandas.Series to Arrow Array during serialization. Arrow will raise errors " +
-        "when detecting unsafe type conversion. When false, disabling Arrow's type " +
+      .doc("When true, Arrow will perform safe type conversion when converting " +
+        "Pandas.Series to Arrow array during serialization. Arrow will raise errors " +
+        "when detecting unsafe type conversion like overflow. When false, disabling Arrow's type " +
         "check and do type conversions anyway.")
       .booleanConf
       .createWithDefault(true)