diff --git a/src/sempy_labs/lakehouse/__init__.py b/src/sempy_labs/lakehouse/__init__.py index 49f96d475..99a75ab77 100644 --- a/src/sempy_labs/lakehouse/__init__.py +++ b/src/sempy_labs/lakehouse/__init__.py @@ -41,6 +41,13 @@ is_schema_enabled, create_schema, ) +from ._partitioning import ( + is_partitioned, + list_partitioned_columns, + get_delta_table_details, + is_over_partitioned, + list_over_partitioned_tables, +) __all__ = [ "get_lakehouse_columns", @@ -68,4 +75,9 @@ "is_schema_enabled", "create_materialized_lake_view", "create_schema", + "is_partitioned", + "list_partitioned_columns", + "get_delta_table_details", + "is_over_partitioned", + "list_over_partitioned_tables", ] diff --git a/src/sempy_labs/lakehouse/_partitioning.py b/src/sempy_labs/lakehouse/_partitioning.py index 9570508df..fdc406207 100644 --- a/src/sempy_labs/lakehouse/_partitioning.py +++ b/src/sempy_labs/lakehouse/_partitioning.py @@ -1,14 +1,51 @@ from typing import Optional, List from uuid import UUID +import pandas as pd from sempy_labs._helper_functions import ( - _create_spark_session, - create_abfss_path, - resolve_workspace_id, - resolve_lakehouse_id, + _create_dataframe, _get_delta_table, ) from sempy._utils._log import log +BYTES_PER_GB = 1024**3 +DELTA_TABLE_DETAIL_COLUMNS = [ + "Table Name", + "Schema Name", + "Size In Bytes", + "Size In GB", + "Files", + "Partition Columns", + "Is Partitioned", +] +DELTA_TABLE_DETAIL_COLUMN_TYPES = { + "Table Name": "str", + "Schema Name": "str", + "Size In Bytes": "int", + "Size In GB": "float", + "Files": "int", + "Partition Columns": "object", + "Is Partitioned": "bool", +} + + +def _is_over_partitioned_by_details( + details: dict, total_table_size_gb: int, average_partition_size_gb: int +) -> bool: + total_size_gb = details["Size In GB"] + partitioned = details["Is Partitioned"] + num_files = details["Files"] + + # Only check if the table is partitioned + if partitioned and num_files > 0: + avg_partition_size_gb = total_size_gb / num_files + + return ( + total_size_gb < total_table_size_gb + or avg_partition_size_gb < average_partition_size_gb + ) + + return False + @log def _get_partitions( @@ -101,6 +138,54 @@ def list_partitioned_columns( return details["partitionColumns"] +@log +def get_delta_table_details( + table: str, + schema: Optional[str] = None, + lakehouse: Optional[str | UUID] = None, + workspace: Optional[str | UUID] = None, +) -> dict: + """ + Gets size and partition details for a delta table. + + Parameters + ---------- + table : str + The name of the delta table. + schema : str, optional + The schema of the table to check. If not provided, the default schema is used. + lakehouse : str | uuid.UUID, default=None + The Fabric lakehouse name or ID. + Defaults to None which resolves to the lakehouse attached to the notebook. + workspace : str | uuid.UUID, default=None + The Fabric workspace name or ID used by the lakehouse. + Defaults to None which resolves to the workspace of the attached lakehouse + or if no lakehouse attached, resolves to the workspace of the notebook. + + Returns + ------- + dict + A dictionary containing table size, file count and partition details. + """ + + details = _get_partitions( + table_name=table, schema_name=schema, lakehouse=lakehouse, workspace=workspace + ) + partition_columns = details.get("partitionColumns", []) + size_bytes = details.get("sizeInBytes", 0) + files = details.get("numFiles", 0) + + return { + "Table Name": table, + "Schema Name": schema, + "Size In Bytes": size_bytes, + "Size In GB": size_bytes / BYTES_PER_GB, + "Files": files, + "Partition Columns": partition_columns, + "Is Partitioned": len(partition_columns) > 0, + } + + @log def is_over_partitioned( table: str, @@ -137,29 +222,70 @@ def is_over_partitioned( True if the table is over-partitioned, False otherwise. """ - workspace_id = resolve_workspace_id(workspace) - lakehouse_id = resolve_lakehouse_id(lakehouse, workspace) - path = create_abfss_path(lakehouse_id, workspace_id, table, schema) - # Get DeltaTable details - spark = _create_spark_session() - details_df = spark.sql(f"DESCRIBE DETAIL delta.`{path}`") - details = details_df.collect()[0].asDict() + details = get_delta_table_details( + table=table, schema=schema, lakehouse=lakehouse, workspace=workspace + ) + return _is_over_partitioned_by_details( + details=details, + total_table_size_gb=total_table_size_gb, + average_partition_size_gb=average_partition_size_gb, + ) - # Extract relevant fields - size_bytes = details["sizeInBytes"] - partition_cols = details["partitionColumns"] - num_files = details["numFiles"] - total_size_gb = size_bytes / (1024**3) +@log +def list_over_partitioned_tables( + schema: Optional[str | List[str]] = None, + lakehouse: Optional[str | UUID] = None, + workspace: Optional[str | UUID] = None, + total_table_size_gb: int = 1000, + average_partition_size_gb: int = 1, +) -> pd.DataFrame: + """ + Lists over-partitioned delta tables in a lakehouse. - # Only check if the table is partitioned - if len(partition_cols) > 0 and num_files > 0: - avg_partition_size_gb = total_size_gb / num_files + Parameters + ---------- + schema : str | typing.List[str], optional + The schema name(s) used to filter tables. + lakehouse : str | uuid.UUID, default=None + The Fabric lakehouse name or ID. + Defaults to None which resolves to the lakehouse attached to the notebook. + workspace : str | uuid.UUID, default=None + The Fabric workspace name or ID used by the lakehouse. + Defaults to None which resolves to the workspace of the attached lakehouse + or if no lakehouse attached, resolves to the workspace of the notebook. + total_table_size_gb : int, default=1000 + Threshold for total table size in GB (default 1TB). + average_partition_size_gb : int, default=1 + Threshold for average partition size in GB. - if ( - total_size_gb < total_table_size_gb - or avg_partition_size_gb < average_partition_size_gb + Returns + ------- + pandas.DataFrame + A table of over-partitioned delta tables. + """ + from sempy_labs.lakehouse._schemas import list_tables + + df = list_tables(lakehouse=lakehouse, workspace=workspace, schema=schema) + over_partitioned_rows = [] + + for _, row in df.query("Format == 'delta'").iterrows(): + table = row["Table Name"] + table_schema = row["Schema Name"] + details = get_delta_table_details( + table=table, + schema=table_schema, + lakehouse=lakehouse, + workspace=workspace, + ) + if _is_over_partitioned_by_details( + details=details, + total_table_size_gb=total_table_size_gb, + average_partition_size_gb=average_partition_size_gb, ): - return True + over_partitioned_rows.append(details) - return False + if not over_partitioned_rows: + return _create_dataframe(columns=DELTA_TABLE_DETAIL_COLUMN_TYPES) + + return pd.DataFrame(over_partitioned_rows, columns=DELTA_TABLE_DETAIL_COLUMNS) diff --git a/tests/test_lakehouse_partitioning.py b/tests/test_lakehouse_partitioning.py new file mode 100644 index 000000000..8ea4d8e01 --- /dev/null +++ b/tests/test_lakehouse_partitioning.py @@ -0,0 +1,72 @@ +import pandas as pd + +import sempy_labs.lakehouse._partitioning as partitioning + + +def test_get_delta_table_details_returns_size_files_and_partition_state(monkeypatch): + def _mock_get_partitions(table_name, schema_name=None, lakehouse=None, workspace=None): + return { + "sizeInBytes": 2147483648, + "numFiles": 8, + "partitionColumns": ["event_date"], + } + + monkeypatch.setattr(partitioning, "_get_partitions", _mock_get_partitions) + + actual = partitioning.get_delta_table_details(table="sales") + + assert actual["Table Name"] == "sales" + assert actual["Size In Bytes"] == 2147483648 + assert actual["Size In GB"] == 2.0 + assert actual["Files"] == 8 + assert actual["Partition Columns"] == ["event_date"] + assert actual["Is Partitioned"] is True + + +def test_is_over_partitioned_uses_thresholds(monkeypatch): + monkeypatch.setattr( + partitioning, + "get_delta_table_details", + lambda **kwargs: {"Size In GB": 100, "Is Partitioned": True, "Files": 200}, + ) + + assert partitioning.is_over_partitioned(table="sales") is True + + +def test_list_over_partitioned_tables_returns_matching_tables(monkeypatch): + def _mock_list_tables(lakehouse=None, workspace=None, schema=None): + return pd.DataFrame( + [ + {"Table Name": "sales", "Schema Name": "dbo", "Format": "delta"}, + {"Table Name": "customers", "Schema Name": "dbo", "Format": "delta"}, + {"Table Name": "staging_sales", "Schema Name": "dbo", "Format": "csv"}, + ] + ) + + monkeypatch.setattr("sempy_labs.lakehouse._schemas.list_tables", _mock_list_tables) + called_tables = [] + + def _mock_get_delta_table_details( + table, schema, lakehouse=None, workspace=None + ): + called_tables.append(table) + return { + "Table Name": table, + "Schema Name": schema, + "Size In Bytes": 2147483648, + "Size In GB": 2, + "Files": 8 if table == "sales" else 1, + "Partition Columns": ["event_date"], + "Is Partitioned": table == "sales", + } + + monkeypatch.setattr( + partitioning, + "get_delta_table_details", + _mock_get_delta_table_details, + ) + + actual = partitioning.list_over_partitioned_tables() + + assert list(actual["Table Name"]) == ["sales"] + assert called_tables == ["sales", "customers"]