run large queries

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

tests/e2e/common/large_queries_mixin.py

@@ -2,6 +2,8 @@
 import math
 import time
 
+import pytest
+
 log = logging.getLogger(__name__)
 
 
@@ -42,7 +44,14 @@ def fetch_rows(self, cursor, row_count, fetchmany_size):
             + "assuming 10K fetch size."
         )
 
-    def test_query_with_large_wide_result_set(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_query_with_large_wide_result_set(self, extra_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8192  # B
         rows = resultSize // width
@@ -52,7 +61,7 @@ def test_query_with_large_wide_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 1000
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             for lz4_compression in [False, True]:
                 cursor.connection.lz4_compression = lz4_compression
                 uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
@@ -68,7 +77,14 @@ def test_query_with_large_wide_result_set(self):
                     assert row[0] == row_id  # Verify no rows are dropped in the middle.
                     assert len(row[1]) == 36
 
-    def test_query_with_large_narrow_result_set(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_query_with_large_narrow_result_set(self, extra_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8  # sizeof(long)
         rows = resultSize / width
@@ -77,12 +93,19 @@ def test_query_with_large_narrow_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 10000000
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             cursor.execute("SELECT * FROM RANGE({rows})".format(rows=rows))
             for row_id, row in enumerate(self.fetch_rows(cursor, rows, fetchmany_size)):
                 assert row[0] == row_id
 
-    def test_long_running_query(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_long_running_query(self, extra_params):
         """Incrementally increase query size until it takes at least 3 minutes,
         and asserts that the query completes successfully.
         """
@@ -92,7 +115,7 @@ def test_long_running_query(self):
         duration = -1
         scale0 = 10000
         scale_factor = 1
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             while duration < min_duration:
                 assert scale_factor < 1024, "Detected infinite loop"
                 start = time.time()