add ansi-safe casting (#303)

* add ansi-safe casting

* update broken sql

* add unit tests

* add unit tests for actions.py ansi error check

* add 2 tests for general raise case in actions.py

* unit tests: fixes and cover general raise condition

* fix test_actions unit tests

* add code coverage

* update docs, rename tests

Author: yankomo

docs/user_guide/data_quality_rules.md

@@ -106,6 +106,7 @@ ADD CONSTRAINT action CHECK (
 **Tip:**  
 - Match your expectation format to the rule type for correct validation.
 - Use `row_dq` for per-row checks, `agg_dq` for summary statistics, and `query_dq` for advanced SQL-based checks.
+- If running in Databricks serverless compute, expectations need to be ANSI compliant. See the ANSI Mode section of the [Serverless doc](serverless.md) for details.
 
 
 Below are the details and examples for each rule type:

docs/user_guide/serverless.md

@@ -66,3 +66,9 @@ def process_data():
 result_df = process_data()
 ```
 
+### IMPORTANT: ANSI Mode
+When using serverless compute on Databricks, ANSI mode is **on by default**, which enforces stricter standards. If you use serverless compute you have two options:
+- Turn ANSI mode off: set `spark.sql.ansi.enabled` to `false`
+- Or, make sure your expectations are written in an ANSI-compliant way
+    - Use `try_cast` instead of `CAST`
+    - Please see Databricks documentation on ANSI mode for more guidance on other requirements for ANSI compliance

spark_expectations/core/exceptions.py

@@ -80,3 +80,13 @@ class SparkExpectationsInvalidAggDQExpectationException(Exception):
     """
     Throw this exception when an invalid agg_dq expectation is encountered
     """
+
+def raise_if_ansi_exception(e: Exception, rule_name: str, rule_expectation: str) -> None:
+    """Check if exception is likely cast error due to ANSI mode. Add that info to error message and raise."""
+    if "CAST_INVALID_INPUT" in str(e):
+        raise SparkExpectationsMiscException(
+            f"Cast error while evaluating rule '{rule_name}' with expectation/SQL '{rule_expectation}'. "
+            f"This may be caused by rule expectation/SQL incompatibility with spark.sql.ansi.enabled=true "
+            f"(default on Databricks Serverless). Please either update your expectation/SQL "
+            f"to be compliant with ANSI mode, disable ANSI mode, or run not on serverless."
+        ) from e

spark_expectations/sinks/utils/report.py

@@ -14,17 +14,15 @@
     size,
     concat_ws,
     abs,
-    least,
-    greatest,
     get_json_object,
     lit,
-    trim,
 )
-from pyspark.sql.types import DoubleType, StringType, DecimalType, TimestampType
+from pyspark.sql.types import DoubleType, StringType, TimestampType
 
 from spark_expectations import _log
 from spark_expectations.core.context import SparkExpectationsContext
 from spark_expectations.core.exceptions import SparkExpectationsMiscException
+from spark_expectations.utils.udf import safe_cast
 
 
 @dataclass
@@ -176,7 +174,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumn(
                     "total_records_only_nbr",
                     when(col("_total_records_str") == "", lit(None).cast("bigint"))
-                    .otherwise(col("_total_records_str").cast("bigint")),
+                    .otherwise(safe_cast(self.spark, "_total_records_str", "bigint")),
                 )
                 .withColumn(
                     "_valid_records_str",
@@ -185,7 +183,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumn(
                     "valid_records_only_nbr",
                     when(col("_valid_records_str") == "", lit(None).cast("bigint"))
-                    .otherwise(col("_valid_records_str").cast("bigint")),
+                    .otherwise(safe_cast(self.spark, "_valid_records_str", "bigint")),
                 )
                 .drop("_total_records_str", "_valid_records_str")
                 .withColumn(
@@ -194,27 +192,30 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                         col("total_records_only_nbr").isNull() & col("valid_records_only_nbr").isNull(),
                         lit(100),
                     )
-                    .when(
-                        col("total_records_only_nbr").isNull() & col("valid_records_only_nbr").isNull(),
-                        lit(100),
-                    )
                     .when(
                         col("total_records_only_nbr").isNotNull() & col("valid_records_only_nbr").isNull(),
                         lit(0),
                     )
                     .otherwise(
                         coalesce(
-                            (
-                                100
+                            safe_cast(
+                                self.spark,
+                                """
+                                100 
                                 * least(
-                                    abs(trim(col("valid_records_only_nbr"))),
-                                    abs(trim(col("total_records_only_nbr"))),
-                                )
-                                / greatest(
-                                    abs(trim(col("valid_records_only_nbr"))),
-                                    abs(trim(col("total_records_only_nbr"))),
-                                )
-                            ).cast(DecimalType(20, 2)),
+                                        abs(valid_records_only_nbr),
+                                        abs(total_records_only_nbr)
+                                    ) / nullif(
+                                            greatest(
+                                                abs(valid_records_only_nbr),
+                                                abs(total_records_only_nbr)
+                                            ),
+                                            0
+                                        )
+                                """,
+                                "decimal(20, 2)"
+                            )
+                            ,
                             lit(0),
                         )
                     ),
@@ -225,10 +226,6 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                         col("total_records_only_nbr").isNull() & col("valid_records_only_nbr").isNull(),
                         lit(0),
                     )
-                    .when(
-                        col("total_records_only_nbr").isNull() & col("valid_records_only_nbr").isNull(),
-                        lit(0),
-                    )
                     .when(
                         col("total_records_only_nbr").isNotNull() & col("valid_records_only_nbr").isNull(),
                         lit(100),
@@ -255,11 +252,11 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                     abs(
                         coalesce(
                             coalesce(
-                                trim(col("total_records_only_nbr")).cast("bigint"),
+                                col("total_records_only_nbr"),
                                 lit(0),
                             )
                             - coalesce(
-                                trim(col("valid_records_only_nbr")).cast("bigint"),
+                                col("valid_records_only_nbr"),
                                 lit(0),
                             ),
                             lit(0),
@@ -333,7 +330,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumnRenamed("source_dq_error_row_count", "failed_records")
                 .withColumn(
                     "success_percentage",
-                    (col("valid_records").cast("double") / col("total_records").cast("double")) * 100,
+                    (safe_cast(self.spark, "valid_records", "double") / safe_cast(self.spark, "total_records", "double")) * 100,
                 )
             )
 

spark_expectations/utils/actions.py

@@ -21,6 +21,7 @@
 from spark_expectations.core.exceptions import (
     SparkExpectationsMiscException,
     SparkExpectOrFailException,
+    raise_if_ansi_exception
 )
 from spark_expectations.utils.udf import get_actions_list, remove_empty_maps
 
@@ -174,8 +175,11 @@ def agg_query_dq_detailed_result(
                         _agg_dq_expectation_aggstring = _agg_dq_expectation_match.group(1)
                         _agg_dq_expectation_expr = _agg_dq_expectation_match.group(2)
                         _agg_dq_expectation_cond_expr = expr(_agg_dq_expectation_aggstring)
-
-                        _agg_dq_actual_count_value_raw = df.agg(_agg_dq_expectation_cond_expr).collect()[0][0]
+                        try:
+                            _agg_dq_actual_count_value_raw = df.agg(_agg_dq_expectation_cond_expr).collect()[0][0]
+                        except Exception as e:
+                            raise_if_ansi_exception(e, _dq_rule["rule"], _dq_rule["expectation"])
+                            raise
 
                         # Handle NoneType (nulls)
                         if _agg_dq_actual_count_value_raw is None:
@@ -252,7 +256,11 @@ def agg_query_dq_detailed_result(
 
                         _agg_dq_expectation_cond_expr = expr(_agg_dq_expectation_aggstring)
 
-                        _agg_dq_actual_count_value = int(df.agg(_agg_dq_expectation_cond_expr).collect()[0][0])
+                        try:
+                            _agg_dq_actual_count_value = int(df.agg(_agg_dq_expectation_cond_expr).collect()[0][0])
+                        except Exception as e:
+                            raise_if_ansi_exception(e, _dq_rule["rule"], _agg_dq_expectation_aggstring)
+                            raise
 
                         _agg_dq_expression_str_lower = (
                             str(_agg_dq_actual_count_value) + _agg_dq_expectation_expr_lowerbound
@@ -319,7 +327,11 @@ def agg_query_dq_detailed_result(
                     )
                 ):
                     for _key, _querydq_query in sub_key_value.items():
-                        _querydq_df = _context.spark.sql(_dq_rule["expectation" + "_" + _key])
+                        try:
+                            _querydq_df = _context.spark.sql(_dq_rule["expectation" + "_" + _key])
+                        except Exception as e:
+                            raise_if_ansi_exception(e, _dq_rule["rule"], _dq_rule["expectation" + "_" + _key])
+                            raise
                         querydq_output.append(
                             (
                                 _context.get_run_id,
@@ -357,7 +369,15 @@ def agg_query_dq_detailed_result(
                         def execute_sql_and_get_result(
                             _se_context: SparkExpectationsContext, query: str
                         ) -> Union[int, float, str]:
-                            return _se_context.spark.sql(f"SELECT ({query}) AS OUTPUT").collect()[0][0] if query else 0
+                            if query:
+                                try:
+                                    result = _se_context.spark.sql(f"SELECT ({query}) AS OUTPUT").collect()[0][0]
+                                except Exception as e:
+                                    raise_if_ansi_exception(e, _dq_rule["rule"], f"SELECT ({query}) AS OUTPUT")
+                                    raise
+                            else:
+                                result = 0
+                            return result
 
                         # function to get the query outputs
                         _querydq_source_query_output = execute_sql_and_get_result(_context, match.group(1))
@@ -380,8 +400,12 @@ def execute_sql_and_get_result(
 
                 _querydq_status_query = "SELECT (" + str(_dq_rule["expectation"]) + ") AS OUTPUT"
 
-                _query_dq_result = int(_context.spark.sql(_querydq_status_query).collect()[0][0])
-
+                try:
+                    _query_dq_result = int(_context.spark.sql(_querydq_status_query).collect()[0][0])
+                except Exception as e:
+                    raise_if_ansi_exception(e, _dq_rule["rule"], _querydq_status_query)
+                    raise
+                
                 status = "pass" if _query_dq_result else "fail"
 
                 if _source_dq_status:

spark_expectations/utils/udf.py

@@ -1,5 +1,5 @@
-from pyspark.sql import Column
-from pyspark.sql.functions import filter, size, transform, when, lit, array
+from pyspark.sql import SparkSession, Column
+from pyspark.sql.functions import filter, size, transform, when, lit, array, expr
 
 
 def remove_empty_maps(column: Column) -> Column:
@@ -40,3 +40,20 @@ def get_actions_list(column: Column) -> Column:
     column = remove_passing_status_maps(column)
     action_if_failed = transform(column, lambda x: x["action_if_failed"])
     return when(size(action_if_failed) == 0, array(lit("ignore"))).otherwise(action_if_failed)  # pragma: no cover
+
+def safe_cast(spark: SparkSession, column: str, target_type: str) -> Column:
+    """
+    Checks if ANSI mode is enabled. If enabled, uses try_cast to cast the column to the target type. If not, uses cast.
+    Args:
+        spark: SparkSession
+        column: column to cast (provided as a string)
+        target_type: target type to cast to
+
+    Returns:
+        Column: the casted column
+    """
+    ansi_enabled = spark.conf.get("spark.sql.ansi.enabled", "false").lower() == "true"
+    if ansi_enabled:
+        return expr(f"try_cast({column} as {target_type})")
+    else: 
+        return expr(f"cast({column} as {target_type})")

tests/integration/utils/test_udf.py

@@ -1,5 +1,6 @@
+from pyspark.sql.types import LongType, DoubleType
 from spark_expectations.core import get_spark_session
-from spark_expectations.utils.udf import remove_empty_maps, get_actions_list
+from spark_expectations.utils.udf import remove_empty_maps, get_actions_list, safe_cast
 
 spark = get_spark_session()
 
@@ -57,3 +58,42 @@ def test_get_actions_list():
 
     for itr in range(0, 4):
         assert results[itr].actions == expected_output[itr]
+
+def test_safe_cast_ansi_disabled():
+    # with default spark.sql.ansi.enabled=false, safe cast should return
+    #  a column equavalent to cast(col as type)
+    spark.conf.set("spark.sql.ansi.enabled", "false")
+
+    columns = ["Name", "Value"]
+    data = [("thing", "123")]
+    df = spark.createDataFrame(data, schema=columns)
+
+    result_df = df.withColumn("casted_value", safe_cast(spark, "Value", "bigint"))
+    result = result_df.select("casted_value").collect()[0]["casted_value"]
+
+    assert result == 123
+    assert result_df.schema["casted_value"].dataType == LongType()
+
+
+def test_safe_cast_ansi_enabled():
+    # with spark.sql.ansi.enabled=true, safe cast should succeed if it's a castable value and should 
+    #  return Null if it's not a valid cast (without safe cast it would throw an exception insead)
+    spark.conf.set("spark.sql.ansi.enabled", "true")
+
+    columns = ["Name", "Value", "Color"]
+    data = [("mug", "10", "red")]
+    df = spark.createDataFrame(data, schema=columns)
+
+    success_result_df = df.withColumn("success_casted_value", safe_cast(spark, "Value", "double"))
+    success_result = success_result_df.select("success_casted_value").collect()[0]["success_casted_value"]
+
+    fail_result_df = df.withColumn("fail_casted_value", safe_cast(spark, "Color", "double"))
+    fail_result = fail_result_df.select("fail_casted_value").collect()[0]["fail_casted_value"]
+
+    # re-set ansi mode so it doesn't affect other tests
+    spark.conf.set("spark.sql.ansi.enabled", "false")
+
+    assert success_result == 10
+    assert success_result_df.schema["success_casted_value"].dataType == DoubleType()
+
+    assert fail_result is None

tests/unit/core/test_expectations.py

@@ -3,7 +3,7 @@
 """
 import pytest
 
-from spark_expectations.core.exceptions import SparkExpectationsUserInputOrConfigInvalidException
+from spark_expectations.core.exceptions import SparkExpectationsMiscException, SparkExpectationsUserInputOrConfigInvalidException, raise_if_ansi_exception
 from spark_expectations.core.expectations import SparkExpectations, WrappedDataFrameWriter
 
 
@@ -472,3 +472,30 @@ def test_add_hash_columns_trims_whitespace_produces_valid_hash(
         assert row["id_hash"] is not None
         assert len(row["id_hash"]) == 32
         assert all(c in "0123456789abcdef" for c in row["id_hash"])
+
+
+class TestRaiseIfAnsiException:
+    """Test cases for SparkExpectations exception raise_if_ansi_exception method."""
+
+    def test_raise_if_ansi_exception(self):
+        ansi_exception = Exception("SparkExpectationsMiscException: error occurred while processing spark expectations"
+        " error occurred while executing func_process error occurred while running expectations error occurred"
+        " while running agg_query_dq_detailed_result [CAST_INVALID_INPUT] The value '' of the type \"STRING\""
+        " cannot be cast to \"BIGINT\" because it is malformed. Correct the value as per the syntax, or change its"
+        " target type. Use `try_cast` to tolerate malformed input and return NULL instead.")
+
+        rule_name = "Column A Greater than Column B"
+        rule_expectation = "[col_A] > [col_B]"
+
+        with pytest.raises(SparkExpectationsMiscException) as exception_info:
+            raise_if_ansi_exception(ansi_exception, rule_name, rule_expectation)
+        
+        assert f"{rule_name}" in str(exception_info.value)
+        assert f"{rule_expectation}" in str(exception_info.value)
+
+    def test_no_raise_ansi_exception(self):
+        non_ansi_exception = Exception("Exception message about some things that are not related to ANSI casting.")
+
+        result = raise_if_ansi_exception(non_ansi_exception, "Foo", "Bar")
+
+        assert result is None