lakehouse_engine.utils.schema_utils

  1"""Utilities to facilitate dataframe schema management."""
  2
  3from logging import Logger
  4from typing import Any, List, Optional
  5
  6from pyspark.sql.functions import col
  7from pyspark.sql.types import StructType
  8
  9from lakehouse_engine.core.definitions import InputSpec
 10from lakehouse_engine.core.exec_env import ExecEnv
 11from lakehouse_engine.utils.logging_handler import LoggingHandler
 12from lakehouse_engine.utils.storage.file_storage_functions import FileStorageFunctions
 13
 14
 15class SchemaUtils(object):
 16    """Schema utils that help retrieve and manage schemas of dataframes."""
 17
 18    _logger: Logger = LoggingHandler(__name__).get_logger()
 19
 20    @staticmethod
 21    def from_file(file_path: str, disable_dbfs_retry: bool = False) -> StructType:
 22        """Get a spark schema from a file (spark StructType json file) in a file system.
 23
 24        Args:
 25            file_path: path of the file in a file system. [Check here](
 26                https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/types/StructType.html).
 27            disable_dbfs_retry: optional flag to disable file storage dbfs.
 28
 29        Returns:
 30            Spark schema struct type.
 31        """
 32        return StructType.fromJson(
 33            FileStorageFunctions.read_json(file_path, disable_dbfs_retry)
 34        )
 35
 36    @staticmethod
 37    def from_file_to_dict(file_path: str, disable_dbfs_retry: bool = False) -> Any:
 38        """Get a dict with the spark schema from a file in a file system.
 39
 40        Args:
 41            file_path: path of the file in a file system. [Check here](
 42                https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/types/StructType.html).
 43            disable_dbfs_retry: optional flag to disable file storage dbfs.
 44
 45        Returns:
 46             Spark schema in a dict.
 47        """
 48        return FileStorageFunctions.read_json(file_path, disable_dbfs_retry)
 49
 50    @staticmethod
 51    def from_dict(struct_type: dict) -> StructType:
 52        """Get a spark schema from a dict.
 53
 54        Args:
 55            struct_type: dict containing a spark schema structure. [Check here](
 56                https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/types/StructType.html).
 57
 58        Returns:
 59             Spark schema struct type.
 60        """
 61        return StructType.fromJson(struct_type)
 62
 63    @staticmethod
 64    def from_table_schema(table: str) -> StructType:
 65        """Get a spark schema from a table.
 66
 67        Args:
 68            table: table name from which to inherit the schema.
 69
 70        Returns:
 71            Spark schema struct type.
 72        """
 73        return ExecEnv.SESSION.read.table(table).schema
 74
 75    @classmethod
 76    def from_input_spec(cls, input_spec: InputSpec) -> Optional[StructType]:
 77        """Get a spark schema from an input specification.
 78
 79        This covers scenarios where the schema is provided as part of the input
 80        specification of the algorithm. Schema can come from the table specified in the
 81        input specification (enforce_schema_from_table) or by the dict with the spark
 82        schema provided there also.
 83
 84        Args:
 85            input_spec: input specification.
 86
 87        Returns:
 88            spark schema struct type.
 89        """
 90        if input_spec.enforce_schema_from_table:
 91            cls._logger.info(
 92                f"Reading schema from table: {input_spec.enforce_schema_from_table}"
 93            )
 94            return SchemaUtils.from_table_schema(input_spec.enforce_schema_from_table)
 95        elif input_spec.schema_path:
 96            cls._logger.info(f"Reading schema from file: {input_spec.schema_path}")
 97            return SchemaUtils.from_file(
 98                input_spec.schema_path, input_spec.disable_dbfs_retry
 99            )
100        elif input_spec.schema:
101            cls._logger.info(
102                f"Reading schema from configuration file: {input_spec.schema}"
103            )
104            return SchemaUtils.from_dict(input_spec.schema)
105        else:
106            cls._logger.info("No schema was provided... skipping enforce schema")
107            return None
108
109    @staticmethod
110    def _get_prefix_alias(num_chars: int, prefix: str, shorten_names: bool) -> str:
111        """Get prefix alias for a field."""
112        return (
113            f"""{'_'.join(
114                [item[:num_chars] for item in prefix.split('.')]
115            )}_"""
116            if shorten_names
117            else f"{prefix}_".replace(".", "_")
118        )
119
120    @staticmethod
121    def schema_flattener(
122        schema: StructType,
123        prefix: str = None,
124        level: int = 1,
125        max_level: int = None,
126        shorten_names: bool = False,
127        alias: bool = True,
128        num_chars: int = 7,
129        ignore_cols: List = None,
130    ) -> List:
131        """Recursive method to flatten the schema of the dataframe.
132
133        Args:
134            schema: schema to be flattened.
135            prefix: prefix of the struct to get the value for. Only relevant
136                for being used in the internal recursive logic.
137            level: level of the depth in the schema being flattened. Only relevant
138                for being used in the internal recursive logic.
139            max_level: level until which you want to flatten the schema. Default: None.
140            shorten_names: whether to shorten the names of the prefixes of the fields
141                being flattened or not. Default: False.
142            alias: whether to define alias for the columns being flattened or
143                not. Default: True.
144            num_chars: number of characters to consider when shortening the names of
145                the fields. Default: 7.
146            ignore_cols: columns which you don't want to flatten. Default: None.
147
148        Returns:
149            A function to be called in .transform() spark function.
150        """
151        cols = []
152        ignore_cols = ignore_cols if ignore_cols else []
153        for field in schema.fields:
154            name = prefix + "." + field.name if prefix else field.name
155            field_type = field.dataType
156
157            if (
158                isinstance(field_type, StructType)
159                and name not in ignore_cols
160                and (max_level is None or level <= max_level)
161            ):
162                cols += SchemaUtils.schema_flattener(
163                    schema=field_type,
164                    prefix=name,
165                    level=level + 1,
166                    max_level=max_level,
167                    shorten_names=shorten_names,
168                    alias=alias,
169                    num_chars=num_chars,
170                    ignore_cols=ignore_cols,
171                )
172            else:
173                if alias and prefix:
174                    prefix_alias = SchemaUtils._get_prefix_alias(
175                        num_chars, prefix, shorten_names
176                    )
177                    cols.append(col(name).alias(f"{prefix_alias}{field.name}"))
178                else:
179                    cols.append(col(name))
180        return cols