Ignore partition fields that are dropped from the current-schema

Author: Fokko

api/src/main/java/org/apache/iceberg/PartitionSpec.java

@@ -131,6 +131,12 @@ public StructType partitionType() {
           for (PartitionField field : fields) {
             Type sourceType = schema.findType(field.sourceId());
             Type resultType = field.transform().getResultType(sourceType);
+
+            // When the source field has been dropped we cannot determine the type
+            if (resultType == null) {
+              resultType = Types.UnknownType.get();
+            }
+
             structFields.add(Types.NestedField.optional(field.fieldId(), field.name(), resultType));
           }
 
@@ -633,8 +639,6 @@ static void checkCompatibility(PartitionSpec spec, Schema schema) {
       // We don't care about the source type since a VoidTransform is always compatible and skip the
       // checks
       if (!transform.equals(Transforms.alwaysNull())) {
-        ValidationException.check(
-            sourceType != null, "Cannot find source column for partition field: %s", field);
         ValidationException.check(
             sourceType.isPrimitiveType(),
             "Cannot partition by non-primitive source field: %s",

core/src/main/java/org/apache/iceberg/Partitioning.java

@@ -239,7 +239,8 @@ public static StructType groupingKeyType(Schema schema, Collection<PartitionSpec
    */
   public static StructType partitionType(Table table) {
     Collection<PartitionSpec> specs = table.specs().values();
-    return buildPartitionProjectionType("table partition", specs, allFieldIds(specs));
+    return buildPartitionProjectionType(
+        "table partition", specs, allActiveFieldIds(table.schema(), specs));
   }
 
   /**
@@ -346,10 +347,11 @@ private static boolean compatibleTransforms(Transform<?, ?> t1, Transform<?, ?>
         || t2.equals(Transforms.alwaysNull());
   }
 
-  // collects IDs of all partition field used across specs
-  private static Set<Integer> allFieldIds(Collection<PartitionSpec> specs) {
+  // collects IDs of all partition field used across specs that are in the current schema
+  private static Set<Integer> allActiveFieldIds(Schema schema, Collection<PartitionSpec> specs) {
     return FluentIterable.from(specs)
         .transformAndConcat(PartitionSpec::fields)
+        .filter(field -> schema.findField(field.sourceId()) != null)
         .transform(PartitionField::fieldId)
         .toSet();
   }

core/src/main/java/org/apache/iceberg/avro/AvroSchemaUtil.java

@@ -188,6 +188,8 @@ static Schema toOption(Schema schema) {
       Preconditions.checkArgument(
           isOptionSchema(schema), "Union schemas are not supported: %s", schema);
       return schema;
+    } else if (schema.getType() == Schema.Type.NULL) {
+      return schema;
     } else {
       return Schema.createUnion(NULL, schema);
     }

core/src/main/java/org/apache/iceberg/avro/TypeToSchema.java

@@ -49,6 +49,7 @@ abstract class TypeToSchema extends TypeUtil.SchemaVisitor<Schema> {
   private static final Schema UUID_SCHEMA =
       LogicalTypes.uuid().addToSchema(Schema.createFixed("uuid_fixed", null, null, 16));
   private static final Schema BINARY_SCHEMA = Schema.create(Schema.Type.BYTES);
+  private static final Schema NULL_SCHEMA = Schema.create(Schema.Type.NULL);
 
   static {
     TIMESTAMP_SCHEMA.addProp(AvroSchemaUtil.ADJUST_TO_UTC_PROP, false);
@@ -243,6 +244,9 @@ public Schema primitive(Type.PrimitiveType primitive) {
                         null,
                         TypeUtil.decimalRequiredBytes(decimal.precision())));
         break;
+      case UNKNOWN:
+        primitiveSchema = NULL_SCHEMA;
+        break;
       default:
         throw new UnsupportedOperationException("Unsupported type ID: " + primitive.typeId());
     }

spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestAlterTablePartitionFields.java

@@ -20,13 +20,17 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
+import java.util.List;
 import org.apache.iceberg.Parameter;
 import org.apache.iceberg.ParameterizedTestExtension;
 import org.apache.iceberg.Parameters;
 import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.TestHelpers;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.spark.SparkCatalogConfig;
 import org.apache.iceberg.spark.source.SparkTable;
 import org.apache.spark.sql.connector.catalog.CatalogManager;
@@ -582,4 +586,44 @@ private void createTable(String schema, String spec) {
           tableName, schema, spec, TableProperties.FORMAT_VERSION, formatVersion);
     }
   }
+
+  private void runCreateAndDropPartitionField(
+      String column, String partitionType, List<Object[]> expected, String predicate) {
+    sql("DROP TABLE IF EXISTS %s", tableName);
+    sql(
+        "CREATE TABLE %s (col_int INTEGER, col_ts TIMESTAMP_NTZ, col_long BIGINT) USING ICEBERG TBLPROPERTIES ('format-version' = %d)",
+        tableName, formatVersion);
+    sql("INSERT INTO %s VALUES (1000, CAST('2024-03-01 19:25:00' as TIMESTAMP), 2100)", tableName);
+    sql("ALTER TABLE %s ADD PARTITION FIELD %s AS col2_partition", tableName, partitionType);
+    sql("INSERT INTO %s VALUES (2000, CAST('2024-04-01 19:25:00' as TIMESTAMP), 2200)", tableName);
+    sql("ALTER TABLE %s DROP PARTITION FIELD col2_partition", tableName);
+    sql("INSERT INTO %s VALUES (3000, CAST('2024-05-01 19:25:00' as TIMESTAMP), 2300)", tableName);
+    sql("ALTER TABLE %s DROP COLUMN %s", tableName, column);
+
+    assertEquals(
+        "Should return correct data",
+        expected,
+        sql("SELECT * FROM %s WHERE %s ORDER BY col_int", tableName, predicate));
+  }
+
+  @TestTemplate
+  public void testDropPartitionAndUnderlyingField() {
+    String predicateLong = "col_ts >= '2024-04-01 19:25:00'";
+    List<Object[]> expectedLong =
+        Lists.newArrayList(
+            new Object[] {2000, LocalDateTime.ofEpochSecond(1711999500, 0, ZoneOffset.UTC)},
+            new Object[] {3000, LocalDateTime.ofEpochSecond(1714591500, 0, ZoneOffset.UTC)});
+    runCreateAndDropPartitionField("col_long", "col_long", expectedLong, predicateLong);
+    runCreateAndDropPartitionField(
+        "col_long", "truncate(2, col_long)", expectedLong, predicateLong);
+    runCreateAndDropPartitionField("col_long", "bucket(16, col_long)", expectedLong, predicateLong);
+
+    String predicateTs = "col_long >= 2200";
+    List<Object[]> expectedTs =
+        Lists.newArrayList(new Object[] {2000, 2200L}, new Object[] {3000, 2300L});
+    runCreateAndDropPartitionField("col_ts", "col_ts", expectedTs, predicateTs);
+    runCreateAndDropPartitionField("col_ts", "year(col_ts)", expectedTs, predicateTs);
+    runCreateAndDropPartitionField("col_ts", "month(col_ts)", expectedTs, predicateTs);
+    runCreateAndDropPartitionField("col_ts", "day(col_ts)", expectedTs, predicateTs);
+  }
 }