dont raise on dataframe insert if columns are superset of schema (#340)

* dont raise on dataframe insert if columns are superset of schema
mymarilyn · Oct 23, 2022 · e815179 · e815179
1 parent 3be308c
commit e815179
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/clickhouse_driver/client.py b/clickhouse_driver/client.py
@@ -507,11 +507,11 @@ def insert_dataframe(
             rv = None
             if sample_block:
                 columns = [x[0] for x in sample_block.columns_with_types]
-                if len(columns) != dataframe.shape[1]:
-                    msg = 'Expected {} columns, got {}'.format(
-                        len(columns), dataframe.shape[1]
-                    )
-                    raise ValueError(msg)
+                # raise if any columns are missing from the dataframe
+                diff = set(columns) - set(dataframe.columns)
+                if len(diff):
+                    msg = "DataFrame missing required columns: {}"
+                    raise ValueError(msg.format(list(diff)))
 
                 data = [dataframe[column].values for column in columns]
                 rv = self.send_data(sample_block, data, columnar=True)

diff --git a/tests/numpy/test_generic.py b/tests/numpy/test_generic.py
@@ -148,12 +148,21 @@ def test_empty_frame_shape(self):
 
         self.assertEqual(df.shape, (0, 2))
 
-    def test_data_less_columns_then_expected(self):
+    def test_data_less_columns_than_expected(self):
         with self.create_table('a Int8, b Int8'):
             with self.assertRaises(ValueError) as e:
                 df = pd.DataFrame([1, 2, 3], columns=['a'])
                 self.client.insert_dataframe('INSERT INTO test VALUES', df)
-            self.assertEqual(str(e.exception), 'Expected 2 columns, got 1')
+            expected = "DataFrame missing required columns: ['b']"
+            self.assertEqual(str(e.exception), expected)
+
+    def test_data_different_columns_than_expected(self):
+        with self.create_table('a Int8, b Int8'):
+            with self.assertRaises(ValueError) as e:
+                df = pd.DataFrame([[1, 2], [3, 4]], columns=['a', 'c'])
+                self.client.insert_dataframe('INSERT INTO test VALUES', df)
+            expected = "DataFrame missing required columns: ['b']"
+            self.assertEqual(str(e.exception), expected)
 
 
 class NoNumPyTestCase(BaseTestCase):