perf(pyspark): ensure that pyspark DDL doesn't use VALUES

Author: cpcloud

ibis/backends/pyspark/__init__.py

@@ -88,17 +88,15 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 class PySparkTableSetFormatter(TableSetFormatter):
     def _format_in_memory_table(self, op):
-        names = op.schema.names
-        rows = ", ".join(
-            f"({', '.join(map(repr, row))})"
-            for row in op.data.itertuples(index=False)
-        )
-        signature = ", ".join(map(self._quote_identifier, names))
-        name = self._quote_identifier(op.name or "_")
-        return f"(VALUES {rows} AS {name} ({signature}))"
+        # we don't need to compile the table to a VALUES statement because the
+        # table has been registered already by createOrReplaceTempView.
+        #
+        # The only place where the SQL API is currently used is DDL operations
+        return op.name
 
 
 class PySparkCompiler(Compiler):
+    cheap_in_memory_tables = True
     table_set_formatter_class = PySparkTableSetFormatter
 
 
@@ -463,6 +461,8 @@ def create_table(
                     table_name, format=format, mode=mode
                 )
                 return
+            else:
+                self._register_in_memory_tables(obj)
 
             ast = self.compiler.to_ast(obj)
             select = ast.queries[0]
@@ -487,6 +487,10 @@ def create_table(
 
         return self.raw_sql(statement.compile())
 
+    def _register_in_memory_table(self, table_op):
+        spark_df = self.compile(table_op.to_expr())
+        spark_df.createOrReplaceTempView(table_op.name)
+
     def create_view(
         self,
         name: str,

ibis/backends/pyspark/compiler.py

@@ -2168,17 +2168,16 @@ def compile_random(*args, **kwargs):
     return F.rand()
 
 
+@compiles(ops.InMemoryTable)
 @compiles(PandasInMemoryTable)
 def compile_in_memory_table(t, expr, scope, timecontext, session, **kwargs):
     op = expr.op()
     fields = [
         pt.StructField(name, ibis_dtype_to_spark_dtype(dtype), dtype.nullable)
         for name, dtype in op.schema.items()
     ]
-    return session.createDataFrame(
-        data=op.data._df,
-        schema=pt.StructType(fields),
-    )
+    schema = pt.StructType(fields)
+    return session.createDataFrame(data=op.data.to_frame(), schema=schema)
 
 
 @compiles(ops.BitwiseAnd)