Always fallback to generic columns #254

mymarilyn · Oct 28, 2021 · ffa6c38 · ffa6c38
1 parent 59504c1
commit ffa6c38
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 54 deletions.
diff --git a/clickhouse_driver/columns/numpy/service.py b/clickhouse_driver/columns/numpy/service.py
@@ -1,40 +1,20 @@
 from ... import errors
-from ..arraycolumn import create_array_column
 from .datecolumn import NumpyDateColumn
 from .datetimecolumn import create_numpy_datetime_column
-from ..decimalcolumn import create_decimal_column
-from ..enumcolumn import create_enum_column
 from .floatcolumn import NumpyFloat32Column, NumpyFloat64Column
 from .intcolumn import (
     NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
     NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column
 )
 from .lowcardinalitycolumn import create_numpy_low_cardinality_column
-from ..nothingcolumn import NothingColumn
-from ..nullcolumn import NullColumn
-from ..nullablecolumn import create_nullable_column
-from ..simpleaggregatefunctioncolumn import (
-    create_simple_aggregate_function_column
-)
 from .stringcolumn import create_string_column
-from ..tuplecolumn import create_tuple_column
-from ..uuidcolumn import UUIDColumn
-from ..intervalcolumn import (
-    IntervalYearColumn, IntervalMonthColumn, IntervalWeekColumn,
-    IntervalDayColumn, IntervalHourColumn, IntervalMinuteColumn,
-    IntervalSecondColumn
-)
-from ..ipcolumn import IPv4Column, IPv6Column
+from ..nullablecolumn import create_nullable_column
 
 column_by_type = {c.ch_type: c for c in [
     NumpyDateColumn,
     NumpyFloat32Column, NumpyFloat64Column,
     NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
-    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column,
-    NothingColumn, NullColumn, UUIDColumn,
-    IntervalYearColumn, IntervalMonthColumn, IntervalWeekColumn,
-    IntervalDayColumn, IntervalHourColumn, IntervalMinuteColumn,
-    IntervalSecondColumn, IPv4Column, IPv6Column
+    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column
 ]}
 
 
@@ -45,36 +25,18 @@ def create_column_with_options(x):
     if spec == 'String' or spec.startswith('FixedString'):
         return create_string_column(spec, column_options)
 
-    elif spec.startswith('Enum'):
-        return create_enum_column(spec, column_options)
-
     elif spec.startswith('DateTime'):
         return create_numpy_datetime_column(spec, column_options)
 
-    elif spec.startswith('Decimal'):
-        return create_decimal_column(spec, column_options)
-
-    elif spec.startswith('Array'):
-        return create_array_column(spec, create_column_with_options)
-
-    elif spec.startswith('Tuple'):
-        return create_tuple_column(spec, create_column_with_options)
-
     elif spec.startswith('Nullable'):
         return create_nullable_column(spec, create_column_with_options)
 
     elif spec.startswith('LowCardinality'):
         return create_numpy_low_cardinality_column(spec,
                                                    create_column_with_options)
-
-    elif spec.startswith('SimpleAggregateFunction'):
-        return create_simple_aggregate_function_column(
-            spec, create_column_with_options)
-
     else:
-        try:
+        if spec in column_by_type:
             cls = column_by_type[spec]
             return cls(**column_options)
 
-        except KeyError as e:
-            raise errors.UnknownTypeError('Unknown type {}'.format(e.args[0]))
+        raise errors.UnknownTypeError('Unknown type {}'.format(spec))
diff --git a/clickhouse_driver/columns/service.py b/clickhouse_driver/columns/service.py
@@ -48,7 +48,12 @@ def get_column_by_spec(spec, column_options):
 
     if use_numpy:
         from .numpy.service import get_numpy_column_by_spec
-        return get_numpy_column_by_spec(spec, column_options)
+
+        try:
+            return get_numpy_column_by_spec(spec, column_options)
+        except errors.UnknownTypeError:
+            # falling back to generic columns
+            pass
 
     def create_column_with_options(x):
         return get_column_by_spec(x, column_options)
@@ -89,8 +94,8 @@ def create_column_with_options(x):
             cls = column_by_type[spec]
             return cls(**column_options)
 
-        except KeyError as e:
-            raise errors.UnknownTypeError('Unknown type {}'.format(e.args[0]))
+        except KeyError:
+            raise errors.UnknownTypeError('Unknown type {}'.format(spec))
 
 
 def read_column(context, column_spec, n_items, buf):

diff --git a/tests/numpy/columns/test_other.py b/tests/numpy/columns/test_other.py
@@ -1,11 +1,5 @@
 from clickhouse_driver import errors
-
-try:
-    from clickhouse_driver.columns.numpy.service import \
-        get_numpy_column_by_spec
-except ImportError:
-    get_numpy_column_by_spec = None
-
+from clickhouse_driver.columns.service import get_column_by_spec
 from clickhouse_driver.context import Context
 
 from tests.numpy.testcase import NumpyBaseTestCase
@@ -14,8 +8,8 @@
 class OtherColumnsTestCase(NumpyBaseTestCase):
     def get_column(self, spec):
         ctx = Context()
-        ctx.client_settings = {'strings_as_bytes': False}
-        return get_numpy_column_by_spec(spec, {'context': ctx})
+        ctx.client_settings = {'strings_as_bytes': False, 'use_numpy': True}
+        return get_column_by_spec(spec, {'context': ctx})
 
     def test_enum(self):
         col = self.get_column("Enum8('hello' = 1, 'world' = 2)")
@@ -37,6 +31,10 @@ def test_simple_aggregation_function(self):
         col = self.get_column('SimpleAggregateFunction(any, Int32)')
         self.assertIsNotNone(col)
 
+    def test_map(self):
+        col = self.get_column('Map(String, String)')
+        self.assertIsNotNone(col)
+
     def test_get_unknown_column(self):
         with self.assertRaises(errors.UnknownTypeError) as e:
             self.get_column('Unicorn')