dasl-client 1.0.22__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)
@@ -240,60 +320,40 @@ class PreviewEngine:
             d(f"{name}", 2)
             display(df)
         d("Gold", 1)
-        for name, df in gold.items():
-            d(f"{name}", 2)
+        for full_name, df in gold.items():
+            d(f"{full_name}", 2)
             d("Stage output", 3)
             display(df)
 
             # NOTE: Name is stored as Gold_name/Silver_input. So we need to get just the Gold table
             # name that we are comparing the dataframe metadata to.
-            name = name.split("/")[0]
+            name = full_name.split("/")[0]
             fqn_gold_table_name = f"{self.force_apply_backticks(gold_table_catalog)}.{self.force_apply_backticks(gold_table_schema)}.{self.force_apply_backticks(name)}"
 
             if not self._spark.catalog.tableExists(f"{fqn_gold_table_name}"):
                 raise UnknownGoldTableError(name, gold_table_schema)
 
-            # Performs the type check.
+            # Create a temporary table to perform the type check
             delta_df = self._spark.table(f"{fqn_gold_table_name}").limit(0)
-            unioned_df = delta_df.unionByName(df, allowMissingColumns=True)
-
-            # Now we check no new columns.
-            if not set(df.columns).issubset(delta_df.columns):
-                raise GoldTableCompatibilityError(
-                    f"Extra columns provided: {', '.join([col for col in df.columns if col not in delta_df.columns])}"
-                )
+            delta_df.write.mode("overwrite").save(
+                f"{self._ds_params.get_autoloader_temp_schema_location()}/{full_name}"
+            )
 
-            # Now we check no new fields in STRUCT columns.
-            for field in delta_df.schema.fields:
-                if isinstance(field.dataType, StructType) and field.name in df.columns:
-                    # Retrieve the corresponding field from the DataFrame's schema.
-                    df_field = next(f for f in df.schema.fields if f.name == field.name)
-                    check_struct_compatibility(field, df_field)
+            # Update the params to indicate we've added a testing temp gold table
+            self._ds_params.add_gold_schema_table(full_name)
 
-            # Check nullable columns exist, and data what we are inserting is set.
-            non_nullable_cols = [
-                field.name for field in delta_df.schema.fields if not field.nullable
-            ]
-            null_checks = [
-                sum_(when(col_(col).isNull(), 1).otherwise(0)).alias(col)
-                for col in non_nullable_cols
-            ]
-            null_counts = df.select(null_checks).collect()[0].asDict()
-            cols_with_nulls = []
+            # Perform the type checks by trying to insert data into the table
             try:
-                cols_with_nulls = [
-                    col_name for col_name, count in null_counts.items() if count > 0
-                ]
-            except TypeError:
-                # There were no records returned and so null_counts == None.
-                pass
-            if cols_with_nulls:
+                df.write.mode("append").save(
+                    f"{self._ds_params.get_autoloader_temp_schema_location()}/{full_name}"
+                )
+            except Exception as e:
                 raise GoldTableCompatibilityError(
-                    f"Record with null data found for non-nullable columns: {', '.join([col for col in cols_with_nulls])}"
+                    f"Preset gold table '{full_name}' did not match the gold schema for {fqn_gold_table_name}: {repr(e)}"
                 )
 
             d("Resultant gold table preview", 3)
-            display(unioned_df)
+            display(df)
 
     def is_backtick_escaped(self, name: str) -> bool:
         """
@@ -346,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(