From 9d3ac8cb39725d6bdc836c9003b903c81fac7a28 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 11:30:45 +0530 Subject: [PATCH 1/7] Update the Pyarrow concat --- src/databricks/sql/client.py | 4 ++-- src/databricks/sql/utils.py | 4 ++-- tests/unit/test_cloud_fetch_queue.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 7e3aa02a9..22700d591 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1335,7 +1335,7 @@ 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 +1391,7 @@ 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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index cd655c4ea..b7b04a944 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -265,7 +265,7 @@ 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 +292,7 @@ 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 diff --git a/tests/unit/test_cloud_fetch_queue.py b/tests/unit/test_cloud_fetch_queue.py index 01d8a79b9..7f19f5a57 100644 --- a/tests/unit/test_cloud_fetch_queue.py +++ b/tests/unit/test_cloud_fetch_queue.py @@ -185,7 +185,7 @@ def test_next_n_rows_more_than_one_table(self, mock_create_next_table): assert ( result == pyarrow.concat_tables( - [self.make_arrow_table(), self.make_arrow_table()] + [self.make_arrow_table(), self.make_arrow_table()],promote_options="default" )[:7] ) @@ -309,7 +309,7 @@ def test_remaining_rows_multiple_tables_fully_returned( assert ( result == pyarrow.concat_tables( - [self.make_arrow_table(), self.make_arrow_table()] + [self.make_arrow_table(), self.make_arrow_table()],promote_options="default" )[3:] ) From 621b627c2c85d38d08388d672bc56c6fecad8787 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 11:39:09 +0530 Subject: [PATCH 2/7] Updated the tests python versions --- .github/workflows/code-quality-checks.yml | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index b4d4776e7..38081796c 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -4,14 +4,13 @@ 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", "3.12"] steps: #---------------------------------------------- # check-out repo and set-up python @@ -62,7 +61,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.8, 3.9, "3.10"] + python-version: [3.9, "3.10", "3.11", "3.12"] steps: #---------------------------------------------- # check-out repo and set-up python @@ -114,7 +113,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.8, 3.9, "3.10"] + python-version: [3.9, "3.10", "3.11", "3.12"] steps: #---------------------------------------------- # check-out repo and set-up python From 4911bf1de718382fdd44cf7120e3e2e0219fcde3 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 11:56:27 +0530 Subject: [PATCH 3/7] Updated github actions --- .github/workflows/code-quality-checks.yml | 15 +++++++++------ .github/workflows/integration.yml | 5 +++-- .github/workflows/publish-test.yml | 5 +++-- .github/workflows/publish.yml | 5 +++-- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index 38081796c..1ae590dde 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -16,10 +16,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 ${{ matrix.python-version }} id: setup-python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} #---------------------------------------------- @@ -28,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 @@ -67,10 +68,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 ${{ matrix.python-version }} id: setup-python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} #---------------------------------------------- @@ -79,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 @@ -119,10 +121,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 ${{ matrix.python-version }} id: setup-python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} #---------------------------------------------- @@ -131,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 diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index c0122f255..4d10464bd 100644 --- a/.github/workflows/integration.yml +++ b/.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 diff --git a/.github/workflows/publish-test.yml b/.github/workflows/publish-test.yml index 203c6571d..b7ffee9f4 100644 --- a/.github/workflows/publish-test.yml +++ b/.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 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 7960e9ed0..c592756b8 100644 --- a/.github/workflows/publish.yml +++ b/.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 From ef48d3e68b3887c638cd849b3c23d9499e21dc1f Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 12:02:32 +0530 Subject: [PATCH 4/7] removed python3.12 --- .github/workflows/code-quality-checks.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index 1ae590dde..946374d1f 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -10,7 +10,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.9, "3.10", "3.11", "3.12"] + python-version: [3.9, "3.10", "3.11"] steps: #---------------------------------------------- # check-out repo and set-up python @@ -62,7 +62,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.9, "3.10", "3.11", "3.12"] + python-version: [3.9, "3.10", "3.11"] steps: #---------------------------------------------- # check-out repo and set-up python @@ -115,7 +115,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.9, "3.10", "3.11", "3.12"] + python-version: [3.9, "3.10", "3.11"] steps: #---------------------------------------------- # check-out repo and set-up python From 8457cb917e879bee2b317d140945b62af7196f5b Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 12:06:20 +0530 Subject: [PATCH 5/7] Black formatted --- src/databricks/sql/client.py | 8 ++++++-- src/databricks/sql/utils.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 22700d591..59c74c6ba 100755 --- a/src/databricks/sql/client.py +++ b/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],promote_options="default") + 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],promote_options="default") + results = pyarrow.concat_tables( + [results, partial_results], promote_options="default" + ) self._next_row_index += partial_results.num_rows return results diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index b7b04a944..48e0aa77e 100644 --- a/src/databricks/sql/utils.py +++ b/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],promote_options="default") + 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],promote_options="default") + 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 From e067f2dcd8b177482d47412959c810855803983d Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 12:57:08 +0530 Subject: [PATCH 6/7] Fixed tests --- tests/e2e/test_complex_types.py | 10 +++++----- tests/e2e/test_parameterized_queries.py | 8 ++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/e2e/test_complex_types.py b/tests/e2e/test_complex_types.py index 446a6b50a..77d634242 100644 --- a/tests/e2e/test_complex_types.py +++ b/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, map_col MAP, struct_col STRUCT @@ -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) diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py index d346ad5c6..ddf8b33fd 100644 --- a/tests/e2e/test_parameterized_queries.py +++ b/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: From 833ea1c610c1e1ea590162392cedc05636d550aa Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 2 Mar 2026 14:25:37 +0530 Subject: [PATCH 7/7] Fixed more tests --- tests/e2e/common/large_queries_mixin.py | 6 +++--- tests/e2e/common/staging_ingestion_tests.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/e2e/common/large_queries_mixin.py b/tests/e2e/common/large_queries_mixin.py index 3a35d8014..5fe292d91 100644 --- a/tests/e2e/common/large_queries_mixin.py +++ b/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) diff --git a/tests/e2e/common/staging_ingestion_tests.py b/tests/e2e/common/staging_ingestion_tests.py index 008055e33..39aea3feb 100644 --- a/tests/e2e/common/staging_ingestion_tests.py +++ b/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)