Add SQL config for Arrow safe type check.

Author: viirya

python/pyspark/serializers.py

@@ -245,7 +245,7 @@ def __repr__(self):
         return "ArrowStreamSerializer"
 
 
-def _create_batch(series, timezone):
+def _create_batch(series, timezone, runner_conf):
     """
     Create an Arrow record batch from the given pandas.Series or list of Series, with optional type.
 
@@ -284,12 +284,17 @@ def create_array(s, t):
         elif LooseVersion(pa.__version__) < LooseVersion("0.11.0"):
             # 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=True)
+            array = pa.Array.from_pandas(s, mask=mask, type=t, safe=enabledArrowSafeTypeCheck)
         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 " + \
-                        "conversions warned by Arrow."
+                        "conversions warned by Arrow. Arrow safe type check can be " + \
+                        "disabled by using SQL config " + \
+                        "`spark.sql.execution.pandas.arrowSafeTypeConversion`."
             raise RuntimeError(error_msg % (s.dtype, t), e)
         return array
 
@@ -302,9 +307,10 @@ class ArrowStreamPandasSerializer(Serializer):
     Serializes Pandas.Series as Arrow data with Arrow streaming format.
     """
 
-    def __init__(self, timezone):
+    def __init__(self, timezone, runner_conf):
         super(ArrowStreamPandasSerializer, self).__init__()
         self._timezone = timezone
+        self._runner_conf = runner_conf
 
     def arrow_to_pandas(self, arrow_column):
         from pyspark.sql.types import from_arrow_type, \
@@ -324,7 +330,7 @@ def dump_stream(self, iterator, stream):
         writer = None
         try:
             for series in iterator:
-                batch = _create_batch(series, self._timezone)
+                batch = _create_batch(series, self._timezone, self._runner_conf)
                 if writer is None:
                     write_int(SpecialLengths.START_ARROW_STREAM, stream)
                     writer = pa.RecordBatchStreamWriter(stream, batch.schema)

python/pyspark/sql/tests/test_pandas_udf.py

@@ -212,13 +212,18 @@ def test_pandas_udf_detect_unsafe_type_conversion(self):
         def udf(column):
             return pd.Series(np.linspace(0, 1, 3))
 
-        udf_df = df.select(['A']).withColumn('udf', udf('A'))
-
         # Since 0.11.0, PyArrow supports the feature to raise an error for unsafe cast.
         if LooseVersion(pa.__version__) >= LooseVersion("0.11.0"):
-            with self.assertRaisesRegexp(Exception,
-                                         "Exception thrown when converting pandas.Series"):
-                    udf_df.collect()
+            with self.sql_conf({
+                    "spark.sql.execution.pandas.arrowSafeTypeConversion": True}):
+                with self.assertRaisesRegexp(Exception,
+                                             "Exception thrown when converting pandas.Series"):
+                    df.select(['A']).withColumn('udf', udf('A')).collect()
+
+        # Disabling Arrow safe type check.
+        with self.sql_conf({
+                "spark.sql.execution.pandas.arrowSafeTypeConversion": False}):
+            df.select(['A']).withColumn('udf', udf('A')).collect()
 
 
 if __name__ == "__main__":

python/pyspark/worker.py

@@ -253,7 +253,7 @@ 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)
+        ser = ArrowStreamPandasSerializer(timezone, runner_conf)
     else:
         ser = BatchedSerializer(PickleSerializer(), 100)
 

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

@@ -1331,6 +1331,16 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  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 " +
+        "check and do type conversions anyway.")
+      .booleanConf
+      .createWithDefault(true)
+
   val REPLACE_EXCEPT_WITH_FILTER = buildConf("spark.sql.optimizer.replaceExceptWithFilter")
     .internal()
     .doc("When true, the apply function of the rule verifies whether the right node of the" +
@@ -2005,6 +2015,8 @@ class SQLConf extends Serializable with Logging {
   def pandasGroupedMapAssignColumnsByName: Boolean =
     getConf(SQLConf.PANDAS_GROUPED_MAP_ASSIGN_COLUMNS_BY_NAME)
 
+  def arrowSafeTypeConversion: Boolean = getConf(SQLConf.PANDAS_ARROW_SAFE_TYPE_CONVERSION)
+
   def replaceExceptWithFilter: Boolean = getConf(REPLACE_EXCEPT_WITH_FILTER)
 
   def decimalOperationsAllowPrecisionLoss: Boolean = getConf(DECIMAL_OPERATIONS_ALLOW_PREC_LOSS)

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowUtils.scala

@@ -133,6 +133,8 @@ object ArrowUtils {
     }
     val pandasColsByName = Seq(SQLConf.PANDAS_GROUPED_MAP_ASSIGN_COLUMNS_BY_NAME.key ->
       conf.pandasGroupedMapAssignColumnsByName.toString)
-    Map(timeZoneConf ++ pandasColsByName: _*)
+    val arrowSafeTypeCheck = Seq(SQLConf.PANDAS_ARROW_SAFE_TYPE_CONVERSION.key ->
+      conf.arrowSafeTypeConversion.toString)
+    Map(timeZoneConf ++ pandasColsByName ++ arrowSafeTypeCheck: _*)
   }
 }