feat(postgres): bulk-load DataFrames with COPY

docs/integrations/engines/postgres.md

@@ -23,3 +23,15 @@ pip install "sqlmesh[postgres]"
 | `role`             | The role to use for authentication with the Postgres server                     | string | N        |
 | `sslmode`          | The security of the connection to the Postgres server                           | string | N        |
 | `application_name` | The name of the application to use for the connection                           | string | N        |
+
+### DataFrame loading
+
+When using the `psycopg2` driver, SQLMesh loads pandas DataFrames (such as those returned by Python models and used for seeds) with `COPY ... FROM STDIN (FORMAT csv)` instead of batched `INSERT ... VALUES` statements, which is much faster for large DataFrames. The data is copied into a temporary staging table, which is created in the same schema as the target table and then cast into the target table's column types. The staging table is dropped afterwards, including when the load fails. Because it is created in the target's schema, the connecting user needs the `CREATE` privilege on that schema, as for other SQLMesh DDL.
+
+No configuration is required and there is no option to disable this behavior. SQLMesh automatically falls back to `INSERT ... VALUES` when:
+
+- A column has a type or values that cannot be loaded with `COPY` identically. Supported columns are integers (up to 64 bits signed, up to 32 bits unsigned), `float64`, booleans, timestamps (with or without time zone), and strings. Other types, such as `float32`, 64-bit unsigned integers, `Decimal`, dates, bytes, dicts, lists, and categoricals, trigger the fallback.
+- The DataFrame has a single column and one of its values is `\.`.
+- The driver does not support `copy_expert` (for example, a driver other than `psycopg2`).
+
+The RisingWave engine always uses `INSERT ... VALUES`.

sqlmesh/core/engine_adapter/base.py

@@ -2921,9 +2921,15 @@ def _get_temp_table(
         Returns the name of the temp table that should be used for the given table name.
         """
         table = t.cast(exp.Table, exp.to_table(table).copy())
-        table.set(
-            "this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}", quoted=quoted)
-        )
+        prefix, suffix = "__temp_", f"_{random_id(short=True)}"
+        name = table.name
+        if self.MAX_IDENTIFIER_LENGTH is not None:
+            # Cut the middle of the name so that the temp table name is a valid identifier
+            keep = max(self.MAX_IDENTIFIER_LENGTH - len(prefix) - len(suffix), 0)
+            if len(name) > keep:
+                tail = keep // 2
+                name = name[: keep - tail] + (name[-tail:] if tail else "")
+        table.set("this", exp.to_identifier(f"{prefix}{name}{suffix}", quoted=quoted))
 
         if table_only:
             table.set("db", None)

sqlmesh/core/engine_adapter/postgres.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import csv
+import io
 import logging
 import re
 import typing as t
@@ -14,14 +16,59 @@
     logical_merge,
     GrantsFromInfoSchemaMixin,
 )
-from sqlmesh.core.engine_adapter.shared import set_catalog
+from sqlmesh.core.engine_adapter.shared import SourceQuery, set_catalog
+from sqlmesh.utils import random_id
 
 if t.TYPE_CHECKING:
+    import pandas as pd
+
     from sqlmesh.core._typing import TableName
-    from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
+    from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
 
 logger = logging.getLogger(__name__)
 
+# pandas dtypes that can be bulk loaded with COPY, mapped to the type of the staging table column.
+# The staging column types are derived from the DataFrame (not from the target table) and the final
+# `SELECT` casts them to the target types, which mirrors how the generic `VALUES` implementation types
+# literals by their Python value and then casts. In particular floats are staged as an unbounded
+# NUMERIC, because the literals emitted by the `VALUES` implementation are NUMERIC and e.g. casting
+# 2.5 to an integer rounds differently for NUMERIC (3) and DOUBLE PRECISION (2).
+_COPY_DTYPE_TO_STAGING_TYPE: t.Dict[str, str] = {
+    **{
+        dtype: "BIGINT"
+        for dtype in (
+            "int8",
+            "int16",
+            "int32",
+            "int64",
+            "uint8",
+            "uint16",
+            "uint32",
+            "Int8",
+            "Int16",
+            "Int32",
+            "Int64",
+            "UInt8",
+            "UInt16",
+            "UInt32",
+        )
+    },
+    "float64": "NUMERIC",
+    "Float64": "NUMERIC",
+    "bool": "BOOLEAN",
+    "boolean": "BOOLEAN",
+}
+# What pandas' `infer_dtype` reports for object / string columns that can be copied. Python numbers in an
+# object column (e.g. seeds, where NaN is replaced with None) are staged as NUMERIC, which is lossless.
+_COPY_INFERRED_TO_STAGING_TYPE: t.Dict[str, str] = {
+    "string": "TEXT",
+    "empty": "TEXT",  # only nulls
+    "boolean": "BOOLEAN",
+    "integer": "NUMERIC",
+    "floating": "NUMERIC",
+    "mixed-integer-float": "NUMERIC",
+}
+
 
 @set_catalog()
 class PostgresEngineAdapter(
@@ -85,6 +132,149 @@ def _fetch_native_df(
             self._connection_pool.commit()
         return df
 
+    def _df_to_source_queries(
+        self,
+        df: DF,
+        target_columns_to_types: t.Dict[str, exp.DataType],
+        batch_size: int,
+        target_table: TableName,
+        source_columns: t.Optional[t.List[str]] = None,
+    ) -> t.List[SourceQuery]:
+        """Bulk loads the DataFrame into a staging table with COPY instead of using `VALUES` batches.
+
+        The generic implementation is used when the DataFrame cannot be loaded safely with COPY.
+        """
+        import pandas as pd
+
+        assert isinstance(df, pd.DataFrame)
+
+        # pandas truncates `na_rep` to 32 characters when writing float columns, so keep the marker short
+        null_marker = f"__sqlmesh_null_{random_id(short=True)}__"
+        # Same column selection and order as the generic implementation
+        ordered_df = df[list(source_columns or target_columns_to_types)]
+        staging_columns_to_types = (
+            None
+            if ordered_df.empty or not hasattr(self.cursor, "copy_expert")
+            else self._get_copy_staging_columns_to_types(ordered_df, null_marker)
+        )
+        if staging_columns_to_types is None:
+            return super()._df_to_source_queries(
+                df,
+                target_columns_to_types,
+                batch_size,
+                target_table,
+                source_columns=source_columns,
+            )
+
+        staging_table = self._get_temp_table(target_table or "pandas")
+
+        def query_factory() -> Query:
+            # The factory can be called multiple times. The staging table is only loaded once, which
+            # assumes that every call would have produced the same result.
+            if not self.table_exists(staging_table):
+                try:
+                    self.create_table(staging_table, staging_columns_to_types)
+                    self._copy_df(staging_table, ordered_df, null_marker, batch_size)
+                except Exception:
+                    # Callers such as `replace_query` invoke the factory outside of `with source_query`,
+                    # so `cleanup_func` won't run and a committed staging table would leak.
+                    try:
+                        self._drop_copy_staging_table(staging_table)
+                    except Exception:
+                        logger.warning("Failed to drop staging table %s", staging_table.sql())
+                    raise
+            return exp.select(
+                *self._casted_columns(target_columns_to_types, source_columns=source_columns)
+            ).from_(staging_table)
+
+        return [
+            SourceQuery(
+                query_factory=query_factory,
+                cleanup_func=lambda: self._drop_copy_staging_table(staging_table),
+            )
+        ]
+
+    def _drop_copy_staging_table(self, staging_table: exp.Table) -> None:
+        try:
+            self.drop_table(staging_table)
+        except Exception as ex:
+            # The cleanup runs while an exception is propagating. If it left the transaction aborted
+            # the DROP fails with "current transaction is aborted", which would hide the original
+            # error. The staging table was created in that transaction, so its rollback removes it.
+            sqlstate = getattr(ex, "pgcode", None) or getattr(ex, "sqlstate", None)
+            if sqlstate != "25P02" or not self._connection_pool.is_transaction_active:
+                raise
+            logger.debug("Skipped dropping %s from an aborted transaction", staging_table.sql())
+
+    @staticmethod
+    def _get_copy_staging_columns_to_types(
+        df: pd.DataFrame, null_marker: str
+    ) -> t.Optional[t.Dict[str, exp.DataType]]:
+        """Returns the staging table column types for the DataFrame, or None if it can't be loaded with COPY.
+
+        None is returned when a column has a dtype or values that wouldn't be read back by COPY exactly
+        as the generic `VALUES` implementation would write them, e.g. dicts, bytes, Decimals, dates or
+        categoricals, or when a string would be misread by the CSV parser.
+        """
+        import pandas as pd
+
+        staging_columns_to_types: t.Dict[str, exp.DataType] = {}
+        for column, series in df.items():
+            dtype = series.dtype
+            staging_type: t.Optional[str]
+            if isinstance(dtype, pd.DatetimeTZDtype):
+                staging_type = "TIMESTAMPTZ"
+            elif pd.api.types.is_datetime64_dtype(dtype):
+                staging_type = "TIMESTAMP"
+            elif str(dtype) in _COPY_DTYPE_TO_STAGING_TYPE:
+                staging_type = _COPY_DTYPE_TO_STAGING_TYPE[str(dtype)]
+            elif dtype == object or isinstance(dtype, pd.StringDtype):
+                staging_type = _COPY_INFERRED_TO_STAGING_TYPE.get(
+                    pd.api.types.infer_dtype(series, skipna=True)
+                )
+                if staging_type == "TEXT":
+                    if series.eq(null_marker).any():
+                        return None
+                    # On PostgreSQL <= 17 an unquoted `\.` line ends the data, which happens if it is the only column
+                    if len(df.columns) == 1 and series.eq("\\.").any():
+                        return None
+            else:
+                staging_type = None
+            if staging_type is None:
+                return None
+            staging_columns_to_types[str(column)] = exp.DataType.build(
+                staging_type, dialect=PostgresEngineAdapter.DIALECT
+            )
+        return staging_columns_to_types
+
+    def _copy_df(
+        self, table: exp.Table, df: pd.DataFrame, null_marker: str, batch_size: int
+    ) -> None:
+        """Loads the DataFrame into the table with `COPY ... FROM STDIN` in chunks of `batch_size` rows.
+
+        This is the only place that calls the driver's COPY API (`psycopg2`'s `copy_expert`).
+        The copy runs on the adapter's cursor so it is part of the current transaction.
+        """
+        columns = ", ".join(self._to_sql(exp.column(c, quoted=True)) for c in df.columns)
+        copy_sql = (
+            f"COPY {self._to_sql(table)} ({columns}) "
+            f"FROM STDIN WITH (FORMAT csv, NULL '{null_marker}')"
+        )
+        chunk_size = batch_size or max(len(df.index), 1)
+        for start in range(0, len(df.index), chunk_size):
+            buffer = io.StringIO()
+            df.iloc[start : start + chunk_size].to_csv(
+                buffer,
+                header=False,
+                index=False,
+                na_rep=null_marker,
+                quoting=csv.QUOTE_MINIMAL,
+                lineterminator="\n",
+            )
+            buffer.seek(0)
+            # A text buffer is encoded by the driver with the connection's client_encoding
+            self.cursor.copy_expert(copy_sql, buffer)
+
     def _create_table_like(
         self,
         target_table_name: TableName,

sqlmesh/core/engine_adapter/risingwave.py

@@ -6,6 +6,7 @@
 
 from sqlglot import exp
 
+from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
 from sqlmesh.core.engine_adapter.shared import (
     set_catalog,
@@ -34,6 +35,9 @@ class RisingwaveEngineAdapter(PostgresEngineAdapter):
     MAX_IDENTIFIER_LENGTH = None
     SUPPORTS_GRANTS = False
 
+    # COPY into a staging table hasn't been verified on RisingWave, so DataFrames keep using `VALUES`
+    _df_to_source_queries = EngineAdapter._df_to_source_queries
+
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]: