semantic-link-labs 0.9.9__py3-none-any.whl → 0.9.11__py3-none-any.whl

sempy_labs/_helper_functions.py

@@ -8,7 +8,7 @@ from sempy.fabric.exceptions import FabricHTTPException, WorkspaceNotFoundExcept
 import pandas as pd
 from functools import wraps
 import datetime
-from typing import Optional, Tuple, List
+from typing import Optional, Tuple, List, Dict
 from uuid import UUID
 import sempy_labs._icons as icons
 from azure.core.credentials import TokenCredential, AccessToken
@@ -74,6 +74,15 @@ def create_abfss_path(
     return path
 
 
+def create_abfss_path_from_path(
+    lakehouse_id: UUID, workspace_id: UUID, file_path: str
+) -> str:
+
+    fp = _get_default_file_path()
+
+    return f"abfss://{workspace_id}@{fp}/{lakehouse_id}/{file_path}"
+
+
 def _get_default_file_path() -> str:
 
     default_file_storage = _get_fabric_context_setting(name="fs.defaultFS")
@@ -266,7 +275,7 @@ def create_item(
         lro_return_status_code=True,
     )
     print(
-        f"{icons.green_dot} The '{name}' {item_type} has been successfully created within the in the '{workspace_name}' workspace."
+        f"{icons.green_dot} The '{name}' {item_type} has been successfully created within the '{workspace_name}' workspace."
     )
 
 
@@ -278,10 +287,9 @@ def get_item_definition(
     return_dataframe: bool = True,
     decode: bool = True,
 ):
-
     from sempy_labs._utils import item_types
 
-    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+    workspace_id = resolve_workspace_id(workspace)
     item_id = resolve_item_id(item, type, workspace_id)
     item_type_url = item_types.get(type)[1]
     path = item_types.get(type)[2]
@@ -304,92 +312,11 @@ def get_item_definition(
         p.get("payload") for p in result["definition"]["parts"] if p.get("path") == path
     )
     if decode:
-        json.loads(_decode_b64(value))
+        return json.loads(_decode_b64(value))
     else:
         return value
 
 
-def resolve_item_id(
-    item: str | UUID, type: Optional[str] = None, workspace: Optional[str | UUID] = None
-) -> UUID:
-
-    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
-    item_id = None
-
-    if _is_valid_uuid(item):
-        # Check (optional)
-        item_id = item
-        try:
-            _base_api(
-                request=f"/v1/workspaces/{workspace_id}/items/{item_id}",
-                client="fabric_sp",
-            )
-        except FabricHTTPException:
-            raise ValueError(
-                f"{icons.red_dot} The '{item_id}' item was not found in the '{workspace_name}' workspace."
-            )
-    else:
-        if type is None:
-            raise ValueError(
-                f"{icons.red_dot} The 'type' parameter is required if specifying an item name."
-            )
-        responses = _base_api(
-            request=f"/v1/workspaces/{workspace_id}/items?type={type}",
-            client="fabric_sp",
-            uses_pagination=True,
-        )
-        for r in responses:
-            for v in r.get("value", []):
-                display_name = v.get("displayName")
-                if display_name == item:
-                    item_id = v.get("id")
-                    break
-
-    if item_id is None:
-        raise ValueError(
-            f"{icons.red_dot} There's no item '{item}' of type '{type}' in the '{workspace_name}' workspace."
-        )
-
-    return item_id
-
-
-def resolve_item_name_and_id(
-    item: str | UUID, type: Optional[str] = None, workspace: Optional[str | UUID] = None
-) -> Tuple[str, UUID]:
-
-    workspace_id = resolve_workspace_id(workspace)
-    item_id = resolve_item_id(item=item, type=type, workspace=workspace_id)
-    item_name = (
-        _base_api(
-            request=f"/v1/workspaces/{workspace_id}/items/{item_id}", client="fabric_sp"
-        )
-        .json()
-        .get("displayName")
-    )
-
-    return item_name, item_id
-
-
-def resolve_item_name(item_id: UUID, workspace: Optional[str | UUID] = None) -> str:
-
-    workspace_id = resolve_workspace_id(workspace)
-    try:
-        item_name = (
-            _base_api(
-                request=f"/v1/workspaces/{workspace_id}/items/{item_id}",
-                client="fabric_sp",
-            )
-            .json()
-            .get("displayName")
-        )
-    except FabricHTTPException:
-        raise ValueError(
-            f"{icons.red_dot} The '{item_id}' item was not found in the '{workspace_id}' workspace."
-        )
-
-    return item_name
-
-
 def resolve_lakehouse_name_and_id(
     lakehouse: Optional[str | UUID] = None, workspace: Optional[str | UUID] = None
 ) -> Tuple[str, UUID]:
@@ -663,11 +590,13 @@ def save_as_delta_table(
     workspace: Optional[str | UUID] = None,
 ):
     """
-    Saves a pandas dataframe as a delta table in a Fabric lakehouse.
+    Saves a pandas or spark dataframe as a delta table in a Fabric lakehouse.
+
+    This function may be executed in either a PySpark or pure Python notebook. If executing in a pure Python notebook, the dataframe must be a pandas dataframe.
 
     Parameters
     ----------
-    dataframe : pandas.DataFrame
+    dataframe : pandas.DataFrame | spark.Dataframe
         The dataframe to be saved as a delta table.
     delta_table_name : str
         The name of the delta table.
@@ -686,19 +615,6 @@ def save_as_delta_table(
         or if no lakehouse attached, resolves to the workspace of the notebook.
     """
 
-    from pyspark.sql.types import (
-        StringType,
-        IntegerType,
-        FloatType,
-        DateType,
-        StructType,
-        StructField,
-        BooleanType,
-        LongType,
-        DoubleType,
-        TimestampType,
-    )
-
     (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
     (lakehouse_name, lakehouse_id) = resolve_lakehouse_name_and_id(
         lakehouse=lakehouse, workspace=workspace_id
@@ -717,52 +633,101 @@ def save_as_delta_table(
             f"{icons.red_dot} Invalid 'delta_table_name'. Delta tables in the lakehouse cannot have spaces in their names."
         )
 
-    spark = _create_spark_session()
+    import pyarrow as pa
+    from pyspark.sql.types import (
+        StringType,
+        IntegerType,
+        FloatType,
+        DateType,
+        StructType,
+        StructField,
+        BooleanType,
+        LongType,
+        DoubleType,
+        TimestampType,
+    )
 
-    type_mapping = {
-        "string": StringType(),
-        "str": StringType(),
-        "integer": IntegerType(),
-        "int": IntegerType(),
-        "float": FloatType(),
-        "date": DateType(),
-        "bool": BooleanType(),
-        "boolean": BooleanType(),
-        "long": LongType(),
-        "double": DoubleType(),
-        "timestamp": TimestampType(),
-    }
+    def get_type_mapping(pure_python):
+        common_mapping = {
+            "string": ("pa", pa.string(), StringType()),
+            "str": ("pa", pa.string(), StringType()),
+            "integer": ("pa", pa.int32(), IntegerType()),
+            "int": ("pa", pa.int32(), IntegerType()),
+            "float": ("pa", pa.float32(), FloatType()),
+            "double": ("pa", pa.float64(), DoubleType()),
+            "long": ("pa", pa.int64(), LongType()),
+            "bool": ("pa", pa.bool_(), BooleanType()),
+            "boolean": ("pa", pa.bool_(), BooleanType()),
+            "date": ("pa", pa.date32(), DateType()),
+            "timestamp": ("pa", pa.timestamp("us"), TimestampType()),
+        }
+        return {k: v[1] if pure_python else v[2] for k, v in common_mapping.items()}
 
-    if isinstance(dataframe, pd.DataFrame):
-        dataframe.columns = [col.replace(" ", "_") for col in dataframe.columns]
-        if schema is None:
-            spark_df = spark.createDataFrame(dataframe)
+    def build_schema(schema_dict, type_mapping, use_arrow=True):
+        if use_arrow:
+            fields = [
+                pa.field(name, type_mapping.get(dtype.lower()))
+                for name, dtype in schema_dict.items()
+            ]
+            return pa.schema(fields)
         else:
-            schema_map = StructType(
+            return StructType(
                 [
-                    StructField(column_name, type_mapping[data_type], True)
-                    for column_name, data_type in schema.items()
+                    StructField(name, type_mapping.get(dtype.lower()), True)
+                    for name, dtype in schema_dict.items()
                 ]
             )
-            spark_df = spark.createDataFrame(dataframe, schema_map)
+
+    # Main logic
+    schema_map = None
+    if schema is not None:
+        use_arrow = _pure_python_notebook()
+        type_mapping = get_type_mapping(use_arrow)
+        schema_map = build_schema(schema, type_mapping, use_arrow)
+
+    if isinstance(dataframe, pd.DataFrame):
+        dataframe.columns = [col.replace(" ", "_") for col in dataframe.columns]
+        if _pure_python_notebook():
+            spark_df = dataframe
+        else:
+            spark = _create_spark_session()
+            if schema is None:
+                spark_df = spark.createDataFrame(dataframe)
+            else:
+                spark_df = spark.createDataFrame(dataframe, schema_map)
     else:
         for col_name in dataframe.columns:
             new_name = col_name.replace(" ", "_")
             dataframe = dataframe.withColumnRenamed(col_name, new_name)
         spark_df = dataframe
 
-    filePath = create_abfss_path(
+    file_path = create_abfss_path(
         lakehouse_id=lakehouse_id,
         lakehouse_workspace_id=workspace_id,
         delta_table_name=delta_table_name,
     )
 
-    if merge_schema:
-        spark_df.write.mode(write_mode).format("delta").option(
-            "mergeSchema", "true"
-        ).save(filePath)
+    if _pure_python_notebook():
+        from deltalake import write_deltalake
+
+        write_args = {
+            "table_or_uri": file_path,
+            "data": spark_df,
+            "mode": write_mode,
+            "schema": schema_map,
+        }
+
+        if merge_schema:
+            write_args["schema_mode"] = "merge"
+
+        write_deltalake(**write_args)
     else:
-        spark_df.write.mode(write_mode).format("delta").save(filePath)
+        writer = spark_df.write.mode(write_mode).format("delta")
+        if merge_schema:
+            writer = writer.option("mergeSchema", "true")
+
+        writer.save(file_path)
+
     print(
         f"{icons.green_dot} The dataframe has been saved as the '{delta_table_name}' table in the '{lakehouse_name}' lakehouse within the '{workspace_name}' workspace."
     )
@@ -898,6 +863,87 @@ def resolve_workspace_name_and_id(
     return workspace_name, workspace_id
 
 
+def resolve_item_id(
+    item: str | UUID, type: Optional[str] = None, workspace: Optional[str | UUID] = None
+) -> UUID:
+
+    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
+    item_id = None
+
+    if _is_valid_uuid(item):
+        # Check (optional)
+        item_id = item
+        try:
+            _base_api(
+                request=f"/v1/workspaces/{workspace_id}/items/{item_id}",
+                client="fabric_sp",
+            )
+        except FabricHTTPException:
+            raise ValueError(
+                f"{icons.red_dot} The '{item_id}' item was not found in the '{workspace_name}' workspace."
+            )
+    else:
+        if type is None:
+            raise ValueError(
+                f"{icons.red_dot} The 'type' parameter is required if specifying an item name."
+            )
+        responses = _base_api(
+            request=f"/v1/workspaces/{workspace_id}/items?type={type}",
+            client="fabric_sp",
+            uses_pagination=True,
+        )
+        for r in responses:
+            for v in r.get("value", []):
+                display_name = v.get("displayName")
+                if display_name == item:
+                    item_id = v.get("id")
+                    break
+
+    if item_id is None:
+        raise ValueError(
+            f"{icons.red_dot} There's no item '{item}' of type '{type}' in the '{workspace_name}' workspace."
+        )
+
+    return item_id
+
+
+def resolve_item_name_and_id(
+    item: str | UUID, type: Optional[str] = None, workspace: Optional[str | UUID] = None
+) -> Tuple[str, UUID]:
+
+    workspace_id = resolve_workspace_id(workspace)
+    item_id = resolve_item_id(item=item, type=type, workspace=workspace_id)
+    item_name = (
+        _base_api(
+            request=f"/v1/workspaces/{workspace_id}/items/{item_id}", client="fabric_sp"
+        )
+        .json()
+        .get("displayName")
+    )
+
+    return item_name, item_id
+
+
+def resolve_item_name(item_id: UUID, workspace: Optional[str | UUID] = None) -> str:
+
+    workspace_id = resolve_workspace_id(workspace)
+    try:
+        item_name = (
+            _base_api(
+                request=f"/v1/workspaces/{workspace_id}/items/{item_id}",
+                client="fabric_sp",
+            )
+            .json()
+            .get("displayName")
+        )
+    except FabricHTTPException:
+        raise ValueError(
+            f"{icons.red_dot} The '{item_id}' item was not found in the '{workspace_id}' workspace."
+        )
+
+    return item_name
+
+
 def _extract_json(dataframe: pd.DataFrame) -> dict:
 
     payload = dataframe["payload"].iloc[0]
@@ -1497,32 +1543,108 @@ def generate_guid():
 
 def _get_column_aggregate(
     table_name: str,
-    column_name: str = "RunId",
+    column_name: str | List[str] = "RunId",
     lakehouse: Optional[str | UUID] = None,
     workspace: Optional[str | UUID] = None,
     function: str = "max",
     default_value: int = 0,
-) -> int:
+    schema_name: Optional[str] = None,
+) -> int | Dict[str, int]:
 
-    from pyspark.sql.functions import approx_count_distinct
-    from pyspark.sql import functions as F
+    workspace_id = resolve_workspace_id(workspace)
+    lakehouse_id = resolve_lakehouse_id(lakehouse, workspace_id)
+    path = create_abfss_path(lakehouse_id, workspace_id, table_name, schema_name)
+    df = _read_delta_table(path)
 
-    function = function.upper()
-    (workspace_name, workspace_id) = resolve_workspace_name_and_id(workspace)
-    lakehouse_id = resolve_lakehouse_id(lakehouse, workspace)
-    path = create_abfss_path(lakehouse_id, workspace_id, table_name)
+    function = function.lower()
 
-    spark = _create_spark_session()
-    df = spark.read.format("delta").load(path)
+    if isinstance(column_name, str):
+        column_name = [column_name]
+
+    if _pure_python_notebook():
+        import polars as pl
+
+        if not isinstance(df, pd.DataFrame):
+            df.to_pandas()
+
+        df = pl.from_pandas(df)
 
-    if function in {"COUNTDISTINCT", "DISTINCTCOUNT"}:
-        result = df.select(F.count_distinct(F.col(column_name)))
-    elif "APPROX" in function:
-        result = df.select(approx_count_distinct(column_name))
+        def get_expr(col):
+            col_dtype = df.schema[col]
+
+            if "approx" in function:
+                return pl.col(col).unique().count().alias(col)
+            elif "distinct" in function:
+                if col_dtype == pl.Decimal:
+                    return pl.col(col).cast(pl.Float64).n_unique().alias(col)
+                else:
+                    return pl.col(col).n_unique().alias(col)
+            elif function == "sum":
+                return pl.col(col).sum().alias(col)
+            elif function == "min":
+                return pl.col(col).min().alias(col)
+            elif function == "max":
+                return pl.col(col).max().alias(col)
+            elif function == "count":
+                return pl.col(col).count().alias(col)
+            elif function in {"avg", "mean"}:
+                return pl.col(col).mean().alias(col)
+            else:
+                raise ValueError(f"Unsupported function: {function}")
+
+        exprs = [get_expr(col) for col in column_name]
+        aggs = df.select(exprs).to_dict(as_series=False)
+
+        if len(column_name) == 1:
+            result = aggs[column_name[0]][0] or default_value
+        else:
+            result = {col: aggs[col][0] for col in column_name}
     else:
-        result = df.selectExpr(f"{function}({column_name})")
+        from pyspark.sql.functions import (
+            count,
+            sum,
+            min,
+            max,
+            avg,
+            approx_count_distinct,
+            countDistinct,
+        )
+
+        result = None
+        if "approx" in function:
+            spark_func = approx_count_distinct
+        elif "distinct" in function:
+            spark_func = countDistinct
+        elif function == "count":
+            spark_func = count
+        elif function == "sum":
+            spark_func = sum
+        elif function == "min":
+            spark_func = min
+        elif function == "max":
+            spark_func = max
+        elif function == "avg":
+            spark_func = avg
+        else:
+            raise ValueError(f"Unsupported function: {function}")
 
-    return result.collect()[0][0] or default_value
+        agg_exprs = []
+        for col in column_name:
+            agg_exprs.append(spark_func(col).alias(col))
+
+        aggs = df.agg(*agg_exprs).collect()[0]
+        if len(column_name) == 1:
+            result = aggs[0] or default_value
+        else:
+            result = {col: aggs[col] for col in column_name}
+
+    return result
+
+
+def _create_spark_dataframe(df: pd.DataFrame):
+
+    spark = _create_spark_session()
+    return spark.createDataFrame(df)
 
 
 def _make_list_unique(my_list):
@@ -1617,6 +1739,9 @@ def _process_and_display_chart(df, title, widget):
     df["Start"] = df["Start"] - Offset
     df["End"] = df["End"] - Offset
 
+    unique_objects = df["Object Name"].nunique()
+    height = min(max(400, unique_objects * 30), 1000)
+
     # Vega-Lite spec for Gantt chart
     spec = (
         """{
@@ -1626,7 +1751,9 @@ def _process_and_display_chart(df, title, widget):
         + df.to_json(orient="records")
         + """ },
         "width": 700,
-        "height": 400,
+        "height": """
+        + str(height)
+        + """,
         "mark": "bar",
         "encoding": {
             "y": {
@@ -1687,6 +1814,7 @@ def _convert_data_type(input_data_type: str) -> str:
         "double": "Double",
         "float": "Double",
         "binary": "Boolean",
+        "long": "Int64",
     }
 
     if "decimal" in input_data_type:
@@ -1842,6 +1970,18 @@ def _update_dataframe_datatypes(dataframe: pd.DataFrame, column_map: dict):
                 dataframe[column] = dataframe[column].fillna(0).astype(int)
             elif data_type in ["str", "string"]:
                 dataframe[column] = dataframe[column].astype(str)
+            # Avoid having empty lists or lists with a value of None.
+            elif data_type in ["list"]:
+                dataframe[column] = dataframe[column].apply(
+                    lambda x: (
+                        None
+                        if (type(x) == list and len(x) == 1 and x[0] == None)
+                        or (type(x) == list and len(x) == 0)
+                        else x
+                    )
+                )
+            elif data_type in ["dict"]:
+                dataframe[column] = dataframe[column]
             else:
                 raise NotImplementedError
 
@@ -1878,18 +2018,58 @@ def _create_spark_session():
     return SparkSession.builder.getOrCreate()
 
 
-def _read_delta_table(path: str):
+def _get_delta_table(path: str) -> str:
+
+    from delta import DeltaTable
 
     spark = _create_spark_session()
 
-    return spark.read.format("delta").load(path)
+    return DeltaTable.forPath(spark, path)
 
 
-def _delta_table_row_count(table_name: str) -> int:
+def _read_delta_table(path: str, to_pandas: bool = True, to_df: bool = False):
 
-    spark = _create_spark_session()
+    if _pure_python_notebook():
+        from deltalake import DeltaTable
+
+        df = DeltaTable(table_uri=path)
+        if to_pandas:
+            df = df.to_pandas()
+    else:
+        spark = _create_spark_session()
+        df = spark.read.format("delta").load(path)
+        if to_df:
+            df = df.toDF()
+
+    return df
+
+
+def _read_delta_table_history(path) -> pd.DataFrame:
+
+    if _pure_python_notebook():
+        from deltalake import DeltaTable
+
+        df = pd.DataFrame(DeltaTable(table_uri=path).history())
+    else:
+        from delta import DeltaTable
+
+        spark = _create_spark_session()
+        delta_table = DeltaTable.forPath(spark, path)
+        df = delta_table.history().toPandas()
+
+    return df
+
+
+def _delta_table_row_count(path: str) -> int:
 
-    return spark.table(table_name).count()
+    if _pure_python_notebook():
+        from deltalake import DeltaTable
+
+        dt = DeltaTable(path)
+        arrow_table = dt.to_pyarrow_table()
+        return arrow_table.num_rows
+    else:
+        return _read_delta_table(path).count()
 
 
 def _run_spark_sql_query(query):
@@ -2070,3 +2250,23 @@ def _xml_to_dict(element):
             element.text.strip() if element.text and element.text.strip() else None
         )
     return data
+
+
+def file_exists(file_path: str) -> bool:
+    """
+    Check if a file exists in the given path.
+
+    Parameters
+    ----------
+    file_path : str
+        The path to the file.
+
+    Returns
+    -------
+    bool
+        True if the file exists, False otherwise.
+    """
+
+    import notebookutils
+
+    return len(notebookutils.fs.ls(file_path)) > 0

sempy_labs/_kql_databases.py

@@ -6,6 +6,8 @@ from sempy_labs._helper_functions import (
     _create_dataframe,
     delete_item,
     create_item,
+    resolve_item_id,
+    resolve_workspace_id,
 )
 from uuid import UUID
 import sempy_labs._icons as icons
@@ -121,3 +123,19 @@ def delete_kql_database(
         )
 
     delete_item(item=kql_database, type="KQLDatabase", workspace=workspace)
+
+
+def _resolve_cluster_uri(
+    kql_database: str | UUID, workspace: Optional[str | UUID] = None
+) -> str:
+
+    workspace_id = resolve_workspace_id(workspace=workspace)
+    item_id = resolve_item_id(
+        item=kql_database, type="KQLDatabase", workspace=workspace
+    )
+    response = _base_api(
+        request=f"/v1/workspaces/{workspace_id}/kqlDatabases/{item_id}",
+        client="fabric_sp",
+    )
+
+    return response.json().get("properties", {}).get("queryServiceUri")