support h3 function

src/clickhouse_orm/funcs.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from functools import wraps
 from inspect import signature, Parameter
 from types import FunctionType
@@ -1158,10 +1160,161 @@ class F(Cond, FunctionOperatorsMixin, metaclass=FMeta):  # pylint: disable=R0904
     def hasAny(arr, x):
         return F("hasAny", arr, x)
 
+    @staticmethod
+    def greatCircleDistance(lon1, lat1, lon2, lat2):
+        """
+        Calculates the distance between two points on the Earth’s surface
+        using the great-circle formula.
+        """
+        return F("greatCircleDistance", lon1, lat1, lon2, lat2)
+
+    @staticmethod
+    def geoDistance(lon1, lat1, lon2, lat2):
+        """
+        Similar to greatCircleDistance but calculates the distance on WGS-84 ellipsoid
+        instead of sphere.
+        This is more precise approximation of the Earth Geoid.
+        The performance is the same as for greatCircleDistance (no performance drawback).
+        It is recommended to use geoDistance to calculate the distances on Earth.
+        """
+        return F("geoDistance", lon1, lat1, lon2, lat2)
+
+    @staticmethod
+    def greatCircleAngle(lon1, lat1, lon2, lat2):
+        """
+        Calculates the central angle between two points on the Earth’s surface
+        using the great-circle formula.
+        """
+        return F("greatCircleAngle", lon1, lat1, lon2, lat2)
+
+    @staticmethod
+    def pointInPolygon(x, y, polygon):
+        """
+        Checks whether the point belongs to the polygon on the plane.
+        """
+        return F("pointInPolygon", (x, y), list(polygon))
+
     @staticmethod
     def geohashEncode(x, y, precision=12):
+        """
+        Encodes latitude and longitude as a geohash-string.
+        """
         return F("geohashEncode", x, y, precision)
 
+    @staticmethod
+    def geohashDecode(encoded):
+        """
+        Decodes any geohash-encoded string into longitude and latitude.
+        """
+        return F("geohashDecode", encoded)
+
+    @staticmethod
+    def geohashesInBox(
+        longitude_min,
+        latitude_min,
+        longitude_max,
+        latitude_max,
+        precision,
+    ):
+        """
+        Returns an array of geohash-encoded strings of given precision that fall inside
+        and intersect boundaries of given box, basically a 2D grid flattened into array.
+        """
+        return F("geohashesInBox", longitude_min, latitude_min,
+                 longitude_max, latitude_max, precision)
+
+    @staticmethod
+    def h3IsValid(h3index):
+        """
+        Verifies whether the number is a valid H3 index.
+        """
+        return F("h3IsValid", h3index)
+
+    @staticmethod
+    def h3GetResolution(h3index):
+        """
+        Defines the resolution of the given H3 index.
+        """
+        return F("h3GetResolution", h3index)
+
+    @staticmethod
+    def h3EdgeAngle(resolution):
+        """
+        Calculates the average length of the H3 hexagon edge in grades.
+        """
+        return F("h3EdgeAngle", resolution)
+
+    @staticmethod
+    def h3EdgeLengthM(resolution):
+        """
+        Calculates the average length of the H3 hexagon edge in meters.
+        """
+        return F("h3EdgeLengthM", resolution)
+
+    @staticmethod
+    def h3EdgeLengthKm(resolution):
+        """
+        Calculates the average length of the H3 hexagon edge in kilometers.
+        """
+        return F("h3EdgeLengthKm", resolution)
+
+    @staticmethod
+    def geoToH3(lon, lat, resolution):
+        """
+        Returns H3 point index (lon, lat) with specified resolution.
+        """
+        return F("geoToH3", lon, lat, resolution)
+
+    @staticmethod
+    def h3ToGeo(h3index):
+        """
+        Returns the centroid longitude and latitude corresponding to the provided H3 index.
+        """
+        return F("h3ToGeo", h3index)
+
+    @staticmethod
+    def h3ToGeoBoundary(h3index):
+        """
+        Returns array of pairs (lon, lat), which corresponds to the boundary
+        of the provided H3 index.
+        """
+        return F("h3ToGeoBoundary", h3index)
+
+    @staticmethod
+    def h3GetBaseCell(index):
+        """
+        Returns the base cell number of the H3 index.
+        """
+        return F("h3GetBaseCell", index)
+
+    @staticmethod
+    def h3HexAreaM2(resolution):
+        """
+        Returns average hexagon area in square meters at the given resolution.
+        """
+        return F("h3HexAreaM2", resolution)
+
+    @staticmethod
+    def h3HexAreaKm2(resolution):
+        """
+        Returns average hexagon area in square kilometers at the given resolution.
+        """
+        return F("h3HexAreaKm2", resolution)
+
+    @staticmethod
+    def h3IndexesAreNeighbors(index1, index2):
+        """
+        Returns whether or not the provided H3 indexes are neighbors.
+        """
+        return F("h3IndexesAreNeighbors", index1, index2)
+
+    @staticmethod
+    def h3ToChildren(index, resolution):
+        """
+        Returns an array of child indexes for the given H3 index.
+        """
+        return F("h3ToChildren", index, resolution)
+
     @staticmethod
     def indexOf(arr, x):
         return F("indexOf", arr, x)

tests/test_funcs.py

@@ -518,6 +518,45 @@ class FuncsTestCase(TestCaseWithData):
         self._test_func(F.arrayReduce('min', arr), 1)
         self._test_func(F.arrayReverse(arr), [3, 2, 1])
 
+    def test_geo_functions(self):
+        self._test_func(
+            F.greatCircleDistance(55.755831, 37.617673, -55.755831, -37.617673), 14128353
+        )
+        self._test_func(
+            F.geoDistance(55.755831, 37.617673, -55.755831, -37.617673), 14128353
+        )
+        self._test_func(F.greatCircleAngle(0, 0, 1, 0), 1.0)
+        self._test_func(F.pointInPolygon(3, 3, [(6, 0), (8, 4), (5, 8), (0, 2)]), 1)
+        self._test_func(F.geohashDecode("ezs42"), (-5.60302734375, 42.60498046875))
+        self._test_func(F.geohashEncode(-5.60302734375, 42.593994140625, 0), "ezs42d000000")
+        self._test_func(
+            F.geohashesInBox(24.48, 40.56, 24.785, 40.81, 4),
+            ["sx1q", "sx1r", "sx32", "sx1w", "sx1x", "sx38"],
+        )
+        self._test_func(F.h3IsValid(630814730351855103), 1)
+        self._test_func(F.h3GetResolution(639821929606596015), 14)
+        self._test_func(F.h3EdgeAngle(10), 0.0005927224846720883)
+        self._test_func(F.h3EdgeLengthM(15), 0.509713273)
+        self._test_func(F.h3EdgeLengthKm(15), 0.000509713)
+        self._test_func(F.geoToH3(37.79506683, 55.71290588, 15), 644325524701193974)
+        self._test_func(F.h3ToGeo(644325524701193974), (37.79506616830252, 55.71290243145668))
+        self._test_func(F.h3GetBaseCell(612916788725809151), 12)
+        self._test_func(F.h3HexAreaM2(13), 43.9)
+        self._test_func(F.h3HexAreaKm2(13), 0.0000439)
+        self._test_func(F.h3IndexesAreNeighbors(617420388351344639, 617420388352655359), 1)
+        self._test_func(
+            F.h3ToChildren(599405990164561919, 6),
+            [
+                603909588852408319,
+                603909588986626047,
+                603909589120843775,
+                603909589255061503,
+                603909589389279231,
+                603909589523496959,
+                603909589657714687,
+            ],
+        )
+
     def test_split_and_merge_functions(self):
         self._test_func(F.splitByChar('_', 'a_b_c'), ['a', 'b', 'c'])
         self._test_func(F.splitByString('__', 'a__b__c'), ['a', 'b', 'c'])