Pyarrow Concat to to know merge with promote options as default (#745)

* Update the Pyarrow concat

* Updated the tests python versions

* Updated github actions

* removed python3.12

* Black formatted

* Fixed tests

* Fixed more tests

Author: jprakash-db

.github/workflows/code-quality-checks.yml

@@ -4,23 +4,22 @@ on:
     branches:
       - main
   pull_request:
-    branches:
-      - main
+
 jobs:
   run-unit-tests:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9, "3.10", "3.11"]
+        python-version: [3.9, "3.10", "3.11"]
     steps:
       #----------------------------------------------
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       #----------------------------------------------
@@ -29,6 +28,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true
@@ -62,16 +62,16 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9, "3.10"]
+        python-version: [3.9, "3.10", "3.11"]
     steps:
       #----------------------------------------------
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       #----------------------------------------------
@@ -80,6 +80,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true
@@ -114,16 +115,16 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9, "3.10"]
+        python-version: [3.9, "3.10", "3.11"]
     steps:
       #----------------------------------------------
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Set up python ${{ matrix.python-version }}
         id: setup-python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       #----------------------------------------------
@@ -132,6 +133,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true

.github/workflows/integration.yml

@@ -20,10 +20,10 @@ jobs:
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up python
         id: setup-python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       #----------------------------------------------
@@ -32,6 +32,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true

.github/workflows/publish-test.yml

@@ -9,10 +9,10 @@ jobs:
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Set up python
         id: setup-python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: 3.9
       #----------------------------------------------
@@ -21,6 +21,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true

.github/workflows/publish.yml

@@ -11,10 +11,10 @@ jobs:
       #       check-out repo and set-up python
       #----------------------------------------------
       - name: Check out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
       - name: Set up python
         id: setup-python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: 3.9
       #----------------------------------------------
@@ -23,6 +23,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
+          version: "2.2.1"
           virtualenvs-create: true
           virtualenvs-in-project: true
           installer-parallel: true

src/databricks/sql/client.py

@@ -1335,7 +1335,9 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
-            results = pyarrow.concat_tables([results, partial_results])
+            results = pyarrow.concat_tables(
+                [results, partial_results], promote_options="default"
+            )
             n_remaining_rows -= partial_results.num_rows
             self._next_row_index += partial_results.num_rows
 
@@ -1391,7 +1393,9 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
-            results = pyarrow.concat_tables([results, partial_results])
+            results = pyarrow.concat_tables(
+                [results, partial_results], promote_options="default"
+            )
             self._next_row_index += partial_results.num_rows
 
         return results

src/databricks/sql/utils.py

@@ -265,7 +265,9 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
             # Get remaining of num_rows or the rest of the current table, whichever is smaller
             length = min(num_rows, self.table.num_rows - self.table_row_index)
             table_slice = self.table.slice(self.table_row_index, length)
-            results = pyarrow.concat_tables([results, table_slice])
+            results = pyarrow.concat_tables(
+                [results, table_slice], promote_options="default"
+            )
             self.table_row_index += table_slice.num_rows
 
             # Replace current table with the next table if we are at the end of the current table
@@ -292,7 +294,9 @@ def remaining_rows(self) -> "pyarrow.Table":
             table_slice = self.table.slice(
                 self.table_row_index, self.table.num_rows - self.table_row_index
             )
-            results = pyarrow.concat_tables([results, table_slice])
+            results = pyarrow.concat_tables(
+                [results, table_slice], promote_options="default"
+            )
             self.table_row_index += table_slice.num_rows
             self.table = self._create_next_table()
             self.table_row_index = 0

tests/e2e/common/large_queries_mixin.py

@@ -83,11 +83,11 @@ def test_query_with_large_narrow_result_set(self):
                 assert row[0] == row_id
 
     def test_long_running_query(self):
-        """Incrementally increase query size until it takes at least 4 minutes,
+        """Incrementally increase query size until it takes at least 2 minutes,
         and asserts that the query completes successfully.
         """
         minutes = 60
-        min_duration = 4 * minutes
+        min_duration = 2 * minutes
 
         duration = -1
         scale0 = 10000
@@ -113,5 +113,5 @@ def test_long_running_query(self):
                 duration = time.time() - start
                 current_fraction = duration / min_duration
                 print("Took {} s with scale factor={}".format(duration, scale_factor))
-                # Extrapolate linearly to reach 4 min and add 50% padding to push over the limit
+                # Extrapolate linearly to reach 2 min and add 50% padding to push over the limit
                 scale_factor = math.ceil(1.5 * scale_factor / current_fraction)

tests/e2e/common/staging_ingestion_tests.py

@@ -46,7 +46,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
         ) as conn:
 
             cursor = conn.cursor()
-            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
+            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/13/file1.csv' OVERWRITE"
             cursor.execute(query)
 
         # GET should succeed
@@ -57,7 +57,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
             extra_params={"staging_allowed_local_path": new_temp_path}
         ) as conn:
             cursor = conn.cursor()
-            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/13/file1.csv' TO '{new_temp_path}'"
             cursor.execute(query)
 
         with open(new_fh, "rb") as fp:
@@ -67,7 +67,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
 
         # REMOVE should succeed
 
-        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv'"
+        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/13/file1.csv'"
 
         with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
             cursor = conn.cursor()
@@ -79,7 +79,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
                 Error, match="Staging operation over HTTP was unsuccessful: 404"
             ):
                 cursor = conn.cursor()
-                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/13/file1.csv' TO '{new_temp_path}'"
                 cursor.execute(query)
 
         os.remove(temp_path)

tests/e2e/test_complex_types.py

@@ -14,7 +14,7 @@ def table_fixture(self, connection_details):
             # Create the table
             cursor.execute(
                 """
-                CREATE TABLE IF NOT EXISTS pysql_test_complex_types_table (
+                CREATE TABLE IF NOT EXISTS pysql_test_complex_types_table_deprecated (
                     array_col ARRAY<STRING>,
                     map_col MAP<STRING, INTEGER>,
                     struct_col STRUCT<field1: STRING, field2: INTEGER>
@@ -24,7 +24,7 @@ def table_fixture(self, connection_details):
             # Insert a record
             cursor.execute(
                 """
-                INSERT INTO pysql_test_complex_types_table
+                INSERT INTO pysql_test_complex_types_table_deprecated
                 VALUES (
                     ARRAY('a', 'b', 'c'),
                     MAP('a', 1, 'b', 2, 'c', 3),
@@ -34,7 +34,7 @@ def table_fixture(self, connection_details):
             )
             yield
             # Clean up the table after the test
-            cursor.execute("DROP TABLE IF EXISTS pysql_test_complex_types_table")
+            cursor.execute("DROP TABLE IF EXISTS pysql_test_complex_types_table_deprecated")
 
     @pytest.mark.parametrize(
         "field,expected_type",
@@ -45,7 +45,7 @@ def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
 
         with self.cursor() as cursor:
             result = cursor.execute(
-                "SELECT * FROM pysql_test_complex_types_table LIMIT 1"
+                "SELECT * FROM pysql_test_complex_types_table_deprecated LIMIT 1"
             ).fetchone()
 
         assert isinstance(result[field], expected_type)
@@ -57,7 +57,7 @@ def test_read_complex_types_as_string(self, field, table_fixture):
             extra_params={"_use_arrow_native_complex_types": False}
         ) as cursor:
             result = cursor.execute(
-                "SELECT * FROM pysql_test_complex_types_table LIMIT 1"
+                "SELECT * FROM pysql_test_complex_types_table_deprecated LIMIT 1"
             ).fetchone()
 
         assert isinstance(result[field], str)

tests/e2e/test_parameterized_queries.py

@@ -123,7 +123,7 @@ def inline_table(self, connection_details):
         """
 
         query = """
-            CREATE TABLE IF NOT EXISTS pysql_e2e_inline_param_test_table (
+            CREATE TABLE IF NOT EXISTS pysql_e2e_inline_param_test_table_deprecated (
             null_col INT,
             int_col INT,
             bigint_col BIGINT,
@@ -167,9 +167,9 @@ def _inline_roundtrip(self, params: dict, paramstyle: ParamStyle):
             This is a no-op but is included to make the test-code easier to read.
         """
         target_column = self._get_inline_table_column(params.get("p"))
-        INSERT_QUERY = f"INSERT INTO pysql_e2e_inline_param_test_table (`{target_column}`) VALUES (%(p)s)"
-        SELECT_QUERY = f"SELECT {target_column} `col` FROM pysql_e2e_inline_param_test_table LIMIT 1"
-        DELETE_QUERY = "DELETE FROM pysql_e2e_inline_param_test_table"
+        INSERT_QUERY = f"INSERT INTO pysql_e2e_inline_param_test_table_deprecated (`{target_column}`) VALUES (%(p)s)"
+        SELECT_QUERY = f"SELECT {target_column} `col` FROM pysql_e2e_inline_param_test_table_deprecated LIMIT 1"
+        DELETE_QUERY = "DELETE FROM pysql_e2e_inline_param_test_table_deprecated"
 
         with self.connection(extra_params={"use_inline_params": True}) as conn:
             with conn.cursor() as cursor: