Fix NumPy chunking #243

First element was empty and not data was inserted
mymarilyn · Sep 19, 2021 · 015d42e · 015d42e
1 parent ae4b0ac
commit 015d42e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/clickhouse_driver/numpy/helpers.py b/clickhouse_driver/numpy/helpers.py
@@ -13,7 +13,7 @@ def column_chunks(columns, n):
 
     # create chunk generator for every column
     chunked = [
-        iter(np.array_split(c, range(0, len(c), n)) if len(c) > n else [c])
+        iter(np.array_split(c, len(c) // n) if len(c) > n else [c])
         for c in columns
     ]
 

diff --git a/tests/numpy/test_generic.py b/tests/numpy/test_generic.py
@@ -120,6 +120,14 @@ def test_insert_simple(self):
             df2 = self.client.query_dataframe('SELECT * FROM test ORDER BY a')
             self.assertTrue(df.equals(df2))
 
+    def test_insert_chunking(self):
+        with self.create_table('a Int64'):
+            rv = self.client.execute(
+                'INSERT INTO test VALUES', [np.array(range(3))], columnar=True,
+                settings={'insert_block_size': 1}
+            )
+            self.assertEqual(rv, 3)
+
 
 class NoNumPyTestCase(BaseTestCase):
     def setUp(self):