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

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
 
@@ -510,9 +511,43 @@ 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.enable", "false").lower() == "true" \
+                    and len(data) > 0:
+                from pyspark.serializers import ArrowSerializer
+                from pyspark.sql.types import from_arrow_schema
+                import pyarrow as pa
+
+                # Slice the DataFrame into batches
+                split = -(-len(data) // self.sparkContext.defaultParallelism)  # round int up
+                slices = (data[i:i + split] for i in xrange(0, len(data), split))
+                batches = [pa.RecordBatch.from_pandas(sliced_df, preserve_index=False)
+                           for sliced_df in slices]
+
+                # write batches to temp file, read by JVM (borrowed from context.parallelize)
+                import os
+                from tempfile import NamedTemporaryFile
+                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.org.apache.spark.sql.execution.arrow.ArrowConverters.toDataFrame(
+                    jrdd, schema.json(), self._wrapped._jsqlContext)
+                df = DataFrame(jdf, self._wrapped)
+                df._schema = schema
+                return df
+            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

python/pyspark/sql/tests.py

@@ -3095,16 +3095,27 @@ 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)]
 
     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 createPandasDataFrameFromeData(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 +3132,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.enable", "false")
-        pdf = df.toPandas()
-        self.spark.conf.set("spark.sql.execution.arrow.enable", "true")
+        try:
+            pdf = df.toPandas()
+        finally:
+            self.spark.conf.set("spark.sql.execution.arrow.enable", "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.createPandasDataFrameFromeData()
         df = self.spark.createDataFrame(self.data, schema=self.schema)
         pdf_arrow = df.toPandas()
         self.assertFramesEqual(pdf_arrow, pdf)
@@ -3147,6 +3152,16 @@ def test_filtered_frame(self):
         self.assertEqual(pdf.columns[0], "i")
         self.assertTrue(pdf.empty)
 
+    def test_createDataFrame_toggle(self):
+        pdf = self.createPandasDataFrameFromeData()
+        self.spark.conf.set("spark.sql.execution.arrow.enable", "false")
+        try:
+            df_no_arrow = self.spark.createDataFrame(pdf)
+        finally:
+            self.spark.conf.set("spark.sql.execution.arrow.enable", "true")
+        df_arrow = self.spark.createDataFrame(pdf)
+        self.assertEquals(df_no_arrow.collect(), df_arrow.collect())
+
 
 @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,42 @@ def toArrowType(dt):
     return arrow_type
 
 
+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/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()
     }
   }
+
+  def toDataFrame(
+      arrowRDD: JavaRDD[Array[Byte]],
+      schemaString: String,
+      sqlContext: SQLContext): DataFrame = {
+    val rdd = arrowRDD.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)
+  }
 }