fix(duckdb): partition initial DuckLake writes

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

Author: anxkhn

sqlmesh/core/engine_adapter/duckdb.py

@@ -173,7 +173,12 @@ def _create_table(
         track_rows_processed: bool = True,
         **kwargs: t.Any,
     ) -> None:
-        catalog = self.get_current_catalog()
+        table_name = (
+            table_name_or_schema.this
+            if isinstance(table_name_or_schema, exp.Schema)
+            else table_name_or_schema
+        )
+        catalog = exp.to_table(table_name).catalog or self.get_current_catalog()
         catalog_type_tuple = self.fetchone(
             exp.select("type")
             .from_("duckdb_databases()")
@@ -182,8 +187,22 @@ def _create_table(
         catalog_type = catalog_type_tuple[0] if catalog_type_tuple else None
 
         partitioned_by_exps = None
+        insert_expression = None
         if catalog_type == "ducklake":
             partitioned_by_exps = kwargs.pop("partitioned_by", None)
+            if (
+                partitioned_by_exps
+                and expression is not None
+                and (replace or not exists or not self.table_exists(table_name))
+            ):
+                insert_expression = expression.copy()
+                query = t.cast(exp.Query, expression)
+                expression = (
+                    exp.select("*")
+                    .from_(query.subquery("_sqlmesh_schema_only", copy=False))
+                    .where(exp.false())
+                    .limit(0)
+                )
 
         super()._create_table(
             table_name_or_schema,
@@ -199,12 +218,6 @@ def _create_table(
         )
 
         if partitioned_by_exps:
-            # Schema object contains column definitions, so we extract Table
-            table_name = (
-                table_name_or_schema.this
-                if isinstance(table_name_or_schema, exp.Schema)
-                else table_name_or_schema
-            )
             table_name_str = (
                 table_name.sql(dialect=self.dialect)
                 if isinstance(table_name, exp.Table)
@@ -215,6 +228,12 @@ def _create_table(
             )
             self.execute(f"ALTER TABLE {table_name_str} SET PARTITIONED BY ({partitioned_by_str});")
 
+            if insert_expression is not None:
+                self.execute(
+                    exp.insert(insert_expression, exp.to_table(table_name)),
+                    track_rows_processed=track_rows_processed,
+                )
+
     @property
     def _is_motherduck(self) -> bool:
         return self._extra_config.get("is_motherduck", False)

tests/core/engine_adapter/test_duckdb.py

@@ -1,4 +1,5 @@
 import typing as t
+from datetime import date
 
 import pandas as pd  # noqa: TID253
 import pytest
@@ -154,3 +155,31 @@ def test_ducklake_partitioning(adapter: EngineAdapter, duck_conn, tmp_path):
         f"SELECT * FROM __ducklake_metadata_{catalog}.main.ducklake_partition_info"
     ).fetchdf()
     assert partition_info.shape[0] == 1
+
+
+def test_ducklake_partitioning_on_initial_query(adapter: EngineAdapter, duck_conn, tmp_path):
+    catalog = "ducklake_initial_partition_db"
+
+    duck_conn.install_extension("ducklake")
+    duck_conn.load_extension("ducklake")
+    duck_conn.execute(
+        f"ATTACH 'ducklake:{tmp_path}/{catalog}.ducklake' AS {catalog} "
+        f"(DATA_PATH '{tmp_path}', DATA_INLINING_ROW_LIMIT 0);"
+    )
+
+    adapter.create_schema(f"{catalog}.test_schema")
+    adapter.replace_query(
+        f"{catalog}.test_schema.test_table",
+        parse_one("SELECT 1 AS id, DATE '2000-01-01' AS ds UNION ALL SELECT 2, DATE '2000-01-02'"),
+        partitioned_by=[exp.to_column("ds")],
+    )
+
+    assert adapter.fetchall(f"SELECT * FROM {catalog}.test_schema.test_table ORDER BY id") == [
+        (1, date(2000, 1, 1)),
+        (2, date(2000, 1, 2)),
+    ]
+    partition_ids = duck_conn.execute(
+        f"SELECT partition_id FROM __ducklake_metadata_{catalog}.main.ducklake_data_file"
+    ).fetchall()
+    assert partition_ids
+    assert all(partition_id is not None for (partition_id,) in partition_ids)