Spatial data validations

src/databricks/labs/dqx/checks_resolver.py

@@ -6,6 +6,7 @@
 from contextlib import contextmanager
 
 from databricks.labs.dqx import check_funcs
+from databricks.labs.dqx.geo import check_funcs as geo_check_funcs
 from databricks.labs.dqx.errors import InvalidCheckError
 
 logger = logging.getLogger(__name__)
@@ -30,6 +31,8 @@ def resolve_check_function(
     """
     logger.debug(f"Resolving function: {function_name}")
     func = getattr(check_funcs, function_name, None)  # resolve using predefined checks first
+    if not func:
+        func = getattr(geo_check_funcs, function_name, None)  # resolve using prefedined geo checks
     if not func and custom_check_functions:
         func = custom_check_functions.get(function_name)  # returns None if not found
     if fail_on_missing and not func:

src/databricks/labs/dqx/llm/resources/yaml_checks_examples.yml

@@ -320,6 +320,91 @@
     for_each_column:
     - col3
     - col5
+- criticality: error
+  check:
+    function: is_latitude
+    arguments:
+      column: col2
+- criticality: error
+  check:
+    function: is_longitude
+    arguments:
+      column: col2
+- criticality: error
+  check:
+    function: is_geometry
+    arguments:
+      column: point_geom
+- criticality: error
+  check:
+    function: is_geography
+    arguments:
+      column: point_geom
+- criticality: error
+  check:
+    function: is_point
+    arguments:
+      column: point_geom
+- criticality: error
+  check:
+    function: is_linestring
+    arguments:
+      column: linestring_geom
+- criticality: error
+  check:
+    function: is_polygon
+    arguments:
+      column: polygon_geom
+- criticality: error
+  check:
+    function: is_multipoint
+    arguments:
+      column: multipoint_geom
+- criticality: error
+  check:
+    function: is_multilinestring
+    arguments:
+      column: multilinestring_geom
+- criticality: error
+  check:
+    function: is_multipolygon
+    arguments:
+      column: multipolygon_geom
+- criticality: error
+  check:
+    function: is_geometrycollection
+    arguments:
+      column: geometrycollection_geom
+- criticality: error
+  check:
+    function: is_ogc_valid
+    arguments:
+      column: point_geom
+- criticality: error
+  check:
+    function: is_non_empty_geometry
+    arguments:
+      column: point_geom
+- criticality: error
+  check:
+    function: has_dimension
+    arguments:
+      column: polygon_geom
+      dimension: 2
+- criticality: error
+  check:
+    function: has_x_coordinate_between
+    arguments:
+      column: polygon_geom
+      min_value: 0.0
+      max_value: 10.0
+- criticality: error
+  check:
+    function: has_y_coordinate_between
+    arguments:
+      column: polygon_geom
+      min_value: 0.0
+      max_value: 10.0
 - criticality: error
   check:
     function: is_not_null

tests/conftest.py

@@ -1,4 +1,5 @@
 import os
+import re
 from collections.abc import Callable, Generator
 from dataclasses import replace
 from functools import cached_property
@@ -40,6 +41,44 @@ def set_utc_timezone():
     os.environ.pop("TZ")
 
 
+@pytest.fixture
+def skip_if_runtime_not_geo_compatible(ws, debug_env):
+    """
+    Skip the test if the cluster runtime does not support the required geo functions, i.e.
+    * serverless clusters have the required geo functions
+    * standard clusters require runtime 17.1 or above
+
+    Args:
+        ws (WorkspaceClient): Workspace client to interact with Databricks.
+        debug_env (dict): Test environment variables.
+    """
+    if "DATABRICKS_SERVERLESS_COMPUTE_ID" in debug_env:
+        return  # serverless clusters have the required geo functions
+
+    # standard clusters require runtime 17.1 or above
+    cluster_id = debug_env.get("DATABRICKS_CLUSTER_ID")
+    if not cluster_id:
+        raise ValueError("DATABRICKS_CLUSTER_ID is not set in debug_env")
+
+    # Fetch cluster details
+    cluster_info = ws.clusters.get(cluster_id)
+    runtime_version = cluster_info.spark_version
+
+    if not runtime_version:
+        raise ValueError(f"Unable to retrieve runtime version for cluster {cluster_id}")
+
+    # Extract major and minor version numbers
+    match = re.match(r"(\d+)\.(\d+)", runtime_version)
+    if not match:
+        raise ValueError(f"Invalid runtime version format: {runtime_version}")
+
+    major, minor = [int(x) for x in match.groups()]
+    valid = major > 17 or (major == 17 and minor >= 1)
+
+    if not valid:
+        pytest.skip("This test requires a cluster with runtime 17.1 or above")
+
+
 class CommonUtils:
     def __init__(self, env_or_skip_fixture: Callable[[str], str], ws: WorkspaceClient):
         self._env_or_skip = env_or_skip_fixture