chore: copy create external table patterns

Author: tswast

bigframes/bigquery/__init__.py

@@ -61,7 +61,7 @@
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
 from bigframes.bigquery._operations.sql import sql_scalar
 from bigframes.bigquery._operations.struct import struct
-from bigframes.bigquery.table import create_external_table
+from bigframes.bigquery._operations.table import create_external_table
 from bigframes.core.logging import log_adapter
 
 _functions = [

bigframes/bigquery/table.py bigframes/bigquery/_operations/ddl.pybigframes/bigquery/table.py renamed to bigframes/bigquery/_operations/ddl.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Mapping, Optional, Union
+from typing import Any, Mapping, Optional, Sequence, Union
 
 import bigframes_vendored.constants
 import google.cloud.bigquery
@@ -101,3 +101,85 @@ def create_external_table(
         session.read_gbq_query(sql)
 
     return _get_table_metadata(bqclient=session.bqclient, table_name=table_name)
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_table")
+def load_data(
+    uris: str | Sequence[str],
+    format: str,
+    destination_table: str,
+    *,
+    schema: Optional[Mapping[str, str]] = None,
+    cluster_by: Optional[Sequence[str]] = None,
+    partition_by: Optional[str] = None,
+    options: Optional[dict[str, Any]] = None,
+    load_options: Optional[dict[str, Any]] = None,
+    connection: Optional[str] = None,
+    hive_partition_columns: Optional[Mapping[str, str]] = None,
+    overwrite: bool = False,
+    session: Optional[bigframes.session.Session] = None,
+) -> pd.Series:
+    """
+    Loads data from external files into a BigQuery table using the `LOAD DATA` statement.
+
+    Args:
+        uris (str | List[str]):
+            The fully qualified URIs for the external data locations (e.g., 'gs://bucket/path/file.csv').
+        format (str):
+            The format of the external data (e.g., 'CSV', 'PARQUET', 'AVRO', 'JSON').
+        destination_table (str, optional):
+            The name of the destination table. If not specified, a temporary table will be created.
+        schema (List[google.cloud.bigquery.SchemaField], optional):
+            The schema of the destination table. If not provided, schema auto-detection will be used.
+        cluster_by (List[str], optional):
+            A list of columns to cluster the table by.
+        partition_by (str, optional):
+            The partition expression for the table.
+        options (dict[str, Any], optional):
+            Table options (e.g., {'description': 'my table'}).
+        load_options (dict[str, Any], optional):
+            Options for loading data (e.g., {'skip_leading_rows': 1}).
+        connection (str, optional):
+            The connection name to use for reading external data.
+        hive_partition_columns (List[google.cloud.bigquery.SchemaField], optional):
+            The external partitioning columns. If set to an empty list, partitioning is inferred.
+        overwrite (bool, default False):
+            If True, overwrites the destination table. If False, appends to it.
+        session (bigframes.session.Session, optional):
+            The session to use. If not provided, the default session is used.
+
+    Returns:
+        pandas.Series:
+            A Series with object dtype containing the table metadata. Reference
+            the `BigQuery Table REST API reference
+            <https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table>`_
+            for available fields.
+    """
+    import bigframes.pandas as bpd
+
+    if session is None:
+        session = bpd.get_global_session()
+
+    if isinstance(uris, str):
+        uris = [uris]
+
+    sql = bigframes.core.sql.table.load_data_ddl(
+        destination_table=destination_table,
+        uris=uris,
+        format=format,
+        schema_fields=schema,
+        cluster_by=cluster_by,
+        partition_by=partition_by,
+        table_options=options,
+        load_options=load_options,
+        connection=connection,
+        hive_partition_columns=hive_partition_columns,
+        overwrite=overwrite,
+    )
+
+    # Execute the LOAD DATA statement
+    session.read_gbq_query(sql)
+
+    # Return a DataFrame pointing to the destination table
+    # We use session.read_gbq to ensure it uses the same session
+    return session.read_gbq(destination_table)

bigframes/bigquery/_operations/io.py

@@ -21,7 +21,7 @@
 import decimal
 import json
 import math
-from typing import Any, cast, Collection, Iterable, Mapping, Optional, TYPE_CHECKING, Union
+from typing import cast, Collection, Iterable, Mapping, Optional, TYPE_CHECKING, Union
 
 import shapely.geometry.base  # type: ignore
 
@@ -172,7 +172,7 @@ def create_vector_index_ddl(
     table_name: str,
     column_name: str,
     stored_column_names: Collection[str],
-    options: Mapping[str, Union[str | int | bool | float]] = {},
+    options: Mapping[str, Union[str, int, bool, float]] = {},
 ) -> str:
     """Encode the VECTOR INDEX statement for BigQuery Vector Search."""
 
@@ -275,65 +275,3 @@ def schema_field_to_sql(field: bigquery.SchemaField) -> str:
     if field.description:
         sql += f" OPTIONS(description={simple_literal(field.description)})"
     return sql
-
-
-def load_data_ddl(
-    destination_table: str,
-    uris: list[str],
-    format: str,
-    *,
-    schema_fields: list[bigquery.SchemaField] | None = None,
-    cluster_by: list[str] | None = None,
-    partition_by: str | None = None,
-    table_options: dict[str, Any] | None = None,
-    load_options: dict[str, Any] | None = None,
-    connection: str | None = None,
-    hive_partition_columns: list[bigquery.SchemaField] | None = None,
-    overwrite: bool = False,
-) -> str:
-    """Construct a LOAD DATA DDL statement."""
-    action = "OVERWRITE" if overwrite else "INTO"
-
-    query = f"LOAD DATA {action} {googlesql.identifier(destination_table)}\n"
-
-    if schema_fields:
-        columns_sql = ",\n".join(schema_field_to_sql(field) for field in schema_fields)
-        query += f"(\n{columns_sql}\n)\n"
-
-    if partition_by:
-        query += f"PARTITION BY {partition_by}\n"
-
-    if cluster_by:
-        query += f"CLUSTER BY {', '.join(cluster_by)}\n"
-
-    if table_options:
-        opts_list = []
-        for k, v in table_options.items():
-            opts_list.append(f"{k}={simple_literal(v)}")
-        query += f"OPTIONS({', '.join(opts_list)})\n"
-
-    files_opts = {}
-    if load_options:
-        files_opts.update(load_options)
-
-    files_opts["uris"] = uris
-    files_opts["format"] = format
-
-    files_opts_list = []
-    for k, v in files_opts.items():
-        files_opts_list.append(f"{k}={simple_literal(v)}")
-
-    query += f"FROM FILES({', '.join(files_opts_list)})\n"
-
-    if hive_partition_columns:
-        cols_sql = ",\n".join(
-            schema_field_to_sql(field) for field in hive_partition_columns
-        )
-        query += f"WITH PARTITION COLUMNS (\n{cols_sql}\n)\n"
-    elif hive_partition_columns is not None:
-        query += "WITH PARTITION COLUMNS\n"
-
-    if connection:
-        query += f"WITH CONNECTION {connection}\n"
-
-    return query

bigframes/core/sql/__init__.py

@@ -14,7 +14,12 @@
 
 from __future__ import annotations
 
-from typing import Mapping, Optional, Union
+from typing import Any, Mapping, Optional, Union
+
+from google.cloud import bigquery
+
+import bigframes.core.compile.googlesql as googlesql
+import bigframes.core.sql
 
 
 def create_external_table_ddl(
@@ -66,3 +71,68 @@ def create_external_table_ddl(
         statement.append(f"OPTIONS ({options_str})")
 
     return " ".join(statement)
+
+
+def load_data_ddl(
+    destination_table: str,
+    uris: list[str],
+    format: str,
+    *,
+    schema_fields: list[bigquery.SchemaField] | None = None,
+    cluster_by: list[str] | None = None,
+    partition_by: str | None = None,
+    table_options: dict[str, Any] | None = None,
+    load_options: dict[str, Any] | None = None,
+    connection: str | None = None,
+    hive_partition_columns: list[bigquery.SchemaField] | None = None,
+    overwrite: bool = False,
+) -> str:
+    """Construct a LOAD DATA DDL statement."""
+    action = "OVERWRITE" if overwrite else "INTO"
+
+    query = f"LOAD DATA {action} {googlesql.identifier(destination_table)}\n"
+
+    if schema_fields:
+        columns_sql = ",\n".join(
+            bigframes.core.sql.schema_field_to_sql(field) for field in schema_fields
+        )
+        query += f"(\n{columns_sql}\n)\n"
+
+    if partition_by:
+        query += f"PARTITION BY {partition_by}\n"
+
+    if cluster_by:
+        query += f"CLUSTER BY {', '.join(cluster_by)}\n"
+
+    if table_options:
+        opts_list = []
+        for k, v in table_options.items():
+            opts_list.append(f"{k}={bigframes.core.sql.simple_literal(v)}")
+        query += f"OPTIONS({', '.join(opts_list)})\n"
+
+    files_opts = {}
+    if load_options:
+        files_opts.update(load_options)
+
+    files_opts["uris"] = uris
+    files_opts["format"] = format
+
+    files_opts_list = []
+    for k, v in files_opts.items():
+        files_opts_list.append(f"{k}={bigframes.core.sql.simple_literal(v)}")
+
+    query += f"FROM FILES({', '.join(files_opts_list)})\n"
+
+    if hive_partition_columns:
+        cols_sql = ",\n".join(
+            bigframes.core.sql.schema_field_to_sql(field)
+            for field in hive_partition_columns
+        )
+        query += f"WITH PARTITION COLUMNS (\n{cols_sql}\n)\n"
+    elif hive_partition_columns is not None:
+        query += "WITH PARTITION COLUMNS\n"
+
+    if connection:
+        query += f"WITH CONNECTION {connection}\n"
+
+    return query