Fix nanoseconds in DateTime64 NumPy INSERTs #307

mymarilyn · Jul 10, 2022 · 4bbdef9 · 4bbdef9
1 parent 978d6ac
commit 4bbdef9
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 3 deletions.
diff --git a/clickhouse_driver/columns/numpy/datetimecolumn.py b/clickhouse_driver/columns/numpy/datetimecolumn.py
@@ -68,7 +68,7 @@ class NumpyDateTime64Column(NumpyDateTimeColumnBase):
     dtype = np.dtype(np.uint64)
     datetime_dtype = 'datetime64[ns]'
 
-    max_scale = 6
+    max_scale = 9
 
     def __init__(self, scale=0, **kwargs):
         self.scale = scale
@@ -81,7 +81,7 @@ def read_items(self, n_items, buf):
         items = super(NumpyDateTime64Column, self).read_items(n_items, buf)
 
         seconds = (items // scale).astype('datetime64[s]')
-        microseconds = ((items % scale) * frac_scale).astype('timedelta64[us]')
+        microseconds = ((items % scale) * frac_scale).astype('timedelta64[ns]')
 
         dt = seconds + microseconds
         return self.apply_timezones_after_read(dt)
@@ -98,7 +98,7 @@ def write_items(self, items, buf):
         items = self.apply_timezones_before_write(items)
 
         seconds = items.astype('datetime64[s]')
-        microseconds = (items - seconds).astype(dtype='timedelta64[us]') \
+        microseconds = (items - seconds).astype(dtype='timedelta64[ns]') \
             .astype(np.uint32) // frac_scale
 
         items = seconds.astype(self.dtype) * scale + microseconds

diff --git a/tests/numpy/columns/test_datetime.py b/tests/numpy/columns/test_datetime.py
@@ -147,6 +147,50 @@ def test_datetime64_explicit_frac(self):
                 inserted[0], self.make_numpy_d64ns(['2012-10-25T14:07:19.1'])
             )
 
+    @require_server_version(20, 1, 2)
+    def test_datetime64_nanosecond_precision(self):
+        with self.create_table('a DateTime64(8)'):
+            data = [self.make_numpy_d64ns([
+                '2012-10-25T14:07:19.12345678',
+                '2012-10-25T14:07:19.99999999',
+            ])]
+            self.client.execute(
+                'INSERT INTO test (a) VALUES', data, columnar=True
+            )
+
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                '2012-10-25 14:07:19.12345678\n'
+                '2012-10-25 14:07:19.99999999\n'
+            )
+
+            inserted = self.client.execute(query, columnar=True)
+            self.assertArraysEqual(inserted[0], data[0])
+
+    @require_server_version(20, 1, 2)
+    def test_datetime64_max_precision(self):
+        with self.create_table('a DateTime64(9)'):
+            data = [self.make_numpy_d64ns([
+                '2012-10-25T14:07:19.123456789',
+                '2012-10-25T14:07:19.999999999',
+            ])]
+            self.client.execute(
+                'INSERT INTO test (a) VALUES', data, columnar=True
+            )
+
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                '2012-10-25 14:07:19.123456789\n'
+                '2012-10-25 14:07:19.999999999\n'
+            )
+
+            inserted = self.client.execute(query, columnar=True)
+            self.assertArraysEqual(inserted[0], data[0])
+
     def test_insert_integers_datetime(self):
         with self.create_table('a DateTime'):
             self.client.execute(