Enable INSERT..SELECT and COPY FROM pushdown for partitioned Iceberg tables

Partitioned writes previously went through the row-by-row PartitionedDestReceiver,
which routes each row to a per-partition CSV file before converting to Parquet.
This change enables DuckDB's PARTITION_BY in COPY TO for tables using pushdownable
transforms (identity, year, month, day, hour), letting DuckDB split the data in a
single pass. Bucket and truncate transforms continue to use the existing path.

Key changes:
- Add partition_pushdown.c with transform-to-DuckDB-SQL conversion, query wrapping
  with synthetic partition columns, and Hive-style path parsing for partition values
- Extend WriteQueryResultTo with partitionByExprs parameter that wraps the query
  with synthetic columns AFTER validation/interval wrappers and adds PARTITION_BY
  to the COPY command
- Modify AddQueryResultToTable to detect pushdownable partitions, pass expressions
  through, and parse per-file partition values from DuckDB output paths
- Lift blanket partition blocks in IsPushdownableInsertSelectQuery and
  IsCopyFromPushdownable to allow pushdown when all transforms are supported
- Disable FILE_SIZE_BYTES when PARTITION_BY is used (DuckDB limitation)

Signed-off-by: Marco Slot <marco.slot@snowflake.com>

Author: sfc-gh-mslot

pg_lake_copy/src/copy/copy.c

@@ -41,6 +41,8 @@
 #include "pg_lake/extensions/pg_lake_copy.h"
 #include "pg_lake/extensions/pg_parquet.h"
 #include "pg_lake/extensions/postgis.h"
+#include "pg_lake/fdw/partition_pushdown.h"
+#include "pg_lake/fdw/partition_transform.h"
 #include "pg_lake/fdw/writable_table.h"
 #include "pg_lake/fdw/schema_operations/register_field_ids.h"
 #include "pg_lake/partitioning/partition_by_parser.h"
@@ -651,9 +653,14 @@ IsCopyFromPushdownable(Relation relation, List *columnNameList,
 	if (!RelationColumnsSuitableForPushdown(relation, sourceFormat))
 		return false;
 
+	/*
+	 * Partitioned tables can be pushed down if all transforms are supported
+	 * by DuckDB PARTITION_BY (identity, year, month, day, hour).
+	 */
 	const char *partitionBy = GetIcebergTablePartitionByOption(relationId);
 
-	if (partitionBy != NULL)
+	if (partitionBy != NULL &&
+		GetPartitionByExpressionsForRelation(relationId) == NIL)
 		return false;
 
 	return true;

pg_lake_engine/include/pg_lake/data_file/data_file_stats.h

@@ -92,6 +92,12 @@ typedef struct DataFileStats
 
 	/* for a new data file with row IDs, the start of the range */
 	int64		rowIdStart;
+
+	/*
+	 * partition key values from COPY TO return_stats (NULL if not
+	 * partitioned)
+	 */
+	char	   *partitionKeysText;
 }			DataFileStats;
 
 typedef struct StatsCollector

pg_lake_engine/include/pg_lake/pgduck/map.h

@@ -19,7 +19,7 @@
 
 #include "pg_lake/pgduck/type.h"
 
-Oid			GetOrCreatePGMapType(const char *name);
+extern PGDLLEXPORT Oid GetOrCreatePGMapType(const char *name);
 char	   *GetDuckDBMapDefinitionForPGType(Oid postgresTypeId,
 											CopyDataFormat format);
 

pg_lake_engine/include/pg_lake/pgduck/write_data.h

@@ -56,7 +56,8 @@ extern PGDLLEXPORT StatsCollector * WriteQueryResultTo(char *query,
 													   TupleDesc queryTupleDesc,
 													   List *leafFields,
 													   IcebergOutOfRangePolicy outOfRangePolicy,
-													   bool wrapNativeIntervals);
+													   bool wrapNativeIntervals,
+													   List *partitionByExprs);
 extern PGDLLEXPORT void AppendFields(StringInfo map, DataFileSchema * schema);
 extern PGDLLEXPORT char *TupleDescToColumnMapForWrite(TupleDesc tupleDesc, CopyDataFormat destinationFormat);
 extern PGDLLEXPORT char *TupleDescToProjectionListForWrite(TupleDesc tupleDesc,

pg_lake_engine/pg_lake_engine--3.3--3.4.sql

@@ -1 +1,5 @@
 -- Upgrade script for pg_lake_engine from 3.3 to 3.4
+
+-- Pre-create MAP(TEXT,TEXT) type for partition_keys parsing in partitioned writes.
+-- This avoids runtime type creation during DML operations.
+SELECT map_type.create('text'::regtype, 'text'::regtype);

pg_lake_engine/src/data_file/data_file_stats.c

@@ -192,6 +192,11 @@ GetDataFileStatsListFromPGResult(PGresult *result, List *leafFields, DataFileSch
 			{
 				fileStats->dataFilePath = pstrdup(resultValue);
 			}
+			else if (strcmp(resultColName, "partition_keys") == 0)
+			{
+				if (!PQgetisnull(result, resultRowIndex, resultColIndex))
+					fileStats->partitionKeysText = pstrdup(resultValue);
+			}
 		}
 
 		statsList = lappend(statsList, fileStats);

pg_lake_engine/src/pgduck/write_data.c

@@ -104,7 +104,8 @@ ConvertCSVFileTo(char *csvFilePath, TupleDesc csvTupleDesc, int maxLineSize,
 							  csvTupleDesc,
 							  leafFields,
 							  ICEBERG_OOR_NONE,
-							  false /* wrapNativeIntervals */ );
+							  false /* wrapNativeIntervals */ ,
+							  NIL /* partitionByExprs */ );
 }
 
 
@@ -124,7 +125,8 @@ WriteQueryResultTo(char *query,
 				   TupleDesc queryTupleDesc,
 				   List *leafFields,
 				   IcebergOutOfRangePolicy outOfRangePolicy,
-				   bool wrapNativeIntervals)
+				   bool wrapNativeIntervals,
+				   List *partitionByExprs)
 {
 	if (outOfRangePolicy != ICEBERG_OOR_NONE)
 	{
@@ -139,6 +141,34 @@ WriteQueryResultTo(char *query,
 													   queryHasRowId);
 	}
 
+	/*
+	 * If partition expressions are given, wrap the (already validated) query
+	 * with synthetic partition columns. This must happen AFTER the validation
+	 * and interval wrappers, because those reconstruct the SELECT list from
+	 * queryTupleDesc and would drop any extra columns added earlier.
+	 */
+	if (partitionByExprs != NIL)
+	{
+		StringInfoData wrapped;
+
+		initStringInfo(&wrapped);
+		appendStringInfoString(&wrapped, "SELECT *");
+
+		int			partIndex = 0;
+		ListCell   *exprCell = NULL;
+
+		foreach(exprCell, partitionByExprs)
+		{
+			char	   *expr = strVal(lfirst(exprCell));
+
+			appendStringInfo(&wrapped, ", %s AS __part_%d", expr, partIndex);
+			partIndex++;
+		}
+
+		appendStringInfo(&wrapped, " FROM (%s) __partitioned_source", query);
+		query = wrapped.data;
+	}
+
 	StringInfoData command;
 
 	initStringInfo(&command);
@@ -364,6 +394,24 @@ WriteQueryResultTo(char *query,
 			elog(ERROR, "unexpected format: %s", formatName);
 	}
 
+	/* add PARTITION_BY if partitioning expressions were specified */
+	if (partitionByExprs != NIL)
+	{
+		appendStringInfoString(&command, ", PARTITION_BY (");
+
+		int			numExprs = list_length(partitionByExprs);
+
+		for (int i = 0; i < numExprs; i++)
+		{
+			if (i > 0)
+				appendStringInfoString(&command, ", ");
+
+			appendStringInfo(&command, "__part_%d", i);
+		}
+
+		appendStringInfoString(&command, ")");
+	}
+
 	/* end WITH options */
 	appendStringInfoString(&command, ")");
 

pg_lake_table/include/pg_lake/fdw/partition_pushdown.h

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2025 Snowflake Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "postgres.h"
+
+#include "pg_lake/iceberg/api/partitioning.h"
+#include "pg_lake/iceberg/manifest_spec.h"
+
+extern PGDLLEXPORT List *GetPartitionByExpressionsForRelation(Oid relationId);
+extern Partition * ParsePartitionValuesFromPartitionKeys(char *partitionKeysText,
+														 List *transforms);

pg_lake_table/include/pg_lake/fdw/partition_transform.h

@@ -27,7 +27,7 @@ extern Partition * ComputePartitionTupleForTuple(List *transforms, TupleTableSlo
 extern void *ApplyBucketTransformToColumn(IcebergPartitionTransform * transform,
 										  Datum columnValue, bool isNull,
 										  size_t *bucketSize);
-extern List *CurrentPartitionTransformList(Oid relationId);
+extern PGDLLEXPORT List *CurrentPartitionTransformList(Oid relationId);
 extern IcebergPartitionSpec * GetPartitionSpecIfAlreadyExist(Oid relationId, List *partitionTransforms);
 extern List *AllPartitionTransformList(Oid relationId);
 extern List *GetPartitionTransformsFromSpecFields(Oid relationId, List *specFields);