SNOW-966003: Fix Arrow return value for zero-row queries (#1832)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 
+- v3.7.0(TBD)
+
+  - Added a new boolean parameter `force_return_table` to `SnowflakeCursor.fetch_arrow_all` to force returning `pyarrow.Table` in case of zero rows.
 
 - v3.6.0(December 09,2023)
 

src/snowflake/connector/cursor.py

@@ -1333,14 +1333,29 @@ def fetch_arrow_batches(self) -> Iterator[Table]:
         )
         return self._result_set._fetch_arrow_batches()
 
-    def fetch_arrow_all(self) -> Table | None:
+    @overload
+    def fetch_arrow_all(self, force_return_table: Literal[False]) -> Table | None:
+        ...
+
+    @overload
+    def fetch_arrow_all(self, force_return_table: Literal[True]) -> Table:
+        ...
+
+    def fetch_arrow_all(self, force_return_table: bool = False) -> Table | None:
+        """
+        Args:
+            force_return_table: Set to True so that when the query returns zero rows,
+                an empty pyarrow table will be returned with schema using the highest bit length for each column.
+                Default value is False in which case None is returned in case of zero rows.
+        """
         self.check_can_use_arrow_resultset()
+
         if self._prefetch_hook is not None:
             self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
         self._log_telemetry_job_data(TelemetryField.ARROW_FETCH_ALL, TelemetryData.TRUE)
-        return self._result_set._fetch_arrow_all()
+        return self._result_set._fetch_arrow_all(force_return_table=force_return_table)
 
     def fetch_pandas_batches(self, **kwargs: Any) -> Iterator[DataFrame]:
         """Fetches a single Arrow Table."""

src/snowflake/connector/result_set.py

@@ -9,7 +9,16 @@
 from concurrent.futures import Future
 from concurrent.futures.thread import ThreadPoolExecutor
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Callable, Deque, Iterable, Iterator
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Deque,
+    Iterable,
+    Iterator,
+    Literal,
+    overload,
+)
 
 from .constants import IterUnit
 from .errors import NotSupportedError
@@ -164,13 +173,21 @@ def _fetch_arrow_batches(
         self._can_create_arrow_iter()
         return self._create_iter(iter_unit=IterUnit.TABLE_UNIT, structure="arrow")
 
-    def _fetch_arrow_all(self) -> Table | None:
+    @overload
+    def _fetch_arrow_all(self, force_return_table: Literal[False]) -> Table | None:
+        ...
+
+    @overload
+    def _fetch_arrow_all(self, force_return_table: Literal[True]) -> Table:
+        ...
+
+    def _fetch_arrow_all(self, force_return_table: bool = False) -> Table | None:
         """Fetches a single Arrow Table from all of the ``ResultBatch``."""
         tables = list(self._fetch_arrow_batches())
         if tables:
             return pa.concat_tables(tables)
         else:
-            return None
+            return self.batches[0].to_arrow() if force_return_table else None
 
     def _fetch_pandas_batches(self, **kwargs) -> Iterator[DataFrame]:
         """Fetches Pandas dataframes in batches, where batch refers to Snowflake Chunk.

test/integ/pandas/test_arrow_pandas.py

@@ -1233,6 +1233,20 @@ def test_simple_arrow_fetch(conn_cnx):
             assert lo == rowcount
 
 
+def test_arrow_zero_rows(conn_cnx):
+    with conn_cnx() as cnx:
+        with cnx.cursor() as cur:
+            cur.execute(SQL_ENABLE_ARROW)
+            cur.execute("select 1::NUMBER(38,0) limit 0")
+            table = cur.fetch_arrow_all(force_return_table=True)
+            # Snowflake will return an integer dtype with maximum bit-length if
+            # no rows are returned
+            assert table.schema[0].type == pyarrow.int64()
+            cur.execute("select 1::NUMBER(38,0) limit 0")
+            # test default behavior
+            assert cur.fetch_arrow_all(force_return_table=False) is None
+
+
 @pytest.mark.parametrize("fetch_fn_name", ["to_arrow", "to_pandas", "create_iter"])
 @pytest.mark.parametrize("pass_connection", [True, False])
 def test_sessions_used(conn_cnx, fetch_fn_name, pass_connection):