Support pruning for external tables

Signed-off-by: Aykut Bozkurt <[email protected]>

Author: sfc-gh-abozkurt

pg_lake_iceberg/include/pg_lake/iceberg/api/table_schema.h

@@ -27,6 +27,7 @@ extern PGDLLEXPORT IcebergTableSchema * GetIcebergTableSchemaByIdFromTableMetada
 extern PGDLLEXPORT IcebergTableSchema * GetCurrentIcebergTableSchema(IcebergTableMetadata * metadata);
 extern PGDLLEXPORT List *GetLeafFieldsFromIcebergMetadata(IcebergTableMetadata * metadata);
 extern PGDLLEXPORT List *GetLeafFieldsForIcebergSchema(IcebergTableSchema * schema);
+extern PGDLLEXPORT DataFileSchemaField * GetDataFileSchemaFieldById(DataFileSchema * schema, int fieldId);
 
 /* write api */
 extern PGDLLEXPORT IcebergTableSchema * RebuildIcebergSchemaFromDataFileSchema(Oid foreignTableOid,

pg_lake_iceberg/src/iceberg/api/table_schema.c

@@ -169,6 +169,38 @@ GetLeafFieldsForIcebergSchema(IcebergTableSchema * schema)
 }
 
 
+/*
+ * GetDataFileSchemaFieldById gets the DataFileSchemaField for the given
+ * iceberg field id.
+ */
+DataFileSchemaField *
+GetDataFileSchemaFieldById(DataFileSchema * schema, int fieldId)
+{
+	DataFileSchemaField *schemaField = NULL;
+
+	for (size_t fieldIdx = 0; fieldIdx < schema->nfields; fieldIdx++)
+	{
+		DataFileSchemaField *field = &schema->fields[fieldIdx];
+
+		if (field->id == fieldId)
+		{
+			schemaField = field;
+			break;
+		}
+	}
+
+	if (schemaField == NULL)
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("field ID %d not found",
+						fieldId)));
+	}
+
+	return schemaField;
+}
+
+
 /*
  * GetLeafFieldsForField returns the leaf fields for the given field.
  * It recursively traverses the field tree and collects all leaf fields, and

pg_lake_table/include/pg_lake/fdw/partition_transform.h

@@ -29,6 +29,7 @@ extern void *ApplyBucketTransformToColumn(IcebergPartitionTransform * transform,
 										  size_t *bucketSize);
 extern List *CurrentPartitionTransformList(Oid relationId);
 extern List *AllPartitionTransformList(Oid relationId);
+extern List *PartitionTransformListForSourceFields(Oid relationId, HTAB *sourceFieldsHash);
 extern List *GetPartitionTransformsFromSpecFields(Oid relationId, List *specFields);
 extern void *DeserializePartitionValueFromPGText(IcebergPartitionTransform * transform,
 												 const char *valueText, size_t *valueLength);

pg_lake_table/src/fdw/data_file_pruning.c

@@ -182,7 +182,6 @@ PruneDataFiles(Oid relationId, List *dataFiles, List *baseRestrictInfoList, Prun
 {
 	List	   *retainedDataFiles = NIL;
 	List	   *columnsUsedInFilters = ColumnsUsedInRestrictions(relationId, baseRestrictInfoList);
-	List	   *partitionTransforms = AllPartitionTransformList(relationId);
 	PgLakeTableProperties tableProperties = GetPgLakeTableProperties(relationId);
 
 	if ((!EnableDataFilePruning && !EnablePartitionPruning) ||
@@ -230,6 +229,9 @@ PruneDataFiles(Oid relationId, List *dataFiles, List *baseRestrictInfoList, Prun
 
 	AddFieldIdsUsedInQuery(fieldIdsUsedInQuery, relationId, tableProperties, columnsUsedInFilters);
 
+	/* get partition transforms for the fields used in the query */
+	List	   *partitionTransforms = PartitionTransformListForSourceFields(relationId, fieldIdsUsedInQuery);
+
 	int			dataFileCount = list_length(dataFiles);
 
 	for (int dataFileIndex = 0; dataFileIndex < dataFileCount; ++dataFileIndex)
@@ -254,6 +256,7 @@ PruneDataFiles(Oid relationId, List *dataFiles, List *baseRestrictInfoList, Prun
 			char	   *dataFilePath = (char *) ((DataFile *) dataFile)->file_path;
 
 			columnStats = GetRemoteParquetColumnStats(dataFilePath, leafFields);
+			partition = CopyPartition(&dataFile->partition);
 		}
 		else
 		{

pg_lake_table/src/fdw/partition_transform.c

@@ -25,7 +25,9 @@
 #include "pg_lake/pgduck/numeric.h"
 #include "pg_lake/fdw/partition_transform.h"
 #include "pg_lake/fdw/schema_operations/field_id_mapping_catalog.h"
+#include "pg_lake/fdw/schema_operations/register_field_ids.h"
 #include "pg_lake/iceberg/api/partitioning.h"
+#include "pg_lake/iceberg/catalog.h"
 #include "pg_lake/iceberg/iceberg_field.h"
 #include "pg_lake/iceberg/iceberg_type_binary_serde.h"
 #include "pg_lake/iceberg/iceberg_type_numeric_binary_serde.h"
@@ -35,6 +37,7 @@
 #include "pg_lake/iceberg/hash_utils.h"
 #include "pg_lake/iceberg/truncate_utils.h"
 #include "pg_lake/util/numeric.h"
+#include "pg_lake/util/rel_utils.h"
 
 static PartitionField * ApplyPartitionTransformToTuple(IcebergPartitionTransform * transform,
 													   TupleTableSlot *slot);
@@ -63,6 +66,8 @@ static void ParseTransformName(const char *name, IcebergPartitionTransformType *
 static bool ParseBracketUintSize(const char *name, const char *prefix, size_t *outVal);
 static void *DatumToPartitionValue(IcebergPartitionTransform * transform, Datum columnValue, bool isNull,
 								   size_t *valuelength);
+static List *AllPartitionSpecFieldsForExternalIcebergTable(char *metadataPath);
+static List *GetAllPartitionSpecFields(Oid relationId);
 
 /*
  * ComputePartitionTupleForTuple applies relative partition transforms
@@ -125,18 +130,94 @@ CurrentPartitionTransformList(Oid relationId)
 List *
 AllPartitionTransformList(Oid relationId)
 {
-	List	   *partitionFields = GetAllIcebergSpecPartitionFieldsFromCatalog(relationId);
+	List	   *partitionFields = GetAllPartitionSpecFields(relationId);
 
-	if (partitionFields == NIL)
+	return GetPartitionTransformsFromSpecFields(relationId, partitionFields);
+}
+
+
+/*
+ * PartitionTransformListForSourceFields gets the partition transforms for only the fields
+ * present in the sourceFieldsHash.
+ */
+List *
+PartitionTransformListForSourceFields(Oid relationId, HTAB *sourceFieldsHash)
+{
+	List	   *partitionFields = GetAllPartitionSpecFields(relationId);
+
+	List	   *filteredPartitionFields = NIL;
+	ListCell   *filteredPartitionFieldCell = NULL;
+
+	foreach(filteredPartitionFieldCell, partitionFields)
 	{
-		/* not partitioned */
-		return NIL;
+		IcebergPartitionSpecField *specField = lfirst(filteredPartitionFieldCell);
+
+		bool		found = false;
+
+		hash_search(sourceFieldsHash, &specField->source_id, HASH_FIND, &found);
+
+		if (found)
+		{
+			filteredPartitionFields = lappend(filteredPartitionFields, specField);
+		}
 	}
 
-	return GetPartitionTransformsFromSpecFields(relationId, partitionFields);
+	return GetPartitionTransformsFromSpecFields(relationId, filteredPartitionFields);
 }
 
 
+/*
+ * GetAllPartitionSpecFields gets all the partition spec fields for the given
+ * relationId.
+ */
+static List *
+GetAllPartitionSpecFields(Oid relationId)
+{
+	List	   *partitionFields = NIL;
+
+	if (IsInternalIcebergTable(relationId))
+	{
+		/* internal iceberg table, get from catalog */
+		partitionFields = GetAllIcebergSpecPartitionFieldsFromCatalog(relationId);
+	}
+	else if (IsExternalIcebergTable(relationId))
+	{
+		/* external iceberg table, get from remote metadata */
+		char	   *currentMetadataPath = GetIcebergMetadataLocation(relationId, false);
+
+		partitionFields = AllPartitionSpecFieldsForExternalIcebergTable(currentMetadataPath);
+	}
+
+	return partitionFields;
+}
+
+
+/*
+ * AllPartitionSpecFieldsForExternalIcebergTable gets all the partition spec fields
+ * for the external Iceberg table.
+ */
+static List *
+AllPartitionSpecFieldsForExternalIcebergTable(char *metadataPath)
+{
+	IcebergTableMetadata *metadata = ReadIcebergTableMetadata(metadataPath);
+
+	List	   *allPartitionFields = NIL;
+
+	for (int specIdx = 0; specIdx < metadata->partition_specs_length; specIdx++)
+	{
+		IcebergPartitionSpec *spec = &metadata->partition_specs[specIdx];
+
+		for (int partitionFieldIdx = 0; partitionFieldIdx < spec->fields_length; partitionFieldIdx++)
+		{
+			IcebergPartitionSpecField *field = &spec->fields[partitionFieldIdx];
+
+			allPartitionFields = lappend(allPartitionFields, field);
+		}
+	}
+
+	return allPartitionFields;
+}
+
 
 /*
 * GetPartitionTransformsFromSpecFields is a wrapper around
@@ -184,7 +265,19 @@ GetPartitionTransformFromSpecField(Oid relationId, IcebergPartitionSpecField * s
 		GetAttributeForFieldId(relationId, specField->source_id);
 	transform->columnName = get_attname(relationId, transform->attnum, false);
 	transform->pgType = GetAttributePGType(relationId, transform->attnum);
-	transform->sourceField = GetRegisteredFieldForAttribute(relationId, transform->attnum);
+
+	if (IsInternalIcebergTable(relationId))
+	{
+		transform->sourceField = GetRegisteredFieldForAttribute(relationId, transform->attnum);
+	}
+	else
+	{
+		Assert(IsExternalIcebergTable(relationId));
+
+		DataFileSchema *schema = GetDataFileSchemaForTable(relationId);
+
+		transform->sourceField = GetDataFileSchemaFieldById(schema, specField->source_id);
+	}
 
 	/* parse transform name */
 	ParseTransformName(transform->transformName,

pg_lake_table/src/fdw/schema_operations/field_id_mapping_catalog.c

@@ -41,19 +41,23 @@
 #include "pg_lake/fdw/schema_operations/field_id_mapping_catalog.h"
 #include "pg_lake/fdw/schema_operations/register_field_ids.h"
 #include "pg_lake/iceberg/api/table_metadata.h"
+#include "pg_lake/iceberg/api/table_schema.h"
 #include "pg_lake/iceberg/catalog.h"
 #include "pg_lake/iceberg/iceberg_field.h"
 #include "pg_lake/parquet/leaf_field.h"
 #include "pg_lake/pgduck/map.h"
 #include "pg_lake/pgduck/serialize.h"
 #include "pg_lake/util/array_utils.h"
+#include "pg_lake/util/rel_utils.h"
 #include "pg_lake/util/spi_helpers.h"
 
 static DataFileSchemaField * CreateRegisteredFieldForAttribute(Oid relationId, int spiIndex);
 static void InsertFieldMapping(Oid relationId, int attrIcebergFieldId,
 							   AttrNumber pg_attnum, PGType pgType,
 							   const char *writeDefault, const char *initialDefault,
 							   int parentFieldId);
+static AttrNumber GetAttributeForFieldIdForInternalIcebergTable(Oid relationId, int fieldId);
+static AttrNumber GetAttributeForFieldIdForExternalIcebergTable(char *metadataPath, Oid relationId, int fieldId);
 
 #ifdef USE_ASSERT_CHECKING
 static List *GetAllRegisteredAttnumsForTopLevelColumns(Oid relationId);
@@ -244,6 +248,28 @@ GetRegisteredFieldForAttribute(Oid relationId, AttrNumber attrNo)
 */
 AttrNumber
 GetAttributeForFieldId(Oid relationId, int fieldId)
+{
+	if (IsInternalIcebergTable(relationId))
+	{
+		return GetAttributeForFieldIdForInternalIcebergTable(relationId, fieldId);
+	}
+	else
+	{
+		Assert(IsExternalIcebergTable(relationId));
+
+		char	   *currentMetadataPath = GetIcebergMetadataLocation(relationId, false);
+
+		return GetAttributeForFieldIdForExternalIcebergTable(currentMetadataPath, relationId, fieldId);
+	}
+}
+
+
+/*
+* GetAttributeForFieldIdForInternalIcebergTable gets the attribute number for a given field ID
+* for internal Iceberg tables from catalog.
+*/
+static AttrNumber
+GetAttributeForFieldIdForInternalIcebergTable(Oid relationId, int fieldId)
 {
 	DECLARE_SPI_ARGS(2);
 
@@ -280,6 +306,54 @@ GetAttributeForFieldId(Oid relationId, int fieldId)
 }
 
 
+/*
+* GetAttributeForFieldIdForExternalIcebergTable gets the attribute number for a given field ID
+* for external Iceberg tables from iceberg metadata.
+ */
+static AttrNumber
+GetAttributeForFieldIdForExternalIcebergTable(char *metadataPath, Oid relationId, int fieldId)
+{
+	DataFileSchema *schema = GetDataFileSchemaForExternalIcebergTable(metadataPath);
+
+	DataFileSchemaField *schemaField = GetDataFileSchemaFieldById(schema, fieldId);
+
+	char	   *attrName = pstrdup(schemaField->name);
+
+	DECLARE_SPI_ARGS(2);
+
+	SPI_ARG_VALUE(1, OIDOID, relationId, false);
+	SPI_ARG_VALUE(2, TEXTOID, attrName, false);
+
+	SPI_START();
+
+	/*
+	 * Although this is a read-only query, we need the execution to use the
+	 * current transaction's snapshot (e.g., GetTransactionSnapshot()) to get
+	 * the snapshot that the current transaction modified.
+	 *
+	 * So we trick the SPI_EXECUTE function to think that the query is not
+	 * read-only and read the transaction snapshot.
+	 */
+	bool		readOnly = false;
+
+	SPI_EXECUTE("SELECT attnum FROM pg_attribute "
+				"WHERE attrelid OPERATOR(pg_catalog.=) $1 "
+				"AND attname OPERATOR(pg_catalog.=) $2 "
+				"AND NOT attisdropped", readOnly);
+
+	/* there is a primary key on these filters */
+	Assert(SPI_processed == 1);
+
+	bool		isNull = false;
+	AttrNumber	attrNo = GET_SPI_VALUE(INT2OID, 0, 1, &isNull);
+
+	Assert(!isNull);
+
+	SPI_END();
+
+	return attrNo;
+}
+
 
 /*
  * GetDataFileSchemaForInternalIcebergTable gets a table schema based on the MAPPING_TABLE_NAME.