expand geometric types

CHANGELOG.md

@@ -12,9 +12,13 @@ instead of being passed as ClickHouse server settings. This is in conjunction wi
 The supported method of passing ClickHouse server settings is to prefix such arguments/query parameters with`ch_`.
 
 ## 0.8.7, 2024-11-21
+### Improvement
+- Added basic support for ClickHouse geometric types Ring, Polygon, MultiPolygon, LineString, and MultiLineString.
+Closed 
+- 
 ### Bug Fix
 - Settings/parameters from one Client will no longer leak into later client instantiations.  Fixes
-https://github.com/ClickHouse/clickhouse-connect/issues/426
+https://github.com/ClickHouse/clickhouse-connect/issues/427
 
 ## 0.8.6, 2024-11-01
 ### Bug Fixes

clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py

@@ -206,6 +206,26 @@ class Point(ChSqlaType, UserDefinedType):
     python_type = None
 
 
+class Ring(ChSqlaType, UserDefinedType):
+    python_type = None
+
+
+class Polygon(ChSqlaType, UserDefinedType):
+    python_type = None
+
+
+class MultiPolygon(ChSqlaType, UserDefinedType):
+    python_type = None
+
+
+class LineString(ChSqlaType, UserDefinedType):
+    python_type = None
+
+
+class MultiLineString(ChSqlaType, UserDefinedType):
+    python_type = None
+
+
 class Date(ChSqlaType, SqlaDate):
     pass
 

clickhouse_connect/datatypes/__init__.py

@@ -4,6 +4,7 @@
 import clickhouse_connect.datatypes.special
 import clickhouse_connect.datatypes.string
 import clickhouse_connect.datatypes.temporal
+import clickhouse_connect.datatypes.geometric
 import clickhouse_connect.datatypes.dynamic
 import clickhouse_connect.datatypes.registry
 import clickhouse_connect.datatypes.postinit

clickhouse_connect/datatypes/container.py

@@ -155,13 +155,6 @@ def convert_dict_insert(self, column: Sequence) -> Sequence:
         return col
 
 
-class Point(Tuple):
-
-    def __init__(self, type_def):
-        super().__init__(type_def)
-        self._name_suffix = ''
-
-
 class Map(ClickHouseType):
     _slots = 'key_type', 'value_type'
     python_type = dict

clickhouse_connect/datatypes/geometric.py

@@ -0,0 +1,51 @@
+from typing import Sequence
+
+from clickhouse_connect.datatypes.base import ClickHouseType
+from clickhouse_connect.driver.insert import InsertContext
+from clickhouse_connect.driver.query import QueryContext
+from clickhouse_connect.driver.types import ByteSource
+
+POINT_DATA_TYPE: ClickHouseType
+RING_DATA_TYPE: ClickHouseType
+POLYGON_DATA_TYPE: ClickHouseType
+MULTI_POLYGON_DATA_TYPE: ClickHouseType
+
+
+class Point(ClickHouseType):
+    def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext):
+        return POINT_DATA_TYPE.write_column(column, dest, ctx)
+
+    def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext) -> Sequence:
+        return POINT_DATA_TYPE.read_column_data(source, num_rows, ctx)
+
+
+class Ring(ClickHouseType):
+    def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext):
+        return RING_DATA_TYPE.write_column(column, dest, ctx)
+
+    def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext) -> Sequence:
+        return RING_DATA_TYPE.read_column_data(source, num_rows, ctx)
+
+
+class Polygon(ClickHouseType):
+    def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext):
+        return POLYGON_DATA_TYPE.write_column(column, dest, ctx)
+
+    def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext) -> Sequence:
+        return POLYGON_DATA_TYPE.read_column_data(source, num_rows, ctx)
+
+
+class MultiPolygon(ClickHouseType):
+    def write_column(self, column: Sequence, dest: bytearray, ctx: InsertContext):
+        return MULTI_POLYGON_DATA_TYPE.write_column(column, dest, ctx)
+
+    def read_column_data(self, source: ByteSource, num_rows: int, ctx: QueryContext) -> Sequence:
+        return MULTI_POLYGON_DATA_TYPE.read_column_data(source, num_rows, ctx)
+
+
+class LineString(Ring):
+    pass
+
+
+class MultiLineString(Polygon):
+    pass

clickhouse_connect/datatypes/postinit.py

@@ -1,4 +1,14 @@
-from clickhouse_connect.datatypes import registry, dynamic
+from clickhouse_connect.datatypes import registry, dynamic, geometric
 
 dynamic.SHARED_DATA_TYPE = registry.get_from_name('Array(String, String)')
 dynamic.STRING_DATA_TYPE = registry.get_from_name('String')
+
+point = 'Tuple(Float64, Float64)'
+ring = f'Array({point})'
+polygon = f'Array({ring})'
+multi_polygon = f'Array({polygon})'
+
+geometric.POINT_DATA_TYPE = registry.get_from_name(point)
+geometric.RING_DATA_TYPE = registry.get_from_name(ring)
+geometric.POLYGON_DATA_TYPE = registry.get_from_name(polygon)
+geometric.MULTI_POLYGON_DATA_TYPE = registry.get_from_name(multi_polygon)

tests/integration_tests/test_geometric.py

@@ -0,0 +1,34 @@
+from typing import Callable
+
+from clickhouse_connect.driver import Client
+
+
+def test_point_column(test_client: Client, table_context: Callable):
+    with table_context('point_column_test', ['key Int32', 'point Point']):
+        data = [[1, (3.55, 3.55)], [2, (4.55, 4.55)]]
+        test_client.insert('point_column_test', data)
+
+        query_result = test_client.query('SELECT * FROM point_column_test ORDER BY key').result_rows
+        assert len(query_result) == 2
+        assert query_result[0] == (1, (3.55, 3.55))
+        assert query_result[1] == (2, (4.55, 4.55))
+
+
+def test_ring_column(test_client: Client, table_context: Callable):
+    with table_context('ring_column_test', ['key Int32', 'ring Ring']):
+        data = [[1, [(5.522, 58.472),(3.55, 3.55)]], [2, [(4.55, 4.55)]]]
+        test_client.insert('ring_column_test', data)
+
+        query_result = test_client.query('SELECT * FROM ring_column_test ORDER BY key').result_rows
+        assert len(query_result) == 2
+        assert query_result[0] == (1, [(5.522, 58.472),(3.55, 3.55)])
+        assert query_result[1] == (2, [(4.55, 4.55)])
+
+
+def test_polygon_column(test_client: Client, table_context: Callable):
+    with table_context('polygon_column_test', ['key Int32', 'polygon Polygon']):
+        res = test_client.query("SELECT readWKTPolygon('POLYGON ((-64.8 32.3, -65.5 18.3, -80.3 25.2, -64.8 32.3))') as polygon")
+        pg = res.first_row[0]
+        test_client.insert('polygon_column_test', [(1, pg), (4, pg)])
+        query_result = test_client.query('SELECT key, polygon FROM polygon_column_test WHERE key = 4')
+        assert query_result.first_row[1] == pg

tests/integration_tests/test_native.py

@@ -177,17 +177,6 @@ def test_tuple_inserts(test_client: Client, table_context: Callable):
         assert query_result[2] == query_result[3]
 
 
-def test_point_inserts(test_client: Client, table_context: Callable):
-    with table_context('insert_point_test', ['key Int32', 'point Point']):
-        data = [[1, (3.55, 3.55)], [2, (4.55, 4.55)]]
-        test_client.insert('insert_point_test', data)
-
-        query_result = test_client.query('SELECT * FROM insert_point_test ORDER BY key').result_rows
-        assert len(query_result) == 2
-        assert query_result[0] == (1, (3.55, 3.55))
-        assert query_result[1] == (2, (4.55, 4.55))
-
-
 def test_agg_function(test_client: Client, table_context: Callable):
     with table_context('agg_func_test', ['key Int32',
                                          'str SimpleAggregateFunction(any, String)',