Fix test

Author: ghanse

docs/dqx/docs/reference/benchmarks.mdx

@@ -66,6 +66,10 @@ sidebar_position: 13
 | test_benchmark_is_aggr_not_equal | 0.296462 | 0.296800 | 0.275119 | 0.312035 | 0.013498 | 0.013448 | 0.291054 | 0.304502 | 5 | 0 | 2 | 3.37 |
 | test_benchmark_is_aggr_not_greater_than | 0.307771 | 0.315185 | 0.277924 | 0.316280 | 0.016705 | 0.010701 | 0.304974 | 0.315675 | 5 | 1 | 1 | 3.25 |
 | test_benchmark_is_aggr_not_less_than | 0.296828 | 0.303167 | 0.276023 | 0.314350 | 0.018030 | 0.033665 | 0.278439 | 0.312105 | 5 | 0 | 1 | 3.37 |
+| test_benchmark_is_area_equal_to | 0.209381 | 0.207647 | 0.205255 | 0.216179 | 0.004471 | 0.006593 | 0.206066 | 0.212659 | 5 | 0 | 1 | 4.78 |
+| test_benchmark_is_area_not_equal_to | 0.208875 | 0.207436 | 0.203626 | 0.217694 | 0.005257 | 0.004513 | 0.206265 | 0.210778 | 5 | 1 | 1 | 4.79 |
+| test_benchmark_is_area_not_greater_than | 0.177230 | 0.179352 | 0.161536 | 0.190875 | 0.010356 | 0.013261 | 0.169503 | 0.182763 | 6 | 0 | 2 | 5.64 |
+| test_benchmark_is_area_not_less_than | 0.171868 | 0.166867 | 0.161877 | 0.204155 | 0.015957 | 0.003194 | 0.164123 | 0.167316 | 6 | 1 | 1 | 5.82 |
 | test_benchmark_is_data_fresh | 0.279160 | 0.235545 | 0.231767 | 0.430390 | 0.085563 | 0.072198 | 0.233457 | 0.305655 | 5 | 1 | 1 | 3.58 |
 | test_benchmark_is_data_fresh_per_time_window | 0.259995 | 0.246444 | 0.242483 | 0.291510 | 0.022543 | 0.037112 | 0.243019 | 0.280132 | 5 | 0 | 1 | 3.85 |
 | test_benchmark_is_equal_to | 0.241270 | 0.241646 | 0.226919 | 0.248632 | 0.008799 | 0.010992 | 0.237380 | 0.248371 | 5 | 0 | 1 | 4.14 |
@@ -129,6 +133,10 @@ sidebar_position: 13
 | test_benchmark_is_null_or_empty[col6] | 0.271883 | 0.288303 | 0.233084 | 0.291877 | 0.025875 | 0.038427 | 0.251213 | 0.289639 | 5 | 0 | 1 | 3.68 |
 | test_benchmark_is_null_or_empty[col7] | 0.255479 | 0.255281 | 0.230887 | 0.289014 | 0.022139 | 0.029264 | 0.238686 | 0.267950 | 5 | 0 | 2 | 3.91 |
 | test_benchmark_is_null_or_empty[col8] | 0.219256 | 0.217609 | 0.214083 | 0.226223 | 0.005137 | 0.008602 | 0.215124 | 0.223726 | 5 | 0 | 2 | 4.56 |
+| test_benchmark_is_num_points_equal_to | 0.213472 | 0.208326 | 0.200840 | 0.228556 | 0.011595 | 0.018574 | 0.205502 | 0.224076 | 5 | 0 | 2 | 4.68 |
+| test_benchmark_is_num_points_not_equal_to | 0.211439 | 0.212084 | 0.200625 | 0.223375 | 0.008900 | 0.013585 | 0.204124 | 0.217709 | 5 | 0 | 2 | 4.73 |
+| test_benchmark_is_num_points_not_greater_than | 0.162069 | 0.161908 | 0.149400 | 0.178192 | 0.010833 | 0.014197 | 0.154168 | 0.168365 | 5 | 0 | 2 | 6.17 |
+| test_benchmark_is_num_points_not_less_than | 0.159204 | 0.157405 | 0.151457 | 0.175503 | 0.008775 | 0.008935 | 0.152260 | 0.161195 | 6 | 1 | 1 | 6.28 |
 | test_benchmark_is_ogc_valid | 0.220708 | 0.223267 | 0.206378 | 0.235649 | 0.011210 | 0.015703 | 0.211743 | 0.227446 | 5 | 0 | 2 | 4.53 |
 | test_benchmark_is_older_than_col2_for_n_days | 0.235241 | 0.230978 | 0.224354 | 0.254865 | 0.011884 | 0.013734 | 0.227788 | 0.241522 | 5 | 0 | 1 | 4.25 |
 | test_benchmark_is_older_than_n_days | 0.246935 | 0.248889 | 0.234393 | 0.253353 | 0.007733 | 0.010372 | 0.242547 | 0.252920 | 5 | 0 | 1 | 4.05 |

docs/dqx/docs/reference/quality_checks.mdx

@@ -1,7 +1,11 @@
+from collections.abc import Callable
+import operator as py_operator
+
 from pyspark.sql import Column
 import pyspark.sql.functions as F
+
 from databricks.labs.dqx.rule import register_rule
-from databricks.labs.dqx.check_funcs import make_condition, _get_normalized_column_and_expr
+from databricks.labs.dqx.check_funcs import make_condition, _get_normalized_column_and_expr, _get_limit_expr
 
 POINT_TYPE = "ST_Point"
 LINESTRING_TYPE = "ST_LineString"
@@ -10,6 +14,7 @@
 MULTILINESTRING_TYPE = "ST_MultiLineString"
 MULTIPOLYGON_TYPE = "ST_MultiPolygon"
 GEOMETRYCOLLECTION_TYPE = "ST_GeometryCollection"
+DEFAULT_SRID = 4326
 
 
 @register_rule("row")
@@ -448,3 +453,335 @@ def has_y_coordinate_between(column: str | Column, min_value: float, max_value:
         F.concat_ws("", F.lit("value `"), col_expr.cast("string"), F.lit(condition_str)),
         f"{col_str_norm}_has_y_coordinates_outside_range",
     )
+
+
+@register_rule("row")
+def is_area_equal_to(
+    column: str | Column, value: int | float | str | Column, srid: int | None = 3857, geodesic: bool = False
+) -> Column:
+    """
+    Checks if the areas of values in a geometry or geography column are equal to a specified value. By default, the 2D
+    Cartesian area in WGS84 (Pseudo-Mercator) with units of meters squared is used. An SRID can be specified to
+    transform the input values and compute areas with specific units of measure.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+        srid: Optional integer SRID to use for computing the area of the geometry or geography value (default `None`).
+            If an SRID is provided, the input value is translated and area is calculated using the units of measure of
+            the specified coordinate reference system (e.g. meters squared for `srid=3857`).
+        geodesic: Whether to use the 2D geodesic area (default `False`).
+
+    Returns:
+        Column object indicating whether the area the geometries in the input column are equal to the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_area",
+        spatial_quantity_label="area",
+        spatial_quantity_name="area",
+        compare_op=py_operator.ne,
+        compare_op_label="not equal to",
+        compare_op_name="not_equal_to",
+        srid=srid,
+        geodesic=geodesic,
+    )
+
+
+@register_rule("row")
+def is_area_not_equal_to(
+    column: str | Column, value: int | float | str | Column, srid: int | None = 3857, geodesic: bool = False
+) -> Column:
+    """
+    Checks if the areas of values in a geometry column are not equal to a specified value. By default, the 2D
+    Cartesian area in WGS84 (Pseudo-Mercator) with units of meters squared is used. An SRID can be specified to
+    transform the input values and compute areas with specific units of measure.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+        srid: Optional integer SRID to use for computing the area of the geometry or geography value (default `None`).
+            If an SRID is provided, the input value is translated and area is calculated using the units of measure of
+            the specified coordinate reference system (e.g. meters squared for `srid=3857`).
+        geodesic: Whether to use the 2D geodesic area (default `False`).
+
+    Returns:
+        Column object indicating whether the area the geometries in the input column are not equal to the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_area",
+        spatial_quantity_label="area",
+        spatial_quantity_name="area",
+        compare_op=py_operator.eq,
+        compare_op_label="equal to",
+        compare_op_name="equal_to",
+        srid=srid,
+        geodesic=geodesic,
+    )
+
+
+@register_rule("row")
+def is_area_not_greater_than(
+    column: str | Column, value: int | float | str | Column, srid: int | None = 3857, geodesic: bool = False
+) -> Column:
+    """
+    Checks if the areas of values in a geometry column are not greater than a specified limit. By default, the 2D
+    Cartesian area in WGS84 (Pseudo-Mercator) with units of meters squared is used. An SRID can be specified to
+    transform the input values and compute areas with specific units of measure.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+        srid: Optional integer SRID to use for computing the area of the geometry or geography value (default `None`).
+            If an SRID is provided, the input value is translated and area is calculated using the units of measure of
+            the specified coordinate reference system (e.g. meters squared for `srid=3857`).
+        geodesic: Whether to use the 2D geodesic area (default `False`).
+
+    Returns:
+        Column object indicating whether the area the geometries in the input column is greater than the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_area",
+        spatial_quantity_label="area",
+        spatial_quantity_name="area",
+        compare_op=py_operator.gt,
+        compare_op_label="greater than",
+        compare_op_name="greater_than",
+        srid=srid,
+        geodesic=geodesic,
+    )
+
+
+@register_rule("row")
+def is_area_not_less_than(
+    column: str | Column, value: int | float | str | Column, srid: int | None = 3857, geodesic: bool = False
+) -> Column:
+    """
+    Checks if the areas of values in a geometry column are not less than a specified limit. By default, the 2D
+    Cartesian area in WGS84 (Pseudo-Mercator) with units of meters squared is used. An SRID can be specified to
+    transform the input values and compute areas with specific units of measure.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+        srid: Optional integer SRID to use for computing the area of the geometry or geography value (default `None`).
+            If an SRID is provided, the input value is translated and area is calculated using the units of measure of
+            the specified coordinate reference system (e.g. meters squared for `srid=3857`).
+        geodesic: Whether to use the 2D geodesic area (default `False`).
+
+    Returns:
+        Column object indicating whether the area the geometries in the input column is less than the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_area",
+        spatial_quantity_label="area",
+        spatial_quantity_name="area",
+        compare_op=py_operator.lt,
+        compare_op_label="less than",
+        compare_op_name="less_than",
+        srid=srid,
+        geodesic=geodesic,
+    )
+
+
+@register_rule("row")
+def is_num_points_equal_to(column: str | Column, value: int | float | str | Column) -> Column:
+    """
+    Checks if the number of coordinate pairs in values of a geometry column is equal to a specified value.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+
+    Returns:
+        Column object indicating whether the number of coordinate pairs in the geometries of the input column is
+        equal to the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_npoints",
+        spatial_quantity_label="number of coordinates",
+        spatial_quantity_name="num_points",
+        compare_op=py_operator.ne,
+        compare_op_label="not equal to",
+        compare_op_name="not_equal_to",
+    )
+
+
+@register_rule("row")
+def is_num_points_not_equal_to(column: str | Column, value: int | float | str | Column) -> Column:
+    """
+    Checks if the number of coordinate pairs in values of a geometry column is not equal to a specified value.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+
+    Returns:
+        Column object indicating whether the number of coordinate pairs in the geometries of the input column is not
+        equal to the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_npoints",
+        spatial_quantity_label="number of coordinates",
+        spatial_quantity_name="num_points",
+        compare_op=py_operator.eq,
+        compare_op_label="equal to",
+        compare_op_name="equal_to",
+    )
+
+
+@register_rule("row")
+def is_num_points_not_greater_than(column: str | Column, value: int | float | str | Column) -> Column:
+    """
+    Checks if the number of coordinate pairs in the values of a geometry column is not greater than a specified limit.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+
+    Returns:
+        Column object indicating whether the number of coordinate pairs in the geometries of the input column is
+        greater than the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_npoints",
+        spatial_quantity_label="number of coordinates",
+        spatial_quantity_name="num_points",
+        compare_op=py_operator.gt,
+        compare_op_label="greater than",
+        compare_op_name="greater_than",
+    )
+
+
+@register_rule("row")
+def is_num_points_not_less_than(column: str | Column, value: int | float | str | Column) -> Column:
+    """
+    Checks if the number of coordinate pairs in values of a geometry column is not less than a specified limit.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+
+    Returns:
+        Column object indicating whether the number of coordinate pairs in the geometries of the input column is
+        less than the provided value
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    return _compare_sql_function_result(
+        column,
+        value,
+        spatial_function="st_npoints",
+        spatial_quantity_label="number of coordinates",
+        spatial_quantity_name="num_points",
+        compare_op=py_operator.lt,
+        compare_op_label="less than",
+        compare_op_name="less_than",
+    )
+
+
+def _compare_sql_function_result(
+    column: str | Column,
+    value: int | float | str | Column,
+    spatial_function: str,
+    spatial_quantity_label: str,
+    spatial_quantity_name: str,
+    compare_op: Callable[[Column, Column], Column],
+    compare_op_label: str,
+    compare_op_name: str,
+    srid: int | None = None,
+    geodesic: bool = False,
+) -> Column:
+    """
+    Compares the results from applying a spatial SQL function (e.g. `st_area`) on a geometry column against a limit
+    using the specified comparison operator.
+
+    Args:
+        column: Column to check; can be a string column name or a column expression
+        value: Value to use in the condition as number, column name or sql expression
+        spatial_function: Spatial SQL function as a string (e.g. `st_npoints`)
+        spatial_quantity_label: Spatial quantity label (e.g. `number of coordinates` )
+        spatial_quantity_name: Spatial quantity identifier (e.g. `num_points`)
+        compare_op: Comparison operator (e.g., `operator.gt`, `operator.lt`).
+        compare_op_label: Human-readable label for the comparison (e.g., 'greater than').
+        compare_op_name: Name identifier for the comparison (e.g., 'greater_than').
+        srid: Optional integer SRID for computing measurements on the converted geometry or geography value (default `None`).
+        geodesic: Whether to convert the input column to a geography type for computing geodesic distances.
+
+    Returns:
+        Column object indicating whether the area the geometries in the input column is less than the provided limit
+
+    Note:
+        This function requires Databricks serverless compute or runtime 17.1 or above.
+    """
+    col_str_norm, col_expr_str, col_expr = _get_normalized_column_and_expr(column)
+    value_expr = _get_limit_expr(value)
+    # NOTE: This function is currently only available in Databricks runtime 17.1 or above or in
+    #   Databricks SQL, due to the use of the `try_to_geometry` and `st_area` functions.
+    if geodesic:
+        spatial_conversion_expr = f"try_to_geography({col_str_norm})"
+        spatial_data_type = "geography"
+    elif srid:
+        spatial_conversion_expr = f"st_transform(st_setsrid(try_to_geometry({col_str_norm}), {DEFAULT_SRID}), {srid})"
+        spatial_data_type = "geometry"
+    else:
+        spatial_conversion_expr = f"try_to_geometry({col_str_norm})"
+        spatial_data_type = "geometry"
+
+    is_valid_cond = F.expr(f"{spatial_conversion_expr} IS NULL")
+    is_valid_message = F.concat_ws(
+        "",
+        F.lit("value `"),
+        col_expr.cast("string"),
+        F.lit(f"` in column `{col_expr_str}` is not a valid {spatial_data_type}"),
+    )
+    compare_cond = compare_op(F.expr(f"{spatial_function}({spatial_conversion_expr})"), value_expr)
+    compare_message = F.concat_ws(
+        "",
+        F.lit("value `"),
+        col_expr.cast("string"),
+        F.lit(f"` in column `{col_expr_str}` has {spatial_quantity_label} {compare_op_label} value: "),
+        value_expr.cast("string"),
+    )
+    condition = F.when(col_expr.isNull(), F.lit(None)).otherwise(is_valid_cond | compare_cond)
+
+    return make_condition(
+        condition,
+        F.when(is_valid_cond, is_valid_message).otherwise(compare_message),
+        f"{col_str_norm}_{spatial_quantity_name}_{compare_op_name}_limit",
+    )