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/manifest.c

@@ -254,7 +254,7 @@ SetManifestPartitionSummary(IcebergManifest * manifest, List *manifestEntries, L
 
 		/* unexpected but better than crash */
 		if (partitionTransform == NULL)
-			ereport(ERROR, (errmsg("Could not find partition transform for field %d", partitionFieldId)));
+			ereport(ERROR, (errmsg("could not find partition transform for field %" PRId32, partitionFieldId)));
 
 		PGType		resultPgType = partitionTransform->resultPgType;
 		Field	   *resultField = PostgresTypeToIcebergField(resultPgType, false, NULL);

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_iceberg/src/iceberg/partitioning/partition.c

@@ -229,8 +229,5 @@ FindPartitionTransformById(List *transforms, int32_t partitionFieldId)
 			return transform;
 	}
 
-	ereport(ERROR,
-			(errcode(ERRCODE_INTERNAL_ERROR),
-			 errmsg("could not find partition transform for field id %d",
-					partitionFieldId)));
+	return NULL;
 }

pg_lake_iceberg/tests/conftest.py

@@ -55,31 +55,6 @@ def superuser_conn(postgres):
     conn.close()
 
 
-@pytest.fixture(scope="session")
-def iceberg_extension(postgres):
-    superuser_conn = open_pg_conn()
-
-    run_command(
-        f"""
-        CREATE EXTENSION IF NOT EXISTS pg_lake_iceberg CASCADE;
-    """,
-        superuser_conn,
-    )
-    superuser_conn.commit()
-
-    yield
-    superuser_conn.rollback()
-
-    run_command(
-        f"""
-        DROP EXTENSION pg_lake_iceberg CASCADE;
-    """,
-        superuser_conn,
-    )
-    superuser_conn.commit()
-    superuser_conn.close()
-
-
 @pytest.fixture(scope="module")
 def pgduck_conn(postgres):
     conn = psycopg2.connect(

pg_lake_iceberg/tests/pytests/test_iceberg_avro.py

@@ -7,7 +7,7 @@
 
 
 def test_iceberg_table_with_large_avro_allocs(
-    pg_conn, superuser_conn, extension, with_default_location, app_user
+    pg_conn, superuser_conn, iceberg_extension, with_default_location, app_user
 ):
     # 63 characters is the maximum length of a column name in PostgreSQL
     table_name = "".join(random.choices(string.ascii_lowercase, k=63))
@@ -149,7 +149,7 @@ def test_iceberg_table_with_large_avro_allocs(
 
 
 def test_manifest_file_with_dot_field_name(
-    pg_conn, superuser_conn, s3, extension, create_reserialize_helper_functions
+    pg_conn, superuser_conn, s3, iceberg_extension, create_reserialize_helper_functions
 ):
     key = "sample_avro/dot-field-name.avro"
     local_manifest_path = sample_avro_filepath(f"dot-field-name.avro")

pg_lake_iceberg/tests/pytests/test_iceberg_data_file_stats.py

@@ -164,7 +164,7 @@ def test_pg_lake_iceberg_table_reserialize_data_file_stats_from_metadata(
 def test_pg_lake_iceberg_table_read_data_file_stats_from_catalog(
     pg_lake_table_metadata_location,
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
 ):
@@ -588,7 +588,7 @@ def test_pg_lake_iceberg_table_special_numerics(
 
 def test_pg_lake_iceberg_table_add_drop_columns(
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -671,7 +671,7 @@ def test_pg_lake_iceberg_table_add_drop_columns(
 
 def test_pg_lake_iceberg_table_uuid_column(
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -704,7 +704,7 @@ def test_pg_lake_iceberg_table_uuid_column(
 
 def test_pg_lake_iceberg_table_bytea_column(
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -738,7 +738,7 @@ def test_pg_lake_iceberg_table_bytea_column(
 
 def test_pg_lake_iceberg_table_serial_column(
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -769,7 +769,7 @@ def test_pg_lake_iceberg_table_serial_column(
 
 def test_pg_lake_iceberg_table_random_values(
     pg_conn,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -864,7 +864,7 @@ def test_pg_lake_iceberg_table_random_values(
 def test_pg_lake_iceberg_table_complex_values(
     superuser_conn,
     enable_stats_for_nested_types,
-    extension,
+    iceberg_extension,
     s3,
     create_helper_functions,
     with_default_location,
@@ -1166,7 +1166,12 @@ def spark_table_metadata_location(installcheck, spark_session):
 
 @pytest.fixture(scope="function")
 def enable_stats_for_nested_types(
-    installcheck, extension, app_user, pg_conn, superuser_conn, with_default_location
+    installcheck,
+    iceberg_extension,
+    app_user,
+    pg_conn,
+    superuser_conn,
+    with_default_location,
 ):
     run_command_outside_tx(
         [
@@ -1194,7 +1199,7 @@ def enable_stats_for_nested_types(
 @pytest.fixture(scope="function")
 def pg_lake_table_metadata_location(
     installcheck,
-    extension,
+    iceberg_extension,
     app_user,
     pg_conn,
     superuser_conn,

pg_lake_iceberg/tests/pytests/test_iceberg_metadata_via_spark.py

@@ -85,7 +85,7 @@ def test_reserialize_via_spark(
     installcheck,
     spark_session,
     superuser_conn,
-    extension,
+    iceberg_extension,
     s3,
     spark_table_metadata_location,
     create_reserialize_helper_functions,

pg_lake_iceberg/tests/pytests/test_polaris_conn.py

@@ -252,84 +252,3 @@ def test_http_errors(
     assert "Could not resolve host" in error or "Resolving timed out after" in error
 
     pg_conn.rollback()
-
-
-@pytest.fixture(scope="module")
-def create_http_helper_functions(superuser_conn, iceberg_extension):
-    run_command(
-        f"""
-       CREATE TYPE lake_iceberg.http_result AS (
-            status        int,
-            body          text,
-            resp_headers  text
-        );
-
-        CREATE OR REPLACE FUNCTION lake_iceberg.test_http_get(
-                url     text,
-                headers text[] DEFAULT NULL)
-        RETURNS lake_iceberg.http_result
-        AS 'pg_lake_iceberg', 'test_http_get'
-        LANGUAGE C;
-
-
-        -- HEAD
-        CREATE OR REPLACE FUNCTION lake_iceberg.test_http_head(
-                url     text,
-                headers text[] DEFAULT NULL)
-        RETURNS lake_iceberg.http_result
-        AS 'pg_lake_iceberg', 'test_http_head'
-        LANGUAGE C;
-
-        -- POST
-        CREATE OR REPLACE FUNCTION lake_iceberg.test_http_post(
-                url     text,
-                body    text,
-                headers text[] DEFAULT NULL)
-        RETURNS lake_iceberg.http_result
-        AS 'pg_lake_iceberg', 'test_http_post'
-        LANGUAGE C;
-
-        -- PUT
-        CREATE OR REPLACE FUNCTION lake_iceberg.test_http_put(
-                url     text,
-                body    text,
-                headers text[] DEFAULT NULL)
-        RETURNS lake_iceberg.http_result
-        AS 'pg_lake_iceberg', 'test_http_put'
-        LANGUAGE C;
-
-        -- DELETE
-        CREATE OR REPLACE FUNCTION lake_iceberg.test_http_delete(
-                url     text,
-                headers text[] DEFAULT NULL)
-        RETURNS lake_iceberg.http_result
-        AS 'pg_lake_iceberg', 'test_http_delete'
-        LANGUAGE C;
-
-        -- URL encode function
-        CREATE OR REPLACE FUNCTION lake_iceberg.url_encode(input TEXT)
-        RETURNS text
-         LANGUAGE C
-         IMMUTABLE STRICT
-        AS 'pg_lake_iceberg', $function$url_encode_path$function$;
-
-""",
-        superuser_conn,
-    )
-    superuser_conn.commit()
-
-    yield
-
-    run_command(
-        """
-        DROP FUNCTION IF EXISTS lake_iceberg.url_encode;
-        DROP FUNCTION IF EXISTS lake_iceberg.test_http_get;
-        DROP FUNCTION IF EXISTS lake_iceberg.test_http_head;
-        DROP FUNCTION IF EXISTS lake_iceberg.test_http_post;
-        DROP FUNCTION IF EXISTS lake_iceberg.test_http_put;
-        DROP FUNCTION IF EXISTS lake_iceberg.test_http_delete;
-        DROP TYPE lake_iceberg.http_result;
-                """,
-        superuser_conn,
-    )
-    superuser_conn.commit()

pg_lake_table/include/pg_lake/partitioning/partition_spec_catalog.h

@@ -48,7 +48,7 @@ extern PGDLLEXPORT int GetCurrentSpecId(Oid relationId);
 extern int	GetLargestPartitionFieldId(Oid relationId);
 extern IcebergPartitionSpecField * GetIcebergPartitionFieldFromCatalog(Oid relationId, int fieldId);
 extern List *GetIcebergSpecPartitionFieldsFromCatalog(Oid relationId, int specId);
-extern List *GetAllIcebergSpecPartitionFieldsFromCatalog(Oid relationId);
+extern List *GetAllPartitionSpecFields(Oid relationId);
 extern Partition * GetDataFilePartition(Oid relationId, List *partitionTransforms,
 										const char *path, int32 *partitionSpecId);
 extern HTAB *GetAllPartitionSpecsFromCatalog(Oid relationId);