jqin61 commented on code in PR #555:
URL: https://github.com/apache/iceberg-python/pull/555#discussion_r1544756546
##########
pyiceberg/table/__init__.py:
##########
@@ -3108,3 +3138,127 @@ def snapshots(self) -> "pa.Table":
snapshots,
schema=snapshots_schema,
)
+
+
+@dataclass(frozen=True)
+class TablePartition:
+ partition_key: PartitionKey
+ arrow_table_partition: pa.Table
+
+
+def _get_partition_sort_order(partition_columns: list[str], reverse: bool =
False) -> dict[str, Any]:
+ order = 'ascending' if not reverse else 'descending'
+ null_placement = 'at_start' if reverse else 'at_end'
+ return {'sort_keys': [(column_name, order) for column_name in
partition_columns], 'null_placement': null_placement}
+
+
+def group_by_partition_scheme(
+ iceberg_table_metadata: TableMetadata, arrow_table: pa.Table,
partition_columns: list[str]
+) -> pa.Table:
+ """Given a table sort it by current partition scheme with all transform
functions supported."""
+ from pyiceberg.transforms import IdentityTransform
+
+ supported = {IdentityTransform}
+ if not all(
+ type(field.transform) in supported for field in
iceberg_table_metadata.spec().fields if field in partition_columns
+ ):
+ raise ValueError(
+ f"Not all transforms are supported, get: {[transform in supported
for transform in iceberg_table_metadata.spec().fields]}."
+ )
+
+ # only works for identity
+ sort_options = _get_partition_sort_order(partition_columns, reverse=False)
+ sorted_arrow_table =
arrow_table.sort_by(sorting=sort_options['sort_keys'],
null_placement=sort_options['null_placement'])
+ return sorted_arrow_table
+
+
+def get_partition_columns(iceberg_table_metadata: TableMetadata, arrow_table:
pa.Table) -> list[str]:
+ arrow_table_cols = set(arrow_table.column_names)
+ partition_cols = []
+ for transform_field in iceberg_table_metadata.spec().fields:
Review Comment:
YES 👍 I changed the name in a latter commit in the origin PR. Overlooked it
here. Resolved in the incoming commit.
--
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
