Merge branch 'main' into shuowei-fix-unpivot-empty

Author: shuoweil

bigframes/bigquery/_operations/table.py

@@ -19,8 +19,8 @@
 import google.cloud.bigquery
 import pandas as pd
 
+import bigframes.core.compile.sqlglot.sql as sg_sql
 import bigframes.core.logging.log_adapter as log_adapter
-import bigframes.core.sql.table
 import bigframes.session
 
 
@@ -80,14 +80,16 @@ def create_external_table(
     """
     import bigframes.pandas as bpd
 
-    sql = bigframes.core.sql.table.create_external_table_ddl(
-        table_name=table_name,
-        replace=replace,
-        if_not_exists=if_not_exists,
-        columns=columns,
-        partition_columns=partition_columns,
-        connection_name=connection_name,
-        options=options,
+    sql = sg_sql.to_sql(
+        sg_sql.create_external_table(
+            table_name=table_name,
+            replace=replace,
+            if_not_exists=if_not_exists,
+            columns=columns,
+            partition_columns=partition_columns,
+            connection_name=connection_name,
+            options=options,
+        )
     )
 
     if session is None:

bigframes/core/compile/sqlglot/sql/__init__.py

@@ -22,7 +22,7 @@
     table,
     to_sql,
 )
-from bigframes.core.compile.sqlglot.sql.ddl import load_data
+from bigframes.core.compile.sqlglot.sql.ddl import create_external_table, load_data
 from bigframes.core.compile.sqlglot.sql.dml import insert, replace
 
 __all__ = [
@@ -35,6 +35,7 @@
     "table",
     "to_sql",
     # From ddl.py
+    "create_external_table",
     "load_data",
     # From dml.py
     "insert",

bigframes/core/compile/sqlglot/sql/ddl.py

@@ -22,51 +22,6 @@
 from bigframes.core.compile.sqlglot.sql import base
 
 
-def _loaddata_sql(self: sg.Generator, expression: sge.LoadData) -> str:
-    out = ["LOAD DATA"]
-    if expression.args.get("overwrite"):
-        out.append("OVERWRITE")
-
-    out.append(f"INTO {self.sql(expression, 'this').strip()}")
-
-    # We ignore inpath as it's just a dummy to satisfy sqlglot requirements
-    # but BigQuery uses FROM FILES instead.
-
-    columns = self.sql(expression, "columns").strip()
-    if columns:
-        out.append(columns)
-
-    partition_by = self.sql(expression, "partition_by").strip()
-    if partition_by:
-        out.append(partition_by)
-
-    cluster_by = self.sql(expression, "cluster_by").strip()
-    if cluster_by:
-        out.append(cluster_by)
-
-    options = self.sql(expression, "options").strip()
-    if options:
-        out.append(options)
-
-    from_files = self.sql(expression, "from_files").strip()
-    if from_files:
-        out.append(f"FROM FILES {from_files}")
-
-    with_partition_columns = self.sql(expression, "with_partition_columns").strip()
-    if with_partition_columns:
-        out.append(f"WITH PARTITION COLUMNS {with_partition_columns}")
-
-    connection = self.sql(expression, "connection").strip()
-    if connection:
-        out.append(f"WITH CONNECTION {connection}")
-
-    return " ".join(out)
-
-
-# Register the transform for BigQuery generator
-sg.dialects.bigquery.BigQuery.Generator.TRANSFORMS[sge.LoadData] = _loaddata_sql
-
-
 def load_data(
     table_name: str,
     *,
@@ -84,21 +39,6 @@ def load_data(
     # Quoting is handled by the dialect.
     table_expr = sge.Table(this=base.identifier(table_name))
 
-    sge_columns = (
-        sge.Schema(
-            this=None,
-            expressions=[
-                sge.ColumnDef(
-                    this=base.identifier(name),
-                    kind=sge.DataType.build(typ, dialect="bigquery"),
-                )
-                for name, typ in columns.items()
-            ],
-        )
-        if columns
-        else None
-    )
-
     sge_partition_by = (
         sge.PartitionedByProperty(
             this=base.identifier(partition_by[0])
@@ -115,50 +55,166 @@ def load_data(
         else None
     )
 
-    sge_table_options = (
-        sge.Properties(
-            expressions=[
-                sge.Property(this=base.identifier(k), value=base.literal(v))
-                for k, v in table_options.items()
-            ]
-        )
-        if table_options
-        else None
-    )
-
     sge_from_files = sge.Tuple(
         expressions=[
             sge.Property(this=base.identifier(k), value=base.literal(v))
             for k, v in from_files_options.items()
         ]
     )
 
-    sge_with_partition_columns = (
-        sge.Schema(
-            this=None,
-            expressions=[
-                sge.ColumnDef(
-                    this=base.identifier(name),
-                    kind=sge.DataType.build(typ, dialect="bigquery"),
-                )
-                for name, typ in with_partition_columns.items()
-            ],
-        )
-        if with_partition_columns
-        else None
-    )
-
     sge_connection = base.identifier(connection_name) if connection_name else None
 
     return sge.LoadData(
         this=table_expr,
         overwrite=(write_disposition == "OVERWRITE"),
         inpath=sge.convert("fake"),  # satisfy sqlglot's required inpath arg
-        columns=sge_columns,
+        columns=_get_sge_schema(columns),
         partition_by=sge_partition_by,
         cluster_by=sge_cluster_by,
-        options=sge_table_options,
+        options=_get_sge_properties(table_options),
         from_files=sge_from_files,
-        with_partition_columns=sge_with_partition_columns,
+        with_partition_columns=_get_sge_schema(with_partition_columns),
+        connection=sge_connection,
+    )
+
+
+def create_external_table(
+    table_name: str,
+    *,
+    replace: bool = False,
+    if_not_exists: bool = False,
+    columns: Optional[Mapping[str, str]] = None,
+    partition_columns: Optional[Mapping[str, str]] = None,
+    connection_name: Optional[str] = None,
+    options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None,
+) -> sge.Create:
+    """Generates the CREATE EXTERNAL TABLE DDL statement."""
+    sge_connection = base.identifier(connection_name) if connection_name else None
+
+    table_expr = sge.Table(this=base.identifier(table_name))
+
+    # sqlglot.expressions.Create usually takes 'this' (Table or Schema)
+    sge_schema = _get_sge_schema(columns)
+    this: sge.Table | sge.Schema
+    if sge_schema:
+        sge_schema.set("this", table_expr)
+        this = sge_schema
+    else:
+        this = table_expr
+
+    return sge.Create(
+        this=this,
+        kind="EXTERNAL TABLE",
+        replace=replace,
+        exists_ok=if_not_exists,
+        properties=_get_sge_properties(options),
         connection=sge_connection,
+        partition_columns=_get_sge_schema(partition_columns),
     )
+
+
+def _get_sge_schema(
+    columns: Optional[Mapping[str, str]] = None
+) -> Optional[sge.Schema]:
+    if not columns:
+        return None
+
+    return sge.Schema(
+        this=None,
+        expressions=[
+            sge.ColumnDef(
+                this=base.identifier(name),
+                kind=sge.DataType.build(typ, dialect=base.DIALECT),
+            )
+            for name, typ in columns.items()
+        ],
+    )
+
+
+def _get_sge_properties(
+    options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None
+) -> Optional[sge.Properties]:
+    if not options:
+        return None
+
+    return sge.Properties(
+        expressions=[
+            sge.Property(this=base.identifier(k), value=base.literal(v))
+            for k, v in options.items()
+        ]
+    )
+
+
+def _loaddata_sql(self: sg.Generator, expression: sge.LoadData) -> str:
+    out = ["LOAD DATA"]
+    if expression.args.get("overwrite"):
+        out.append("OVERWRITE")
+
+    out.append(f"INTO {self.sql(expression, 'this').strip()}")
+
+    # We ignore inpath as it's just a dummy to satisfy sqlglot requirements
+    # but BigQuery uses FROM FILES instead.
+
+    columns = self.sql(expression, "columns").strip()
+    if columns:
+        out.append(columns)
+
+    partition_by = self.sql(expression, "partition_by").strip()
+    if partition_by:
+        out.append(partition_by)
+
+    cluster_by = self.sql(expression, "cluster_by").strip()
+    if cluster_by:
+        out.append(cluster_by)
+
+    options = self.sql(expression, "options").strip()
+    if options:
+        out.append(options)
+
+    from_files = self.sql(expression, "from_files").strip()
+    if from_files:
+        out.append(f"FROM FILES {from_files}")
+
+    with_partition_columns = self.sql(expression, "with_partition_columns").strip()
+    if with_partition_columns:
+        out.append(f"WITH PARTITION COLUMNS {with_partition_columns}")
+
+    connection = self.sql(expression, "connection").strip()
+    if connection:
+        out.append(f"WITH CONNECTION {connection}")
+
+    return " ".join(out)
+
+
+def _create_sql(self: sg.Generator, expression: sge.Create) -> str:
+    kind = expression.args.get("kind")
+    if kind != "EXTERNAL TABLE":
+        return self.create_sql(expression)
+
+    out = ["CREATE"]
+    if expression.args.get("replace"):
+        out.append("OR REPLACE")
+    out.append("EXTERNAL TABLE")
+    if expression.args.get("exists_ok"):
+        out.append("IF NOT EXISTS")
+
+    out.append(self.sql(expression, "this"))
+
+    connection = self.sql(expression, "connection").strip()
+    if connection:
+        out.append(f"WITH CONNECTION {connection}")
+
+    partition_columns = self.sql(expression, "partition_columns").strip()
+    if partition_columns:
+        out.append(f"WITH PARTITION COLUMNS {partition_columns}")
+
+    properties = self.sql(expression, "properties").strip()
+    if properties:
+        out.append(properties)
+
+    return " ".join(out)
+
+
+# Register the transform for BigQuery generator
+base.DIALECT.Generator.TRANSFORMS[sge.LoadData] = _loaddata_sql
+base.DIALECT.Generator.TRANSFORMS[sge.Create] = _create_sql