move from_spark, etc. to SparkBackend (hail-is#5019)

* move from_spark, etc. to SparkBackend

* fixed bug

* fixed bug

Author: cseed

hail/python/hail/backend/backend.py

@@ -5,7 +5,9 @@
 from hail.expr.table_type import *
 from hail.expr.matrix_type import *
 from hail.ir.renderer import Renderer
+from hail.table import Table
 
+import pyspark
 
 class Backend(object):
     @abc.abstractmethod
@@ -43,6 +45,20 @@ def matrix_read_type(self, mir):
         jir = self._to_java_ir(mir)
         return tmatrix._from_java(jir.typ())
 
+    def from_spark(self, df, key):
+        return Table._from_java(Env.hail().table.Table.fromDF(Env.hc()._jhc, df._jdf, key))
+
+    def to_spark(self, t, flatten):
+        t = t.expand_types()
+        if flatten:
+            t = t.flatten()
+        return pyspark.sql.DataFrame(t._jt.toDF(Env.hc()._jsql_context), Env.sql_context())
+
+    def to_pandas(self, t, flatten):
+        return self.to_spark(t, flatten).toPandas()
+
+    def from_pandas(self, df, key):
+        return Table.from_spark(Env.sql_context().createDataFrame(df), key)
 
 class ServiceBackend(Backend):
     def __init__(self, host, port=80, scheme='http'):

hail/python/hail/table.py

@@ -2493,8 +2493,7 @@ def from_spark(df, key=[]):
         :class:`.Table`
             Table constructed from the Spark SQL DataFrame.
         """
-
-        return Table._from_java(Env.hail().table.Table.fromDF(Env.hc()._jhc, df._jdf, key))
+        return Env.spark_backend('from_spark').from_spark(df, key)
 
     @typecheck_method(flatten=bool)
     def to_spark(self, flatten=True):
@@ -2513,10 +2512,7 @@ def to_spark(self, flatten=True):
         :class:`.pyspark.sql.DataFrame`
 
         """
-        t = self.expand_types()
-        if flatten:
-            t = t.flatten()
-        return pyspark.sql.DataFrame(t._jt.toDF(Env.hc()._jsql_context), Env.sql_context())
+        return Env.spark_backend('to_spark').to_spark(self, flatten)
 
     @typecheck_method(flatten=bool)
     def to_pandas(self, flatten=True):
@@ -2536,7 +2532,7 @@ def to_pandas(self, flatten=True):
         :class:`.pandas.DataFrame`
 
         """
-        return self.to_spark(flatten).toPandas()
+        return Env.spark_backend('to_pandas').to_pandas(self, flatten)
 
     @staticmethod
     @typecheck(df=pandas.DataFrame,
@@ -2560,7 +2556,7 @@ def from_pandas(df, key=[]):
         -------
         :class:`.Table`
         """
-        return Table.from_spark(Env.sql_context().createDataFrame(df), key)
+        return Env.spark_backend('from_pandas').from_pandas(df, key)
 
     @typecheck_method(other=table_type, tolerance=nullable(numeric), absolute=bool)
     def _same(self, other, tolerance=1e-6, absolute=False):

hail/python/hail/utils/java.py

@@ -65,6 +65,13 @@ def hc():
     def backend():
         return Env.hc()._backend
 
+    def spark_backend(op):
+        b = Env.backend()
+        if isinstance(b, hail.backend.SparkBackend):
+            return b
+        else:
+            raise NotImplementedError(f"{b.__class__.__name__} doesn't support {op}, only SparkBackend")
+
     @staticmethod
     def sql_context():
         return Env.hc()._sql_context