small refactor

Author: gabeiglio

pyiceberg/partitioning.py

@@ -267,20 +267,17 @@ def validate_parents_are_structs(field_id: int) -> None:
                 continue
 
             if not isinstance(field.transform, VoidTransform):
-                if not is_source_column_in_schema:
+                if source_field := schema_fields.get(field.source_id):
+                    source_type = source_field.field_type
+                    if not source_type.is_primitive:
+                        raise ValidationError(f"Cannot partition by non-primitive source field: {source_type}")
+                    if not field.transform.can_transform(source_type):
+                        raise ValidationError(f"Invalid source type {source_type} for transform: {field.transform}")
+                    # The only valid parent types for a PartitionField are StructTypes. This must be checked recursively
+                    validate_parents_are_structs(field.source_id)
+                else:
                     raise ValidationError(f"Cannot find source column for partition field: {field}")
 
-                source_type = schema.find_type(field.source_id)
-
-                if not source_type.is_primitive:
-                    raise ValidationError(f"Cannot partition by non-primitive source field: {source_type}")
-
-                if not field.transform.can_transform(source_type):
-                    raise ValidationError(f"Invalid source type {source_type} for transform: {field.transform}")
-
-                # The only valid parent types for a PartitionField are StructTypes. This must be checked recursively
-                validate_parents_are_structs(field.source_id)
-
 
 UNPARTITIONED_PARTITION_SPEC = PartitionSpec(spec_id=0)