diff --git a/src/snowflake/connector/result_batch.py b/src/snowflake/connector/result_batch.py index 24dd837d3..5077c7892 100644 --- a/src/snowflake/connector/result_batch.py +++ b/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 empty Arrow table based on schema""" + """Returns emtpy 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] diff --git a/src/snowflake/connector/result_set.py b/src/snowflake/connector/result_set.py index 6986a1f6c..226a5b4ab 100644 --- a/src/snowflake/connector/result_set.py +++ b/src/snowflake/connector/result_set.py @@ -189,8 +189,7 @@ def _fetch_pandas_all(self, **kwargs) -> pandas.DataFrame: ignore_index=True, # Don't keep in result batch indexes **kwargs, ) - # Empty dataframe - return self.batches[0].to_pandas() + return pandas.DataFrame(columns=self.batches[0].column_names) def _get_metrics(self) -> dict[str, int]: """Sum up all the chunks' metrics and show them together.""" diff --git a/test/integ/pandas/test_arrow_pandas.py b/test/integ/pandas/test_arrow_pandas.py index e6698eab7..cdfc39ae4 100644 --- a/test/integ/pandas/test_arrow_pandas.py +++ b/test/integ/pandas/test_arrow_pandas.py @@ -1169,18 +1169,6 @@ 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={ @@ -1191,12 +1179,18 @@ def test_pandas_dtypes(conn_cnx): cur.execute( "select 1::integer, 2.3::double, 'foo'::string, current_timestamp()::timestamp where 1=0" ) - expected_types = [numpy.int64, float, object, numpy.datetime64] - assert_pandas_batch_types(cur.fetch_pandas_all(), expected_types) - batches = cur.get_result_batches() + batch = batches[0].to_pandas() + + assert batch.dtypes is not None assert batches[0].to_arrow() is not True - assert_pandas_batch_types(batches[0].to_pandas(), expected_types) + + 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) def test_timestamp_tz(conn_cnx):