Fokko commented on code in PR #245:
URL: https://github.com/apache/iceberg-python/pull/245#discussion_r1466147869
##########
pyiceberg/partitioning.py:
##########
@@ -215,3 +235,111 @@ def assign_fresh_partition_spec_ids(spec: PartitionSpec,
old_schema: Schema, fre
)
)
return PartitionSpec(*partition_fields, spec_id=INITIAL_PARTITION_SPEC_ID)
+
+
+T = TypeVar("T")
+
+
+class PartitionSpecVisitor(Generic[T], ABC):
+ @abstractmethod
+ def identity(self, field_id: int, source_name: str, source_id: int) -> T:
+ """Visit identity partition field."""
+
+ @abstractmethod
+ def bucket(self, field_id: int, source_name: str, source_id: int,
num_buckets: int) -> T:
+ """Visit bucket partition field."""
+
+ @abstractmethod
+ def truncate(self, field_id: int, source_name: str, source_id: int, width:
int) -> T:
+ """Visit truncate partition field."""
+
+ @abstractmethod
+ def year(self, field_id: int, source_name: str, source_id: int) -> T:
+ """Visit year partition field."""
+
+ @abstractmethod
+ def month(self, field_id: int, source_name: str, source_id: int) -> T:
+ """Visit month partition field."""
+
+ @abstractmethod
+ def day(self, field_id: int, source_name: str, source_id: int) -> T:
+ """Visit day partition field."""
+
+ @abstractmethod
+ def hour(self, field_id: int, source_name: str, source_id: int) -> T:
+ """Visit hour partition field."""
+
+ @abstractmethod
+ def always_null(self, field_id: int, source_name: str, source_id: int) ->
T:
+ """Visit void partition field."""
+
+ @abstractmethod
+ def unknown(self, field_id: int, source_name: str, source_id: int,
transform: str) -> T:
+ """Visit unknown partition field."""
+ raise ValueError(f"Unknown transform {transform} is not supported")
+
+
+class _PartitionNameGenerator(PartitionSpecVisitor[str]):
+ def identity(self, field_id: int, source_name: str, source_id: int) -> str:
+ return source_name
+
+ def bucket(self, field_id: int, source_name: str, source_id: int,
num_buckets: int) -> str:
+ return source_name + "_bucket_" + str(num_buckets)
+
+ def truncate(self, field_id: int, source_name: str, source_id: int, width:
int) -> str:
+ return source_name + "_trunc_" + str(width)
+
+ def year(self, field_id: int, source_name: str, source_id: int) -> str:
+ return source_name + "_year"
+
+ def month(self, field_id: int, source_name: str, source_id: int) -> str:
+ return source_name + "_month"
+
+ def day(self, field_id: int, source_name: str, source_id: int) -> str:
+ return source_name + "_day"
+
+ def hour(self, field_id: int, source_name: str, source_id: int) -> str:
+ return source_name + "_hour"
+
+ def always_null(self, field_id: int, source_name: str, source_id: int) ->
str:
+ return source_name + "_null"
+
+ def unknown(self, field_id: int, source_name: str, source_id: int,
transform: str) -> str:
+ return super().unknown(field_id, source_name, source_id, transform)
+
+
+R = TypeVar("R")
+
+
+@singledispatch
+def _visit(spec: PartitionSpec, schema: Schema, visitor:
PartitionSpecVisitor[R]) -> List[R]:
+ results = []
+ for field in spec.fields:
+ results.append(_visit_field(schema, field, visitor))
+ return results
Review Comment:
Nit: using for-comprehension you can directly create the list:
```suggestion
return [_visit_field(schema, field, visitor) for field in spec.fields]
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For additional commands, e-mail: issues-h...@iceberg.apache.org
