fabricks 3.0.11__py3-none-any.whl

fabricks/utils/path.py

@@ -0,0 +1,242 @@
+import os
+from pathlib import Path as PathlibPath
+from typing import List, Optional, Union
+
+from pyspark.sql.dataframe import DataFrame
+
+from fabricks.utils.spark import spark
+
+
+class Path:
+    def __init__(self, path: Union[str, PathlibPath], assume_git: bool = False):
+        self.assume_git = assume_git
+
+        if isinstance(path, PathlibPath):
+            path = path.as_posix()
+
+        new_path = str(path)
+        if new_path.startswith("abfss:/") and not new_path.startswith("abfss://"):
+            new_path = new_path.replace("abfss:/", "abfss://")  # // is replaced by / by pathlibpath
+
+        self.path: str = new_path
+
+    @classmethod
+    def from_uri(
+        cls,
+        uri: str,
+        regex: Optional[dict[str, str]] = None,
+        assume_git: Optional[bool] = False,
+    ):
+        uri = uri.strip()
+        if assume_git is None:
+            assume_git = False
+        if regex:
+            import re
+
+            for key, value in regex.items():
+                uri = re.sub(rf"{key}", value, uri)
+
+        return cls(uri, assume_git=assume_git)
+
+    @property
+    def pathlibpath(self) -> PathlibPath:
+        return PathlibPath(self.string)
+
+    @property
+    def string(self) -> str:
+        return self.path
+
+    def get_container(self) -> str:
+        import re
+
+        assert self.string.startswith("abfss://")
+
+        r = re.compile(r"(?<=abfss:\/\/)(.+?)(?=@)")
+        m = re.findall(r, self.string)[0]
+        return m
+
+    def get_storage_account(self) -> str:
+        import re
+
+        assert self.string.startswith("abfss://")
+
+        r = re.compile(r"(?<=@)(.+?)(?=\.)")
+        m = re.findall(r, self.string)[0]
+        return m
+
+    def get_file_name(self) -> str:
+        return self.pathlibpath.name
+
+    def get_file_system(self) -> str:
+        import re
+
+        assert self.string.startswith("abfss://")
+
+        r = re.compile(r"(?<=\.)(.+)(?=\/)")
+        m = re.findall(r, self.string)[0]
+        return m
+
+    def get_dbfs_mnt_path(self) -> str:
+        mount_point = self.pathlibpath.parts[1].split(".")[0].split("@")[0]
+        rest = self.pathlibpath.parts[2:]
+
+        return str(os.path.join("/dbfs/mnt", mount_point, "/".join(rest)))
+
+    def get_notebook_path(self) -> str:
+        path = self.path.replace("Workspace/", "")
+        if path.endswith(".ipynb"):
+            path = path.replace(".ipynb", "")
+        if path.endswith(".py"):
+            path = path.replace(".py", "")
+        return path
+
+    def get_sql(self) -> str:
+        p = self.string
+        if not p.endswith(".sql"):
+            p += ".sql"
+
+        with open(p, "r") as f:
+            sql = f.read()
+
+        return sql
+
+    def is_sql(self) -> bool:
+        return self.string.endswith(".sql")
+
+    def exists(self) -> bool:
+        try:
+            if self.assume_git:
+                return self.pathlibpath.exists()
+
+            else:
+                from fabricks.utils.spark import dbutils
+
+                assert dbutils is not None, "dbutils not found"
+
+                dbutils.fs.ls(self.string)
+                return True
+
+        except Exception:
+            return False
+
+    def joinpath(self, *other):
+        new_path = self.pathlibpath.joinpath(*other)
+        return Path(path=new_path, assume_git=self.assume_git)
+
+    def append(self, other: str):
+        new_path = self.string + other
+        return Path(path=new_path, assume_git=self.assume_git)
+
+    def parent(self, *other):
+        new_path = self.pathlibpath.parent
+        return Path(path=new_path, assume_git=self.assume_git)
+
+    def get_file_info(self) -> DataFrame:
+        assert not self.assume_git
+
+        rows = self._yield_file_info(self.string)
+        df = spark.createDataFrame(
+            rows,
+            schema=["path", "name", "size", "modification_time"],
+        )
+        return df
+
+    def walk(
+        self,
+        depth: Optional[int] = None,
+        convert: Optional[bool] = False,
+        file_format: Optional[str] = None,
+    ) -> List:
+        out = []
+        if self.exists():
+            if self.pathlibpath.is_file():
+                out = [self.string]
+            elif depth:
+                assert not self.assume_git
+                out = self._list_fs(depth)
+            else:
+                if self.assume_git:
+                    out = list(self._yield_git(self.string))
+                else:
+                    out = list(self._yield_fs(self.string))
+
+        if file_format:
+            out = [o for o in out if o.endswith(".sql")]
+        if convert:
+            out = [Path(o) for o in out]
+        return out
+
+    def _list_fs(self, depth: int) -> List:
+        from databricks.sdk.runtime import dbutils
+
+        paths = dbutils.fs.ls(self.string)
+
+        if depth == 1:
+            children = paths
+        else:
+            i = 1
+            children = []
+            while True:
+                if i == depth:
+                    break
+                else:
+                    children = []
+
+                for path in paths:
+                    children += dbutils.fs.ls(path.path)
+
+                paths = children
+                i += 1
+
+        return [c.path for c in children]
+
+    def _yield_file_info(self, path: str):
+        from databricks.sdk.runtime import dbutils
+
+        for child in dbutils.fs.ls(path):
+            if child.isDir():  # type: ignore
+                yield from self._yield_file_info(child.path)
+            else:
+                yield dbutils.fs.ls(child.path)[0]
+
+    def _yield_fs(self, path: str):
+        from databricks.sdk.runtime import dbutils
+
+        for child in dbutils.fs.ls(path):
+            if child.isDir():  # type: ignore
+                yield from self._yield_fs(child.path)
+            else:
+                yield str(child.path)
+
+    def _yield_git(self, path: Union[str, PathlibPath]):
+        if isinstance(path, str):
+            path = PathlibPath(path)
+
+        for child in path.glob(r"*"):
+            if child.is_dir():
+                yield from self._yield_git(child)
+            else:
+                yield str(child)
+
+    def rm(self):
+        from databricks.sdk.runtime import dbutils
+
+        if self.exists():
+            list(self._rm(self.string))
+            dbutils.fs.rm(self.string, recurse=True)
+
+    def _rm(self, path: str):
+        from databricks.sdk.runtime import dbutils
+
+        try:
+            for child in dbutils.fs.ls(path):
+                if child.isDir():  # type: ignore
+                    yield from self._rm(child.path)
+                else:
+                    yield dbutils.fs.rm(child.path, recurse=True)
+
+        except Exception:
+            return False
+
+    def __str__(self) -> str:
+        return self.string

fabricks/utils/pip.py

@@ -0,0 +1,61 @@
+import subprocess
+from typing import List, Optional, Union
+
+from fabricks.utils.path import Path
+
+
+def pip_package(
+    package: Union[str, List[str]],
+    whl_path: Optional[Path] = None,
+    tgt_path: Optional[Path] = None,
+):
+    if isinstance(package, str):
+        package = [package]
+
+    args = ["pip", "install"]
+
+    if whl_path:
+        w = whl_path.get_dbfs_mnt_path()
+        args += ["--no-index", f"--find-links={w}"]
+
+    if tgt_path:
+        t = tgt_path.get_dbfs_mnt_path()
+        args += ["--target", t]
+
+    for p in package:
+        out = subprocess.run(args + [p], capture_output=True)
+        if out.returncode == 1:
+            raise ValueError(p, out.stderr)
+
+
+def pip_requirements(
+    requirements_path: Path,
+    whl_path: Optional[Path] = None,
+    tgt_path: Optional[Path] = None,
+):
+    r = requirements_path.string
+
+    args = ["pip", "install"]
+
+    if whl_path:
+        w = whl_path.get_dbfs_mnt_path()
+        args += ["--no-index", f"--find-links={w}"]
+
+    if tgt_path:
+        t = tgt_path.get_dbfs_mnt_path()
+        args += ["--target", t]
+
+    out = subprocess.run(args + ["-r", r], capture_output=True)
+    if out.returncode == 1:
+        raise ValueError(r, out.stderr)
+
+
+def pip_wheel(requirement_path: Path, whl_path: Path):
+    import subprocess
+
+    r = requirement_path.string
+    w = whl_path.get_dbfs_mnt_path()
+
+    out = subprocess.run(["pip", "wheel", "--wheel-dir", w, "-r", r], capture_output=True)
+    if out.returncode == 1:
+        raise ValueError(r, out.stderr)

fabricks/utils/pydantic.py

@@ -0,0 +1,94 @@
+from typing import List, Literal, Type, TypeVar, Union, get_args, get_origin
+
+import yaml
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic import parse_obj_as
+from pyspark.sql import DataFrame
+from pyspark.sql.types import (
+    ArrayType,
+    BooleanType,
+    DoubleType,
+    LongType,
+    MapType,
+    Row,
+    StringType,
+    StructField,
+    StructType,
+)
+
+from fabricks.context import SPARK
+
+types_ = {
+    str: StringType(),
+    bool: BooleanType(),
+    float: DoubleType(),
+    int: LongType(),
+    dict: MapType(StringType(), StringType()),
+}
+T = TypeVar("T")
+
+
+def _to_spark_type(type_):
+    if type_ in types_:
+        return types_[type_]
+
+    origin = get_origin(type_)
+    args = get_args(type_)
+    if origin is Literal:
+        return StringType()
+    if origin is list:
+        return ArrayType(_to_spark_type(args[0]))
+    if origin is dict:
+        return MapType(
+            _to_spark_type(args[0]),
+            _to_spark_type(args[1]),
+        )
+
+    if issubclass(type_, PydanticBaseModel):
+        return _schema_pyspark(type_)
+
+    raise ValueError(type_)
+
+
+def _schema_pyspark(model):
+    fields = []
+    for field in model.__fields__.values():
+        type_ = field.outer_type_
+        spark_type_ = _to_spark_type(type_)
+        f = StructField(
+            name=field.name,
+            dataType=spark_type_,  # type: ignore
+            nullable=not field.required,
+        )
+        fields.append(f)
+    return StructType(fields)
+
+
+class FBaseModel(PydanticBaseModel):
+    @classmethod
+    def from_yaml(cls: Type[T], path: str) -> Union[T, List[T]]:
+        with open(path, encoding="utf-8") as f:
+            y = yaml.safe_load(f)
+            if isinstance(y, List):
+                return parse_obj_as(List[cls], y)
+            else:
+                return parse_obj_as(cls, y)
+
+    @classmethod
+    def from_row(cls: Type[T], row: Row) -> T:
+        return parse_obj_as(cls, row.asDict(True))
+
+    @classmethod
+    def from_dataframe(cls: Type[T], df: DataFrame) -> List[T]:
+        return [parse_obj_as(cls, row.asDict(True)) for row in df.collect()]
+
+    def schema_pyspark(self):
+        return _schema_pyspark(self)
+
+    @staticmethod
+    def get_dataframe(data: Union[T, List[T]]) -> DataFrame:
+        if isinstance(data, List):
+            df = SPARK.createDataFrame([d.dict() for d in data], data[0].schema_pyspark())  # type: ignore
+        else:
+            df = SPARK.createDataFrame([data.dict()], data.schema_pyspark())  # type: ignore
+        return df

fabricks/utils/read/__init__.py

@@ -0,0 +1,11 @@
+from fabricks.utils.read.read import read, read_batch, read_stream
+from fabricks.utils.read.read_excel import read_excel
+from fabricks.utils.read.read_yaml import read_yaml
+
+__all__ = [
+    "read_batch",
+    "read_excel",
+    "read_stream",
+    "read_yaml",
+    "read",
+]

fabricks/utils/read/_types.py

@@ -0,0 +1,3 @@
+from typing import Literal
+
+AllowedIOModes = Literal["overwrite", "append"]

fabricks/utils/read/read.py

@@ -0,0 +1,305 @@
+from typing import List, Optional, Union, overload
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import StructType
+
+from fabricks.context import SPARK
+from fabricks.utils.path import Path
+
+
+@overload
+def read_stream(
+    src: Union[Path, str],
+    file_format: str,
+    *,
+    schema: StructType,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read_stream(
+    src: Union[Path, str],
+    file_format: str,
+    schema_path: Union[Path, str],
+    *,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read_stream(
+    src: Union[Path, str],
+    file_format: str,
+    *,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+def read_stream(
+    src: Union[Path, str],
+    file_format: str,
+    schema_path: Optional[Union[Path, str]] = None,
+    hints: Optional[Union[str, List[str]]] = None,
+    schema: Optional[StructType] = None,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame:
+    return _read_stream(
+        src=src,
+        file_format=file_format,
+        schema_path=schema_path,
+        hints=hints,
+        schema=schema,
+        options=options,
+        spark=spark,
+    )
+
+
+def _read_stream(
+    src: Union[Path, str],
+    file_format: str,
+    schema_path: Optional[Union[Path, str]] = None,
+    hints: Optional[Union[str, List[str]]] = None,
+    schema: Optional[StructType] = None,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame:
+    if spark is None:
+        spark = SPARK
+    assert spark is not None
+
+    if file_format == "table":
+        assert isinstance(src, str)
+        return spark.readStream.table(src)
+    else:
+        file_format = "binaryFile" if file_format == "pdf" else file_format
+        if isinstance(src, str):
+            src = Path(src)
+        if file_format == "delta":
+            reader = spark.readStream.format("delta")
+        else:
+            reader = spark.readStream.format("cloudFiles")
+            reader.option("cloudFiles.format", file_format)
+            if schema:
+                reader.schema(schema)
+            else:
+                assert schema_path
+                if isinstance(schema_path, str):
+                    schema_path = Path(schema_path)
+                reader.option("cloudFiles.inferColumnTypes", "true")
+                reader.option("cloudFiles.useIncrementalListing", "true")
+                reader.option("cloudFiles.schemaEvolutionMode", "addNewColumns")
+                reader.option("cloudFiles.schemaLocation", schema_path.string)
+                if hints:
+                    if isinstance(hints, str):
+                        hints = [hints]
+                    reader.option("cloudFiles.schemaHints", f"{' ,'.join(hints)}")
+
+        # default options
+        reader.option("recursiveFileLookup", "true")
+        reader.option("skipChangeCommits", "true")
+        reader.option("ignoreDeletes", "true")
+        if file_format == "csv":
+            reader.option("header", "true")
+        # custom / override options
+        if options:
+            for key, value in options.items():
+                reader.option(key, value)
+
+        df = reader.load(src.string)
+        df = df.withColumnRenamed("_rescued_data", "__rescued_data")
+        return df
+
+
+@overload
+def read_batch(
+    src: Union[Path, str],
+    file_format: str,
+    schema: StructType,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read_batch(
+    src: Union[Path, str],
+    file_format: str,
+    *,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+def read_batch(
+    src: Union[Path, str],
+    file_format: str,
+    schema: Optional[StructType] = None,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame:
+    return _read_batch(
+        src=src,
+        file_format=file_format,
+        schema=schema,
+        options=options,
+        spark=spark,
+    )
+
+
+def _read_batch(
+    src: Union[Path, str],
+    file_format: str,
+    schema: Optional[StructType] = None,
+    options: Optional[dict[str, str]] = None,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame:
+    if spark is None:
+        spark = SPARK
+    assert spark is not None
+
+    if file_format == "table":
+        assert isinstance(src, str)
+        return spark.read.table(src)
+    else:
+        path_glob_filter = file_format
+        file_format = "binaryFile" if file_format == "pdf" else file_format
+        if isinstance(src, str):
+            src = Path(src)
+        reader = spark.read.format(file_format)
+        reader = reader.option("pathGlobFilter", f"*.{path_glob_filter}")
+        if schema:
+            reader = reader.schema(schema)
+        # default options
+        reader = reader.option("recursiveFileLookup", "True")
+        if file_format == "parquet":
+            reader = reader.option("mergeSchema", "true")
+        if file_format == "csv":
+            reader = reader.option("header", "true")
+        # custom / override options
+        if options:
+            for key, value in options.items():
+                reader = reader.option(key, value)
+        return reader.load(src.string)
+
+
+@overload
+def read(
+    stream: bool,
+    table: str,
+    *,
+    metadata: Optional[bool] = False,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read(
+    stream: bool,
+    *,
+    path: Union[Path, str],
+    file_format: str = "delta",
+    metadata: Optional[bool] = False,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read(
+    stream: bool,
+    *,
+    path: Union[Path, str],
+    file_format: str,
+    schema: StructType,
+    options: Optional[dict[str, str]] = None,
+    metadata: Optional[bool] = True,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+@overload
+def read(
+    stream: bool,
+    *,
+    path: Union[Path, str],
+    file_format: str,
+    schema_path: Union[Path, str],
+    options: Optional[dict[str, str]] = None,
+    metadata: Optional[bool] = True,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame: ...
+
+
+def read(
+    stream: bool,
+    table: Optional[str] = None,
+    path: Optional[Union[Path, str]] = None,
+    file_format: Optional[str] = None,
+    schema_path: Optional[Union[Path, str]] = None,
+    schema: Optional[StructType] = None,
+    hints: Optional[Union[str, List[str]]] = None,
+    options: Optional[dict[str, str]] = None,
+    metadata: Optional[bool] = True,
+    spark: Optional[SparkSession] = None,
+) -> DataFrame:
+    if spark is None:
+        spark = SPARK
+    assert spark is not None
+
+    if table is not None:
+        file_format = "table"
+        src = table
+    else:
+        assert path
+        assert file_format
+        src = path
+
+    if stream:
+        df = _read_stream(
+            src=src,
+            file_format=file_format,
+            schema_path=schema_path,
+            hints=hints,
+            schema=schema,
+            options=options,
+            spark=spark,
+        )
+    else:
+        df = _read_batch(
+            src=src,
+            file_format=file_format,
+            schema=schema,
+            options=options,
+            spark=spark,
+        )
+
+    if metadata:
+        if stream and file_format == "delta":
+            df = df.selectExpr(
+                "*",
+                """
+                struct(
+                    cast(null as string) as file_path,
+                    cast(null as string) as file_name,
+                    cast(null as string) as file_size,            
+                    cast(null as string) as file_modification_time
+                    ) as __metadata
+                """,
+            )
+        else:
+            df = df.selectExpr(
+                "*",
+                """
+                struct(
+                    _metadata.file_path as file_path,
+                    _metadata.file_name as file_name,
+                    _metadata.file_size as file_size,            
+                    _metadata.file_modification_time as file_modification_time
+                    ) as __metadata
+                """,
+            )
+    return df