SNOW-105364: added PYTHON_CONNECTOR_QUERY_RESULT_FORMAT

Author: sfc-gh-stakeda

scripts/install.bat

@@ -7,7 +7,7 @@ SET SCRIPT_DIR=%~dp0
 "%PYTHON%/python.exe" -m venv env
 call env\Scripts\activate
 # https://github.com/pypa/pip/issues/6566
-python -m pip install --upgrade pip==18.1
+python -m pip install --upgrade pip
 pip install pandas
 pip install numpy
 pip install pendulum
@@ -25,4 +25,4 @@ set /p connector_whl=<whl_name
 pip install %connector_whl%
 pip list --format=columns
 
-cd %SCRIPT_DIR%/..
+cd %SCRIPT_DIR%/..

test/test_arrow_pandas.py

@@ -24,6 +24,7 @@
     no_arrow_iterator_ext = True
 
 SQL_ENABLE_ARROW = "alter session set query_result_format='ARROW_FORCE';"
+SQL_ENABLE_ARROW_V2 = "alter session set python_connector_query_result_format='ARROW_FORCE';"
 
 EPSILON = 1e-8
 
@@ -229,6 +230,7 @@ def test_semi_struct(conn_cnx):
         # fetch dataframe with new arrow support
         cursor_table = cnx_table.cursor()
         cursor_table.execute(SQL_ENABLE_ARROW)
+        cursor_table.execute(SQL_ENABLE_ARROW_V2)
         cursor_table.execute(sql_text)
         df_new = cursor_table.fetch_pandas_all()
         col_new = df_new.iloc[0]
@@ -430,6 +432,7 @@ def validate_pandas(conn_cnx, sql, cases, col_count, method='one', data_type='fl
         # fetch dataframe with new arrow support
         cursor_table = cnx_table.cursor()
         cursor_table.execute(SQL_ENABLE_ARROW)
+        cursor_table.execute(SQL_ENABLE_ARROW_V2)
         cursor_table.execute(sql)
 
         # build dataframe
@@ -524,6 +527,7 @@ def test_empty(conn_cnx):
     with conn_cnx() as cnx:
         cursor = cnx.cursor()
         cursor.execute(SQL_ENABLE_ARROW)
+        cursor.execute(SQL_ENABLE_ARROW_V2)
         cursor.execute("select seq4() from table(generator(rowcount=>1)) limit 0")
         assert cursor.fetch_pandas_all() is None, 'the result is not none'
 
@@ -553,6 +557,7 @@ def fetch_pandas(conn_cnx, sql, row_count, col_count, method='one'):
             # fetch dataframe by fetching row by row
             cursor_row = cnx_row.cursor()
             cursor_row.execute(SQL_ENABLE_ARROW)
+            cursor_row.execute(SQL_ENABLE_ARROW_V2)
             cursor_row.execute(sql)
 
             # build dataframe
@@ -578,6 +583,7 @@ def fetch_pandas(conn_cnx, sql, row_count, col_count, method='one'):
             # fetch dataframe with new arrow support
             cursor_table = cnx_table.cursor()
             cursor_table.execute(SQL_ENABLE_ARROW)
+            cursor_table.execute(SQL_ENABLE_ARROW_V2)
             cursor_table.execute(sql)
 
             # build dataframe

test/test_arrow_result.py

@@ -413,6 +413,7 @@ def test_select_with_empty_resultset(conn_cnx):
     with conn_cnx() as cnx:
         cursor = cnx.cursor()
         cursor.execute("alter session set query_result_format='ARROW_FORCE'")
+        cursor.execute("alter session set python_connector_query_result_format='ARROW_FORCE'")
         cursor.execute("select seq4() from table(generator(rowcount=>100)) limit 0")
 
         assert cursor.fetchone() is None
@@ -448,10 +449,12 @@ def iterate_over_test_chunk(test_name, conn_cnx, sql_text, row_count, col_count,
             if expected is None:
                 cursor_json = json_cnx.cursor()
                 cursor_json.execute("alter session set query_result_format='JSON'")
+                cursor_json.execute("alter session set python_connector_query_result_format='JSON'")
                 cursor_json.execute(sql_text)
 
             cursor_arrow = arrow_cnx.cursor()
             cursor_arrow.execute("alter session set query_result_format='ARROW_FORCE'")
+            cursor_arrow.execute("alter session set python_connector_query_result_format='ARROW_FORCE'")
             cursor_arrow.execute(sql_text)
 
             if expected is None: