SNOW-641540 Fix dtypes of fetch_pandas_all for empty result (#1226)

* SNOW-641540 Fix dtypes of fetch_pandas_all for empty result

When the result is empty, the pandas dataframe information is lost
when creating the iterator. To keep the information, we store the
to_pandas result, and use it for empty dataframe case.

* Address comment

* use list

* format

Author: sfc-gh-sfan

src/snowflake/connector/result_batch.py

@@ -646,7 +646,7 @@ def _get_arrow_iter(
         return self._create_iter(iter_unit=IterUnit.TABLE_UNIT, connection=connection)
 
     def _create_empty_table(self) -> Table:
-        """Returns emtpy Arrow table based on schema"""
+        """Returns empty Arrow table based on schema"""
         if installed_pandas:
             # initialize pyarrow type array corresponding to FIELD_TYPES
             FIELD_TYPE_TO_PA_TYPE = [e.pa_type() for e in FIELD_TYPES]

src/snowflake/connector/result_set.py

@@ -189,7 +189,8 @@ def _fetch_pandas_all(self, **kwargs) -> pandas.DataFrame:
                 ignore_index=True,  # Don't keep in result batch indexes
                 **kwargs,
             )
-        return pandas.DataFrame(columns=self.batches[0].column_names)
+        # Empty dataframe
+        return self.batches[0].to_pandas()
 
     def _get_metrics(self) -> dict[str, int]:
         """Sum up all the chunks' metrics and show them together."""

test/integ/pandas/test_arrow_pandas.py

@@ -1169,6 +1169,18 @@ def assert_dtype_equal(a, b):
     )
 
 
+def assert_pandas_batch_types(
+    batch: pandas.DataFrame, expected_types: list[type]
+) -> None:
+    assert batch.dtypes is not None
+
+    pandas_dtypes = batch.dtypes
+    # pd.string is represented as an np.object
+    # np.dtype string is not the same as pd.string (python)
+    for pandas_dtype, expected_type in zip(pandas_dtypes, expected_types):
+        assert_dtype_equal(pandas_dtype.type, numpy.dtype(expected_type).type)
+
+
 def test_pandas_dtypes(conn_cnx):
     with conn_cnx(
         session_parameters={
@@ -1179,18 +1191,12 @@ def test_pandas_dtypes(conn_cnx):
             cur.execute(
                 "select 1::integer, 2.3::double, 'foo'::string, current_timestamp()::timestamp where 1=0"
             )
-            batches = cur.get_result_batches()
-            batch = batches[0].to_pandas()
+            expected_types = [numpy.int64, float, object, numpy.datetime64]
+            assert_pandas_batch_types(cur.fetch_pandas_all(), expected_types)
 
-            assert batch.dtypes is not None
+            batches = cur.get_result_batches()
             assert batches[0].to_arrow() is not True
-
-            pandas_dtypes = batch.dtypes
-            expected_types = [numpy.int64, float, object, numpy.datetime64]
-            # pd.string is represented as an np.object
-            # np.dtype string is not the same as pd.string (python)
-            for i, typ in enumerate(expected_types):
-                assert_dtype_equal(pandas_dtypes[i].type, numpy.dtype(typ).type)
+            assert_pandas_batch_types(batches[0].to_pandas(), expected_types)
 
 
 def test_timestamp_tz(conn_cnx):