flatten check_compatible sort

Author: gabeiglio

pyiceberg/table/sorting.py

@@ -171,15 +171,14 @@ def __repr__(self) -> str:
         return f"SortOrder({fields}order_id={self.order_id})"
 
     def check_compatible(self, schema: Schema) -> None:
-        schema_ids = schema._lazy_id_to_field
         for field in self.fields:
-            if source_field := schema_ids.get(field.source_id):
-                if not source_field.field_type.is_primitive:
-                    raise ValidationError(f"Cannot sort by non-primitive source field: {source_field}")
-                if not field.transform.can_transform(source_field.field_type):
-                    raise ValidationError(f"Invalid source type {source_field.field_type} for transform: {field.transform}")
-            else:
+            source_field = schema._lazy_id_to_field.get(field.source_id)
+            if source_field is None:
                 raise ValidationError(f"Cannot find source column for sort field: {field}")
+            if not source_field.field_type.is_primitive:
+                raise ValidationError(f"Cannot sort by non-primitive source field: {source_field}")
+            if not field.transform.can_transform(source_field.field_type):
+                raise ValidationError(f"Invalid source type {source_field.field_type} for transform: {field.transform}")
 
 
 UNSORTED_SORT_ORDER_ID = 0