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,10 +653,19 @@ 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)
-		return false;
+	{
+		List	   *transforms = CurrentPartitionTransformList(relationId);
+
+		if (!AllPartitionTransformsPushdownable(transforms))
+			return false;
+	}
 
 	return true;
 }

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/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 bool AllPartitionTransformsPushdownable(List *transforms);
+extern List *GetPartitionByExpressions(List *transforms);
+extern Partition * ParsePartitionValuesFromPath(char *filePath, 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);

pg_lake_table/src/fdw/partition_pushdown.c

@@ -0,0 +1,218 @@
+/*
+ * 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.
+ */
+
+/*
+ * partition_pushdown.c
+ *
+ * Utilities for pushing down partitioned Iceberg writes to DuckDB using
+ * the PARTITION_BY clause in COPY TO. Supports identity and temporal
+ * (year, month, day, hour) partition transforms.
+ */
+#include "postgres.h"
+
+#include "utils/builtins.h"
+
+#include "pg_lake/fdw/partition_pushdown.h"
+#include "pg_lake/fdw/partition_transform.h"
+#include "pg_lake/iceberg/api/partitioning.h"
+#include "pg_lake/iceberg/manifest_spec.h"
+
+/* prefix for synthetic partition columns in COPY TO queries */
+#define PARTITION_COLUMN_PREFIX "__part_"
+
+
+static char *PartitionTransformToDuckDBExpression(IcebergPartitionTransform * transform);
+
+
+/*
+ * AllPartitionTransformsPushdownable returns true if every transform in the
+ * list can be expressed as a DuckDB SQL expression for PARTITION_BY.
+ */
+bool
+AllPartitionTransformsPushdownable(List *transforms)
+{
+	if (transforms == NIL)
+		return false;
+
+	ListCell   *cell = NULL;
+
+	foreach(cell, transforms)
+	{
+		IcebergPartitionTransform *transform = lfirst(cell);
+		char	   *expr = PartitionTransformToDuckDBExpression(transform);
+
+		if (expr == NULL)
+			return false;
+	}
+
+	return true;
+}
+
+
+/*
+ * PartitionTransformToDuckDBExpression returns a DuckDB SQL expression that
+ * computes the Iceberg partition value for the given transform.
+ *
+ * Returns NULL for transforms that cannot be pushed down (bucket, truncate, void).
+ *
+ * The expressions produce Iceberg-compatible partition values:
+ * - year: integer years since 1970
+ * - month: integer months since Jan 1970
+ * - day: integer days since 1970-01-01
+ * - hour: integer hours since 1970-01-01T00:00:00
+ * - identity: the column value (with type-specific casts for date/timestamp)
+ */
+static char *
+PartitionTransformToDuckDBExpression(IcebergPartitionTransform * transform)
+{
+	const char *col = quote_identifier(transform->columnName);
+	Oid			typeOid = transform->pgType.postgresTypeOid;
+
+	switch (transform->type)
+	{
+		case PARTITION_TRANSFORM_IDENTITY:
+			{
+				/*
+				 * Identity partitions use the column value directly.  DuckDB
+				 * writes the value in its default text format into the
+				 * Hive-style path (e.g. __part_0=2025-01-15), and
+				 * DeserializePartitionValueFromPGText parses it back using
+				 * the transform's result PG type.
+				 */
+				return psprintf("%s", col);
+			}
+
+		case PARTITION_TRANSFORM_YEAR:
+			return psprintf("(year(%s) - 1970)", col);
+
+		case PARTITION_TRANSFORM_MONTH:
+			return psprintf("((year(%s) - 1970) * 12 + month(%s) - 1)", col, col);
+
+		case PARTITION_TRANSFORM_DAY:
+			return psprintf("datediff('day', date '1970-01-01', %s::date)", col);
+
+		case PARTITION_TRANSFORM_HOUR:
+			return psprintf("datediff('hour', timestamp '1970-01-01', %s::timestamp)", col);
+
+		case PARTITION_TRANSFORM_BUCKET:
+		case PARTITION_TRANSFORM_TRUNCATE:
+		case PARTITION_TRANSFORM_VOID:
+			return NULL;
+	}
+
+	return NULL;
+}
+
+
+/*
+ * GetPartitionByExpressions returns a list of String values with the DuckDB
+ * SQL expressions for each partition transform. These are passed to
+ * WriteQueryResultTo which wraps the query with synthetic partition columns
+ * after validation wrapping.
+ */
+List *
+GetPartitionByExpressions(List *transforms)
+{
+	List	   *exprs = NIL;
+	ListCell   *cell = NULL;
+
+	foreach(cell, transforms)
+	{
+		IcebergPartitionTransform *transform = lfirst(cell);
+		char	   *expr = PartitionTransformToDuckDBExpression(transform);
+
+		Assert(expr != NULL);
+
+		exprs = lappend(exprs, makeString(expr));
+	}
+
+	return exprs;
+}
+
+
+/*
+ * ParsePartitionValuesFromPath extracts partition values from the Hive-style
+ * directory path produced by DuckDB COPY TO with PARTITION_BY.
+ *
+ * A path like:
+ *   s3://bucket/data/abc123/__part_0=54/__part_1=us-east/data_0.parquet
+ *
+ * is parsed to extract __part_0=54 and __part_1=us-east, which are then
+ * converted to the proper Iceberg binary format using the partition transforms.
+ */
+Partition *
+ParsePartitionValuesFromPath(char *filePath, List *transforms)
+{
+	int			numTransforms = list_length(transforms);
+	Partition  *partition = palloc0(sizeof(Partition));
+
+	partition->fields = palloc0(sizeof(PartitionField) * numTransforms);
+	partition->fields_length = numTransforms;
+
+	for (int partIndex = 0; partIndex < numTransforms; partIndex++)
+	{
+		IcebergPartitionTransform *transform = list_nth(transforms, partIndex);
+
+		/* build the search key: "__part_N=" */
+		char	   *searchKey = psprintf(PARTITION_COLUMN_PREFIX "%d=", partIndex);
+		int			searchKeyLen = strlen(searchKey);
+
+		/* find this key in the path */
+		char	   *found = strstr(filePath, searchKey);
+
+		if (found == NULL)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_INTERNAL_ERROR),
+					 errmsg("could not find partition key %s in path %s",
+							searchKey, filePath)));
+		}
+
+		/* extract the value (from after '=' up to the next '/' or end) */
+		char	   *valueStart = found + searchKeyLen;
+		char	   *valueEnd = strchr(valueStart, '/');
+		int			valueLen = (valueEnd != NULL) ?
+			(valueEnd - valueStart) : strlen(valueStart);
+
+		char	   *valueText = pnstrdup(valueStart, valueLen);
+
+		/* populate the partition field */
+		PartitionField *field = &partition->fields[partIndex];
+
+		field->field_id = transform->partitionFieldId;
+		field->field_name = pstrdup(transform->partitionFieldName);
+		field->value_type = GetTransformResultAvroType(transform);
+
+		if (strcmp(valueText, "NULL") == 0)
+		{
+			/* NULL partition value */
+			field->value = NULL;
+			field->value_length = 0;
+		}
+		else
+		{
+			/*
+			 * Convert the text value to Iceberg binary format. The text
+			 * representation matches the transform's result PG type.
+			 */
+			field->value = DeserializePartitionValueFromPGText(
+															   transform, valueText, &field->value_length);
+		}
+	}
+
+	return partition;
+}