Test returning data in columns and in np objects

Author: paxcodes

tests/backends/tests.py

@@ -560,6 +560,122 @@ def test_timezone_none_use_tz_false(self):
             connection.init_connection_state()
 
 
+def check_numpy():
+    """Check if numpy is installed."""
+
+    try:
+        import numpy as _
+    except ImportError:
+        return False
+
+    return True
+
+class ColumnarTestCase(TransactionTestCase):
+    available_apps = ["backends"]
+    databases = {"default", "s2r1"}
+
+    def test_columnar_query(self):
+        sql = """
+            SELECT number, number*2, number*3, number*4, number*5 
+            FROM system.numbers 
+            LIMIT 10
+        """
+        with connections["s2r1"].cursor() as cursorWrapper:
+            cursorWrapper.cursor.columnar = True
+            cursorWrapper.execute(sql)
+            self.assertEqual(
+                cursorWrapper.fetchall(),
+                [
+                    (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
+                    (0, 2, 4, 6, 8, 10, 12, 14, 16, 18),
+                    (0, 3, 6, 9, 12, 15, 18, 21, 24, 27),
+                    (0, 4, 8, 12, 16, 20, 24, 28, 32, 36),
+                    (0, 5, 10, 15, 20, 25, 30, 35, 40, 45),
+                ],
+            )
+
+            cursorWrapper.execute(sql)
+            self.assertEqual(
+                cursorWrapper.fetchmany(2),
+                [
+                    (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
+                    (0, 2, 4, 6, 8, 10, 12, 14, 16, 18),
+                ],
+            )
+
+            actual_results = [
+                r
+                for results in iter(lambda: cursorWrapper.fetchmany(2), [])
+                for r in results
+            ]
+            self.assertEqual(
+                actual_results,
+                [
+                    (0, 3, 6, 9, 12, 15, 18, 21, 24, 27),
+                    (0, 4, 8, 12, 16, 20, 24, 28, 32, 36),
+                    (0, 5, 10, 15, 20, 25, 30, 35, 40, 45),
+                ],
+            )
+
+            cursorWrapper.execute(sql)
+            self.assertEqual(
+                cursorWrapper.fetchone(),
+                (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
+            )
+            cursorWrapper.cursor.columnar = False
+    
+
+    @unittest.skipUnless(check_numpy(), "numpy is not installed")
+    def test_use_numpy_query(self):
+        sql = """
+            SELECT toDateTime32('2022-01-01 01:00:05', 'UTC'), number, number*2.5
+            FROM system.numbers 
+            LIMIT 3
+        """
+        import numpy as np
+        
+        with connections["s2r1"].cursor() as cursorWrapper:
+            cursorWrapper.cursor.columnar = True
+            cursorWrapper.cursor.use_numpy = True
+            cursorWrapper.execute(sql)
+            np.testing.assert_equal(
+                cursorWrapper.fetchall(),
+                [
+                    np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
+                    np.array([0, 1, 2], dtype=np.uint64),
+                    np.array([0, 2.5, 5.0], dtype=np.float64)
+                ],
+            )
+
+            cursorWrapper.execute(sql)
+            np.testing.assert_equal(
+                cursorWrapper.fetchmany(2),
+                [
+                    np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
+                    np.array([0, 1, 2], dtype=np.uint64),
+                ],
+            )
+
+            actual_results = [
+                r
+                for results in iter(lambda: cursorWrapper.fetchmany(2), [])
+                for r in results
+            ]
+            np.testing.assert_equal(
+                actual_results,
+                [
+                    np.array([0, 2.5, 5], dtype=np.float64),
+                ],
+            )
+
+            cursorWrapper.execute(sql)
+            np.testing.assert_equal(
+                cursorWrapper.fetchone(),
+                np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
+            )
+            cursorWrapper.cursor.columnar = False
+            cursorWrapper.cursor.use_numpy = False
+
 # These tests aren't conditional because it would require differentiating
 # between MySQL+InnoDB and MySQL+MYISAM (something we currently can't do).
 class FkConstraintsTests(TransactionTestCase):