[SPARK-20791][PYSPARK] Use Arrow to create Spark DataFrame from Pandas

python/pyspark/java_gateway.py

@@ -121,6 +121,7 @@ def killChild():
     java_import(gateway.jvm, "org.apache.spark.mllib.api.python.*")
     # TODO(davies): move into sql
     java_import(gateway.jvm, "org.apache.spark.sql.*")
+    java_import(gateway.jvm, "org.apache.spark.sql.api.python.*")
     java_import(gateway.jvm, "org.apache.spark.sql.hive.*")
     java_import(gateway.jvm, "scala.Tuple2")
 

python/pyspark/sql/dataframe.py

@@ -70,12 +70,12 @@ class DataFrame(object):
     .. versionadded:: 1.3
     """
 
-    def __init__(self, jdf, sql_ctx):
+    def __init__(self, jdf, sql_ctx, schema=None):
         self._jdf = jdf
         self.sql_ctx = sql_ctx
         self._sc = sql_ctx and sql_ctx._sc
         self.is_cached = False
-        self._schema = None  # initialized lazily
+        self._schema = schema  # initialized lazily if None
         self._lazy_rdd = None
 
     @property

python/pyspark/sql/session.py

@@ -23,6 +23,7 @@
 
 if sys.version >= '3':
     basestring = unicode = str
+    xrange = range
 else:
     from itertools import imap as map
 
@@ -414,6 +415,43 @@ def _createFromLocal(self, data, schema):
         data = [schema.toInternal(row) for row in data]
         return self._sc.parallelize(data), schema
 
+    def _createFromPandasWithArrow(self, df, schema):
+        """
+        Create a DataFrame from a given pandas.DataFrame by slicing the into partitions, converting
+        to Arrow data, then reading into the JVM to parallelsize. If a schema is passed in, the
+        data types will be used to coerce the data in Pandas to Arrow conversion.
+        """
+        import os
+        from tempfile import NamedTemporaryFile
+        from pyspark.serializers import ArrowSerializer
+        from pyspark.sql.types import from_arrow_schema, to_arrow_schema
+        import pyarrow as pa
+
+        # Slice the DataFrame into batches
+        step = -(-len(df) // self.sparkContext.defaultParallelism)  # round int up
+        df_slices = (df[start:start + step] for start in xrange(0, len(df), step))
+        arrow_schema = to_arrow_schema(schema) if schema is not None else None
+        batches = [pa.RecordBatch.from_pandas(df_slice, schema=arrow_schema, preserve_index=False)
+                   for df_slice in df_slices]
+
+        # write batches to temp file, read by JVM (borrowed from context.parallelize)
+        tempFile = NamedTemporaryFile(delete=False, dir=self._sc._temp_dir)
+        try:
+            serializer = ArrowSerializer()
+            serializer.dump_stream(batches, tempFile)
+            tempFile.close()
+            readRDDFromFile = self._jvm.PythonRDD.readRDDFromFile
+            jrdd = readRDDFromFile(self._jsc, tempFile.name, len(batches))
+        finally:
+            # readRDDFromFile eagerily reads the file so we can delete right after.
+            os.unlink(tempFile.name)
+
+        # Create the Spark DataFrame, there will be at least 1 batch
+        schema = from_arrow_schema(batches[0].schema)
+        jdf = self._jvm.PythonSQLUtils.arrowPayloadToDataFrame(
+            jrdd, schema.json(), self._wrapped._jsqlContext)
+        return DataFrame(jdf, self._wrapped, schema)
+
     @since(2.0)
     @ignore_unicode_prefix
     def createDataFrame(self, data, schema=None, samplingRatio=None, verifySchema=True):
@@ -510,9 +548,13 @@ def createDataFrame(self, data, schema=None, samplingRatio=None, verifySchema=Tr
         except Exception:
             has_pandas = False
         if has_pandas and isinstance(data, pandas.DataFrame):
-            if schema is None:
-                schema = [str(x) for x in data.columns]
-            data = [r.tolist() for r in data.to_records(index=False)]
+            if self.conf.get("spark.sql.execution.arrow.enabled", "false").lower() == "true" \
+                    and len(data) > 0:
+                return self._createFromPandasWithArrow(data, schema)
+            else:
+                if schema is None:
+                    schema = [str(x) for x in data.columns]
+                data = [r.tolist() for r in data.to_records(index=False)]
 
         if isinstance(schema, StructType):
             verify_func = _make_type_verifier(schema) if verifySchema else lambda _: True
@@ -541,9 +583,7 @@ def prepare(obj):
             rdd, schema = self._createFromLocal(map(prepare, data), schema)
         jrdd = self._jvm.SerDeUtil.toJavaArray(rdd._to_java_object_rdd())
         jdf = self._jsparkSession.applySchemaToPythonRDD(jrdd.rdd(), schema.json())
-        df = DataFrame(jdf, self._wrapped)
-        df._schema = schema
-        return df
+        return DataFrame(jdf, self._wrapped, schema)
 
     @ignore_unicode_prefix
     @since(2.0)

python/pyspark/sql/tests.py

@@ -3095,16 +3095,32 @@ def setUpClass(cls):
             StructField("3_long_t", LongType(), True),
             StructField("4_float_t", FloatType(), True),
             StructField("5_double_t", DoubleType(), True)])
-        cls.data = [("a", 1, 10, 0.2, 2.0),
-                    ("b", 2, 20, 0.4, 4.0),
-                    ("c", 3, 30, 0.8, 6.0)]
+        cls.data = [(u"a", 1, 10, 0.2, 2.0),
+                    (u"b", 2, 20, 0.4, 4.0),
+                    (u"c", 3, 30, 0.8, 6.0)]
+
+    @classmethod
+    def tearDownClass(cls):
+        ReusedPySparkTestCase.tearDownClass()
+        cls.spark.stop()
 
     def assertFramesEqual(self, df_with_arrow, df_without):
         msg = ("DataFrame from Arrow is not equal" +
                ("\n\nWith Arrow:\n%s\n%s" % (df_with_arrow, df_with_arrow.dtypes)) +
                ("\n\nWithout:\n%s\n%s" % (df_without, df_without.dtypes)))
         self.assertTrue(df_without.equals(df_with_arrow), msg=msg)
 
+    def createPandasDataFrameFromData(self):
+        import pandas as pd
+        import numpy as np
+        data_dict = {}
+        for j, name in enumerate(self.schema.names):
+            data_dict[name] = [self.data[i][j] for i in range(len(self.data))]
+        # need to convert these to numpy types first
+        data_dict["2_int_t"] = np.int32(data_dict["2_int_t"])
+        data_dict["4_float_t"] = np.float32(data_dict["4_float_t"])
+        return pd.DataFrame(data=data_dict)
+
     def test_unsupported_datatype(self):
         schema = StructType([StructField("dt", DateType(), True)])
         df = self.spark.createDataFrame([(datetime.date(1970, 1, 1),)], schema=schema)
@@ -3121,21 +3137,15 @@ def test_null_conversion(self):
     def test_toPandas_arrow_toggle(self):
         df = self.spark.createDataFrame(self.data, schema=self.schema)
         self.spark.conf.set("spark.sql.execution.arrow.enabled", "false")
-        pdf = df.toPandas()
-        self.spark.conf.set("spark.sql.execution.arrow.enabled", "true")
+        try:
+            pdf = df.toPandas()
+        finally:
+            self.spark.conf.set("spark.sql.execution.arrow.enabled", "true")
         pdf_arrow = df.toPandas()
         self.assertFramesEqual(pdf_arrow, pdf)
 
     def test_pandas_round_trip(self):
-        import pandas as pd
-        import numpy as np
-        data_dict = {}
-        for j, name in enumerate(self.schema.names):
-            data_dict[name] = [self.data[i][j] for i in range(len(self.data))]
-        # need to convert these to numpy types first
-        data_dict["2_int_t"] = np.int32(data_dict["2_int_t"])
-        data_dict["4_float_t"] = np.float32(data_dict["4_float_t"])
-        pdf = pd.DataFrame(data=data_dict)
+        pdf = self.createPandasDataFrameFromData()
         df = self.spark.createDataFrame(self.data, schema=self.schema)
         pdf_arrow = df.toPandas()
         self.assertFramesEqual(pdf_arrow, pdf)
@@ -3147,6 +3157,29 @@ def test_filtered_frame(self):
         self.assertEqual(pdf.columns[0], "i")
         self.assertTrue(pdf.empty)
 
+    def test_createDataFrame_toggle(self):
+        pdf = self.createPandasDataFrameFromData()
+        self.spark.conf.set("spark.sql.execution.arrow.enabled", "false")
+        try:
+            df_no_arrow = self.spark.createDataFrame(pdf)
+        finally:
+            self.spark.conf.set("spark.sql.execution.arrow.enabled", "true")
+        df_arrow = self.spark.createDataFrame(pdf)
+        self.assertEquals(df_no_arrow.collect(), df_arrow.collect())
+
+    def test_createDataFrame_with_schema(self):
+        pdf = self.createPandasDataFrameFromData()
+        df = self.spark.createDataFrame(pdf, schema=self.schema)
+        self.assertEquals(self.schema, df.schema)
+        pdf_arrow = df.toPandas()
+        self.assertFramesEqual(pdf_arrow, pdf)
+
+    def test_schema_conversion_roundtrip(self):
+        from pyspark.sql.types import from_arrow_schema, to_arrow_schema
+        arrow_schema = to_arrow_schema(self.schema)
+        schema_rt = from_arrow_schema(arrow_schema)
+        self.assertEquals(self.schema, schema_rt)
+
 
 @unittest.skipIf(not _have_pandas or not _have_arrow, "Pandas or Arrow not installed")
 class VectorizedUDFTests(ReusedPySparkTestCase):

python/pyspark/sql/types.py

@@ -1624,6 +1624,50 @@ def to_arrow_type(dt):
     return arrow_type
 
 
+def to_arrow_schema(schema):
+    """ Convert a schema from Spark to Arrow
+    """
+    import pyarrow as pa
+    fields = [pa.field(field.name, to_arrow_type(field.dataType)) for field in schema]
+    return pa.schema(fields)
+
+
+def from_arrow_type(at):
+    """ Convert pyarrow type to Spark data type.
+    """
+    # TODO: newer pyarrow has is_boolean(at) functions that would be better to check type
+    import pyarrow as pa
+    if at == pa.bool_():
+        spark_type = BooleanType()
+    elif at == pa.int8():
+        spark_type = ByteType()
+    elif at == pa.int16():
+        spark_type = ShortType()
+    elif at == pa.int32():
+        spark_type = IntegerType()
+    elif at == pa.int64():
+        spark_type = LongType()
+    elif at == pa.float32():
+        spark_type = FloatType()
+    elif at == pa.float64():
+        spark_type = DoubleType()
+    elif type(at) == pa.DecimalType:
+        spark_type = DecimalType(precision=at.precision, scale=at.scale)
+    elif at == pa.string():
+        spark_type = StringType()
+    else:
+        raise TypeError("Unsupported type in conversion from Arrow: " + str(at))
+    return spark_type
+
+
+def from_arrow_schema(arrow_schema):
+    """ Convert schema from Arrow to Spark.
+    """
+    return StructType(
+        [StructField(field.name, from_arrow_type(field.type), nullable=field.nullable)
+         for field in arrow_schema])
+
+
 def _test():
     import doctest
     from pyspark.context import SparkContext

sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala

@@ -17,9 +17,12 @@
 
 package org.apache.spark.sql.api.python
 
+import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.sql.{DataFrame, SQLContext}
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
 import org.apache.spark.sql.catalyst.expressions.ExpressionInfo
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
+import org.apache.spark.sql.execution.arrow.ArrowConverters
 import org.apache.spark.sql.types.DataType
 
 private[sql] object PythonSQLUtils {
@@ -29,4 +32,19 @@ private[sql] object PythonSQLUtils {
   def listBuiltinFunctionInfos(): Array[ExpressionInfo] = {
     FunctionRegistry.functionSet.flatMap(f => FunctionRegistry.builtin.lookupFunction(f)).toArray
   }
+
+  /**
+   * Python Callable function to convert ArrowPayloads into a [[DataFrame]].
+   *
+   * @param payloadRDD A JavaRDD of ArrowPayloads.
+   * @param schemaString JSON Formatted Schema for ArrowPayloads.
+   * @param sqlContext The active [[SQLContext]].
+   * @return The converted [[DataFrame]].
+   */
+  def arrowPayloadToDataFrame(
+       payloadRDD: JavaRDD[Array[Byte]],
+       schemaString: String,
+       sqlContext: SQLContext): DataFrame = {
+    ArrowConverters.toDataFrame(payloadRDD, schemaString, sqlContext)
+  }
 }

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

@@ -29,6 +29,8 @@ import org.apache.arrow.vector.schema.ArrowRecordBatch
 import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel
 
 import org.apache.spark.TaskContext
+import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.sql.{DataFrame, SQLContext}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.execution.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
 import org.apache.spark.sql.types._
@@ -203,4 +205,16 @@ private[sql] object ArrowConverters {
       reader.close()
     }
   }
+
+  private[sql] def toDataFrame(
+      payloadRDD: JavaRDD[Array[Byte]],
+      schemaString: String,
+      sqlContext: SQLContext): DataFrame = {
+    val rdd = payloadRDD.rdd.mapPartitions { iter =>
+      val context = TaskContext.get()
+      ArrowConverters.fromPayloadIterator(iter.map(new ArrowPayload(_)), context)
+    }
+    val schema = DataType.fromJson(schemaString).asInstanceOf[StructType]
+    sqlContext.internalCreateDataFrame(rdd, schema)
+  }
 }