Preserve empty dataframe shape #321

mymarilyn · Jul 1, 2022 · 4bc8a46 · 4bc8a46
1 parent b51d036
commit 4bc8a46
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/clickhouse_driver/client.py b/clickhouse_driver/client.py
@@ -414,8 +414,9 @@ def query_dataframe(
             settings=settings
         )
 
+        columns = [re.sub(r'\W', '_', name) for name, type_ in columns]
         return pd.DataFrame(
-            {re.sub(r'\W', '_', col[0]): d for d, col in zip(data, columns)}
+            {col: d for d, col in zip(data, columns)}, columns=columns
         )
 
     def insert_dataframe(

diff --git a/tests/numpy/test_generic.py b/tests/numpy/test_generic.py
@@ -141,6 +141,13 @@ def test_insert_not_ordered_columns(self):
             )
             self.assertEqual(rv, n)
 
+    def test_empty_frame_shape(self):
+        df = self.client.query_dataframe(
+            'SELECT number AS a, number AS a FROM system.numbers LIMIT 0'
+        )
+
+        self.assertEqual(df.shape, (0, 2))
+
 
 class NoNumPyTestCase(BaseTestCase):
     def setUp(self):