dasl-client 1.0.23__py3-none-any.whl → 1.0.24__py3-none-any.whl

dasl_client/preset_development/preview_engine.py

@@ -7,6 +7,12 @@ from dasl_client.preset_development.stage import *
 from dasl_client.preset_development.errors import *
 import yaml
 from IPython import get_ipython
+from itertools import count
+
+
+@udf(StringType())
+def constant_udf(*args):
+    return "<sortable_random_id>"
 
 
 class PreviewEngine:
@@ -54,8 +60,9 @@ class PreviewEngine:
                 self._preset.get("silver", None), self._pretransform_name
             )
 
+        self._pre_bronze = None
         self._bronze = None
-        self._pre = None
+        self._pre_silver = None
         self._silver = []
         self._gold = []
         self._result_df_map: Tuple[
@@ -124,13 +131,32 @@ class PreviewEngine:
 
     def _compile_stages(self) -> None:
         """
-        Creates Stage objects, setting pretransform to None if not provided.
+        Creates Stage objects, setting silver pretransform to None if not provided.
         """
+        pre_bronze_field_counter = count()
+        pre_bronze_name_counter = count()
+        pre_bronze_expr_groups = self._preset.get("bronze", {}).get("preTransform", [])
+        if pre_bronze_expr_groups:
+            tables = [
+                {
+                    "name": f"Index {next(pre_bronze_name_counter)}",
+                    "fields": [
+                        {"name": str(next(pre_bronze_field_counter)), "expr": expr}
+                        for expr in expr_group
+                    ],
+                }
+                for expr_group in pre_bronze_expr_groups
+            ]
+            for table in tables:
+                self._pre_bronze = [
+                    Stage(self._spark, "bronze pretransform", table) for table in tables
+                ]
+
         pretransform = None
         if self._pretransform_name:
             for table in self._preset["silver"]["preTransform"]:
                 if table["name"] == self._pretransform_name:
-                    self._pre = Stage(self._spark, "silver pretransform", table)
+                    self._pre_silver = Stage(self._spark, "silver pretransform", table)
                     break
 
         self._silver = [
@@ -151,8 +177,56 @@ class PreviewEngine:
         Returns:
             Dataframes containing the output from each run Stage.
         """
-        if self._pre:
-            df = self._pre.run(df)
+        # If we are in silverbronze mode, and an autoloader has been provided, or we are
+        # not in silverbronze mode, we need to run the preBronze stage.
+        pre_bronze_output = {}
+        if (
+            self._ds_params._mode != "silverbronze"
+            or self._ds_params._autoloader_location
+        ):
+            if self._pre_bronze:
+                for stage in self._pre_bronze:
+                    df = stage.run(df)
+                    pre_bronze_output[stage._name] = df
+        else:
+            # We are in silverbronze mode with no autoloader, so we treat first
+            # silverbronze table as initial df.
+            df = (
+                self._spark.table(self._ds_params._bronze_tables[0].get("name", ""))
+                .drop("dasl_id")
+                .limit(self._ds_params._record_limit)
+            )
+
+        if time_col := self._ds_params._time_column:
+            df = df.filter(
+                f"timestamp({time_col}) >= timestamp('{self._ds_params._start_time}') AND timestamp({time_col}) < timestamp('{self._ds_params._end_time}')"
+            )
+
+        df = df.withColumn("dasl_id", constant_udf())
+
+        self._bronze = df
+
+        # Deal with silverbronze table joins.
+        # Note: We can blind get here as validation should've caught anything missing.
+        if self._ds_params._mode == "silverbronze":
+            if alias := self._ds_params._bronze_tables[0].get("alias", None):
+                df = df.alias(alias)
+            for bronze_table in self._ds_params._bronze_tables[1:]:
+                join_df = (
+                    spark.table(bronze_table["name"])
+                    .drop("dasl_id")
+                    .limit(self._ds_params._record_limit)
+                )
+                if alias := bronze_table.get("alias", None):
+                    join_df = join_df.alias(alias)
+                df = df.join(
+                    join_df,
+                    expr(bronze_table["joinExpr"]),
+                    bronze_table.get("joinType", "left"),
+                )
+
+        if self._pre_silver:
+            df = self._pre_silver.run(df)
 
         silver_output_map = {}
         for table in self._silver:
@@ -166,15 +240,17 @@ class PreviewEngine:
             )
 
         return (
-            (df, silver_output_map, gold_output_map)
-            if self._pre
-            else (None, silver_output_map, gold_output_map)
+            (df, silver_output_map, gold_output_map, pre_bronze_output)
+            if self._pre_silver
+            else (None, silver_output_map, gold_output_map, pre_bronze_output)
         )
 
     def _render_output(
         self,
         input_df: DataFrame,
-        stage_dataframes: Tuple[DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]],
+        stage_dataframes: Tuple[
+            List[DataFrame], DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]
+        ],
         gold_table_catalog: str,
         gold_table_schema: str,
     ) -> None:
@@ -195,7 +271,7 @@ class PreviewEngine:
         def d(txt, lvl) -> None:
             displayHTML(
                 f"""
-                <div style="background-color: 
+                <div style="background-color:
                 background-color: rgb(18, 23, 26); padding: 0; margin: 0;">
                     <h{lvl} style="margin: 0; background-color: rgb(244, 234, 229);">{txt}</h{lvl}>
                 </div>
@@ -227,12 +303,16 @@ class PreviewEngine:
                             prefix=prefix + target_field.name + ".",
                         )
 
-        (pre_df, silver, gold) = stage_dataframes
-        d("Input", 1)
+        (pre_silver, silver, gold, pre_bronze) = stage_dataframes
+        d("Autoloader Input", 1)
         display(input_df)
+        d("Bronze Pre-Transform", 1)
+        for name, df in pre_bronze.items():
+            d(f"{name}", 2)
+            display(df)
         d("Silver Pre-Transform", 1)
-        if pre_df:
-            display(pre_df)
+        if pre_silver:
+            display(pre_silver)
         else:
             d("Skipped", 2)
         d("Silver Transform", 1)
@@ -326,31 +406,32 @@ class PreviewEngine:
             )
 
         # If we are using the autoloader, fetch format from preset and others.
-        if self._ds_params._mode == "autoloader":
+        if self._ds_params._mode == "autoloader" or (
+            self._ds_params._mode == "silverbronze"
+            and self._ds_params._autoloader_location
+        ):
+            if self._preset.get("bronze", {}).get("loadAsSingleVariant", False) == True:
+                self._ds_params._set_load_as_single_variant()
             if not (autoloader_conf := self._preset.get("autoloader", None)):
                 raise MissingAutoloaderConfigError()
             if not (file_format := autoloader_conf.get("format", None)):
                 raise AutoloaderMissingFieldError("format")
-            self._ds_params.set_autoloader_format(file_format)
+            self._ds_params._set_autoloader_format(file_format)
             if schemaFile := autoloader_conf.get("schemaFile", None):
-                self._ds_params.set_autoloader_schema_file(schemaFile)
-            if multiline := autoloader_conf.get("multiline", None):
-                if multiline == "true":
-                    self._ds_params.set_multiline(True)
-                else:
-                    self._ds_params.set_multiline(False)
+                self._ds_params._set_autoloader_schema_file(schemaFile)
             if cloudFiles := autoloader_conf.get("cloudFiles", None):
                 if schema_hints := cloudFiles.get("schemaHints", None):
-                    self._ds_params.set_autoloader_cloudfiles_schema_hints(schema_hints)
+                    self._ds_params._set_autoloader_cloudfiles_schema_hints(
+                        schema_hints
+                    )
                 if schema_hints_file := cloudFiles.get("schemaHintsFile", None):
-                    self._ds_params.set_autoloader_cloudfiles_schema_hint_file(
+                    self._ds_params._set_autoloader_cloudfiles_schema_hint_file(
                         schema_hints_file
                     )
 
         self._compile_stages()
 
         with self._ds_params as df:
-            self._bronze = df
             self._result_df_map = self._run(df)
             if display:
                 self._render_output(

dasl_client/preset_development/preview_parameters.py

@@ -9,11 +9,6 @@ import uuid
 from IPython import get_ipython
 
 
-@udf(StringType())
-def constant_udf(*args):
-    return "<sortable_random_id>"
-
-
 class PreviewParameters:
     """
     This class provides three methods for supplying input records to the preset development environment.
@@ -60,6 +55,78 @@ class PreviewParameters:
         .set_table("system.access.audit")
     ```
 
+    **4. SilverBronze Mode:**
+    "silverbronze" mode, works like a more advanced "table" mode. It allows for joining of multiple
+    tables as input. This mode requires setting bronze table definitions. This mode behaves in 2
+    seperate ways depending on whether an autoloader location is set or not. If an autoloader location
+    is set the first entry in the bronze table definitions is used to name and alias the autoloader's
+    input and these can be used in later join expressions. Used in this way, the autoloader will be
+    loaded as in "autoloader" mode, and run through preBronze stages before being joined with the
+    remainder of the bronze table definitions. This mimics not skipping bronze in a DataSource and
+    joining what was read in silver. If an autoloader location is not set, the behaviour instead
+    attempts to emulate a DataSource set to skip the bronze stage. That is, all preBronze and bronze
+    stages will be skipped, and the name of the first entry in the given bronze table definitions will
+    be read from instead. Any subsequent bronze table definitions will be joined against this table.
+
+    Using no autoloader location (this will read from the first table):
+    ```python
+    bronze_tables = [
+    {
+        "name": "databricks_dev.default.sev_map",
+        "alias": "tab1"
+    },
+    {
+        "name": "databricks_dev.alan_bronze.akamai_waf",
+        "alias": "tab2",
+        "joinExpr": "id::string = tab2.serviceID",
+        "joinType": "left"
+    },
+    {
+        "name": "databricks_dev.alan_silver.cloudflare_hjttp_request",
+        "alias": "tab3",
+        "joinExpr": "tab1.id::string = tab3.ClientRequestsBytes",
+        "joinType": "inner"
+    }
+    ]
+
+    ds_params = (
+        PreviewParameters(spark)
+        .from_silverbronze_tables()
+        .set_bronze_table_definitions(bronze_tables)
+    )
+
+    ps = PreviewEngine(spark, yaml_string, ds_params)
+    ps.evaluate("stage.gold")
+    ```
+
+    Using an autoloader location (this will read from the autoloader and name the df tab1):
+    ```python
+    bronze_tables = [
+    {
+        "name": "tab1"
+    },
+    {
+        "name": "databricks_dev.alan_bronze.akamai_waf",
+        "alias": "tab2",
+        "joinExpr": "id::string = tab2.serviceID",
+        "joinType": "left"
+    },
+    {
+        "name": "databricks_dev.alan_silver.cloudflare_hjttp_request",
+        "alias": "tab3",
+        "joinExpr": "tab1.id::string = tab3.ClientRequestsBytes",
+        "joinType": "inner"
+    }
+    ]
+
+    ds_params = (
+        PreviewParameters(spark)
+        .from_silverbronze_tables()
+        .set_bronze_table_definitions(bronze_tables)
+        .set_autoloader_location("s3://antimatter-dasl-testing/csamples3/mars/area1/")
+    )
+    ```
+
     **Note:**
     When using autoloader mode, this implementation requires a location to store a temporary schema for
     the loaded records. By default, this is set to `"dbfs:/tmp/schemas"`. You can change this using
@@ -94,7 +161,7 @@ class PreviewParameters:
             df (DataFrame): Internal Spark DataFrame loaded using the specified parameters.
         """
         self._spark = spark
-        self._mode = None  # [input, autoloader]
+        self._mode = None  # [input, table, autoloader, silverbronze]
         self._record_limit = 10
         self._autoloader_temp_schema_location = "dbfs:/tmp/schemas"
         self._gold_test_schemas = []
@@ -110,18 +177,97 @@ class PreviewParameters:
         self._cloudfiles_schema_hints = None
         self._cloudfiles_reader_case_sensitive = "true"
         self._cloudfiles_multiline = "true"
+        self._cloudfiles_wholetext = "false"
         self._schema_uuid_str = str(uuid.uuid4())
+        self._single_variant_column = None
 
         self._schema = None
         self._data = None
 
         self._table = None
 
+        self._bronze_tables = None
+
         self._pretransform_name = None
-        self._bronze_pre_transform: Optional[List[str]] = None
 
         self._df = None
 
+    def __create_from_autoloader(self) -> DataFrame:
+        stream_df = (
+            self._spark.readStream.format("cloudFiles")
+            .option("cloudFiles.format", self._autoloader_format)
+            .option("readerCaseSensitive", self._cloudfiles_reader_case_sensitive)
+        )
+
+        # text and wholetext needs to be handled seperately.
+        stream_df = (
+            stream_df.option("multiline", self._cloudfiles_multiline)
+            if self._autoloader_format != "text"
+            else stream_df.option("wholetext", self._cloudfiles_wholetext)
+        )
+
+        if self._single_variant_column:
+            stream_df = stream_df.option(
+                "singleVariantColumn", self._single_variant_column
+            )
+
+        if self._schema_file:
+            with open(self._schema_file, "r") as f:
+                stream_df = stream_df.schema(f.read().strip())
+        else:
+            stream_df = (
+                stream_df.option("inferSchema", "true")
+                .option("cloudFiles.inferColumnTypes", "true")
+                .option(
+                    "cloudFiles.schemaLocation",
+                    f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
+                )
+            )
+
+        if self._cloudfiles_schema_hints:
+            stream_df = stream_df.option(
+                "cloudFiles.schemaHints", self._cloudfiles_schema_hints
+            )
+        elif self._clouldfiles_schema_hints_file:
+            stream_df = stream_df.option(
+                "cloudFiles.schemaHintsFile", self._clouldfiles_schema_hints_file
+            )
+
+        stream_df = stream_df.load(self._autoloader_location).limit(self._record_limit)
+
+        query = (
+            stream_df.writeStream.format("memory")
+            .queryName("batch_data")
+            .trigger(availableNow=True)
+            .start()
+        )
+
+        query.awaitTermination()
+
+    def __create_from_silverbronze_tables_join(self) -> DataFrame:
+        if not self._bronze_tables or not len(self._bronze_tables):
+            raise MissingBronzeTablesError()
+
+        # Validate name and joinExpr are set.
+        for i in range(len(self._bronze_tables)):
+            if not self._bronze_tables[i].get("name", None):
+                raise MissingBronzeTableFieldError("name")
+            if i > 0 and not self._bronze_tables[i].get("joinExpr", None):
+                raise MissingBronzeTableFieldError("joinExpr")
+
+        # If there is an autoloader location given, we create the df now and
+        # then allow preBronze stage to run. Otherwise we skip preBronze stages
+        # and as part of running the silverbronze joins we create the df from
+        # the first entry in the bronze tables list.
+        df = None
+        if self._autoloader_location:
+            self.__create_from_autoloader()
+            df = self._spark.table("batch_data").alias(
+                self._bronze_tables[0].get("name", "")
+            )  # Use first's name.
+
+        return df
+
     def __enter__(self):
         """
         Creates a DataFrame with data using the method specified. In the case of "autoloader",
@@ -137,59 +283,10 @@ class PreviewParameters:
         elif self._mode == "table":
             self._df = self._spark.table(self._table).limit(self._record_limit)
         elif self._mode == "autoloader":
-            stream_df = (
-                self._spark.readStream.format("cloudFiles")
-                .option("cloudFiles.format", self._autoloader_format)
-                .option("multiline", self._cloudfiles_multiline)
-                .option("readerCaseSensitive", self._cloudfiles_reader_case_sensitive)
-            )
-
-            if self._schema_file:
-                with open(self._schema_file, "r") as f:
-                    stream_df = stream_df.schema(f.read().strip())
-            else:
-                stream_df = (
-                    stream_df.option("inferSchema", "true")
-                    .option("cloudFiles.inferColumnTypes", "true")
-                    .option(
-                        "cloudFiles.schemaLocation",
-                        f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
-                    )
-                )
-
-            if self._cloudfiles_schema_hints:
-                stream_df = stream_df.option(
-                    "cloudFiles.schemaHints", self._cloudfiles_schema_hints
-                )
-            elif self._clouldfiles_schema_hints_file:
-                stream_df = stream_df.option(
-                    "cloudFiles.schemaHintsFile", self._clouldfiles_schema_hints_file
-                )
-
-            stream_df = stream_df.load(self._autoloader_location).limit(
-                self._record_limit
-            )
-
-            if self._bronze_pre_transform is not None:
-                stream_df = stream_df.selectExpr(*self._bronze_pre_transform)
-
-            query = (
-                stream_df.writeStream.format("memory")
-                .queryName("batch_data")
-                .trigger(availableNow=True)
-                .start()
-            )
-
-            query.awaitTermination()
-
+            self.__create_from_autoloader()
             self._df = self._spark.table("batch_data")
-
-        if self._time_column:
-            self._df = self._df.filter(
-                f"timestamp({self._time_column}) >= timestamp('{self._start_time}') AND timestamp({self._time_column}) < timestamp('{self._end_time}')"
-            )
-
-        self._df = self._df.withColumn("dasl_id", constant_udf())
+        elif self._mode == "silverbronze":
+            self._df = self.__create_from_silverbronze_tables_join()
 
         return self._df
 
@@ -254,6 +351,36 @@ class PreviewParameters:
         self._mode = "table"
         return self
 
+    def from_silverbronze_tables(self):
+        """
+        Set the data source loader to "bronze tables" mode. Requires a list of bronze table
+        definitions to be provided.
+
+        Returns:
+            PreviewParameters: The current instance with updated configuration.
+        """
+        self._mode = "silverbronze"
+        return self
+
+    def set_bronze_table_definitions(self, definitions: List[Dict[str, str]]):
+        """
+        Set the bronze table definitions for bronze tables mode. `name` and `joinExpr` are
+        required. If `alias` is not provided, one can use the `name` to refer to the table.
+        If `joinType` is not provided, "left" is used as a default value. If pr
+
+        [
+            {
+                "name": "name",
+                "alias": "alias1",
+                "joinType": "inner",
+                "joinExpr": "base_table.col1 = alias1.col1
+            },
+            ...
+        ]
+        """
+        self._bronze_tables = definitions
+        return self
+
     def set_autoloader_temp_schema_location(self, path: str):
         """
         Set the location for the autoloader's streaming mode schema to be created. This is
@@ -311,7 +438,7 @@ class PreviewParameters:
         self._autoloader_location = location
         return self
 
-    def set_autoloader_format(self, file_format: str):
+    def _set_autoloader_format(self, file_format: str):
         """
         Used internally to set the autoloader format.
 
@@ -320,12 +447,16 @@ class PreviewParameters:
         """
         if file_format.lower() == "jsonl":
             self._autoloader_format = "json"
-            self.set_autoloader_multiline(False)
+            self._cloudfiles_multiline = "false"
+            return self
+        if file_format.lower() == "wholetext":
+            self._autoloader_format = "text"
+            self._cloudfiles_wholetext = "true"
             return self
         self._autoloader_format = file_format
         return self
 
-    def set_autoloader_schema_file(self, path: str):
+    def _set_autoloader_schema_file(self, path: str):
         """
         Set the schema file path for "autoloader" mode.
 
@@ -335,7 +466,7 @@ class PreviewParameters:
         self._schema_file = path
         return self
 
-    def set_autoloader_cloudfiles_schema_hint_file(self, path: str):
+    def _set_autoloader_cloudfiles_schema_hint_file(self, path: str):
         """
         Set the cloudFiles schema hints file path for "autoloader" mode.
 
@@ -345,7 +476,7 @@ class PreviewParameters:
         self._clouldfiles_schema_hints_file = path
         return self
 
-    def set_autoloader_cloudfiles_schema_hints(self, cloudfiles_schema_hints: str):
+    def _set_autoloader_cloudfiles_schema_hints(self, cloudfiles_schema_hints: str):
         """
         Set the cloudFiles schema hints string for "autoloader" mode.
 
@@ -355,26 +486,6 @@ class PreviewParameters:
         self._cloudfiles_schema_hints = cloudfiles_schema_hints
         return self
 
-    def set_autoloader_reader_case_sensitive(self, b: bool):
-        """
-        Set the cloudFiles reader case-sensitive boolean for "autoloader" mode.
-
-        Returns:
-            PreviewParameters: The current instance with updated configuration.
-        """
-        self._cloudfiles_reader_case_sensitive = "true" if b else "false"
-        return self
-
-    def set_autoloader_multiline(self, b: bool):
-        """
-        Set the cloudFiles multiline boolean for "autoloader" mode.
-
-        Returns:
-            PreviewParameters: The current instance with updated configuration.
-        """
-        self._cloudfiles_multiline = "true" if b else "false"
-        return self
-
     def set_pretransform_name(self, pretransform_name: str):
         """
         Set the pretransform name to use, if desired. If not set, Silver PreTransform
@@ -386,16 +497,6 @@ class PreviewParameters:
         self._pretransform_name = pretransform_name
         return self
 
-    def set_bronze_pre_transform(self, expr: List[str]):
-        """
-        Sets a pre-transform expression that will run before data is written to bronze
-
-        Returns:
-            PreviewParameters: The current instance with updated configuration.
-        """
-        self._bronze_pre_transform = expr
-        return self
-
     def set_date_range(self, column: str, start_time: str, end_time: str):
         """
         Set the TIMESTAMP column and date range to use as the input data filter to
@@ -431,6 +532,17 @@ class PreviewParameters:
         self._table = table_name
         return self
 
+    def _set_load_as_single_variant(self, col_name: Optional[str] = None):
+        """
+        Enable loadAsSingleVariant mode. This will ingest data into a single VARIANT-typed column.
+        The default name of that column is `data`.
+
+        Returns:
+            PreviewParameters: The current instance with updated configuration.
+        """
+        self._single_variant_column = col_name if col_name is not None else "data"
+        return self
+
     def add_gold_schema_table(self, gold_schema_table_name: str):
         """
         Add a gold schema temporary table name that will need to be cleaned