fix(snowflake): request microsecond timestamps when fetching Arrow data

ibis/backends/snowflake/__init__.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import glob
+import inspect
 import itertools
 import json
 import os
@@ -38,7 +39,7 @@
 from ibis.backends.sql.compilers.base import STAR
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Iterator, Mapping
+    from collections.abc import Callable, Generator, Iterator, Mapping
     from urllib.parse import ParseResult
 
     import pandas as pd
@@ -154,6 +155,25 @@ def download_file(url: str) -> Generator[str]:
 }
 
 
+def _arrow_fetch_kwargs(fetch: Callable[..., Any]) -> dict[str, bool]:
+    """Keyword arguments for a cursor's Arrow fetch method.
+
+    Snowflake picks the timestamp unit of each Arrow batch from the values in
+    it, so batches of one result can disagree (``timestamp[ns]`` in one,
+    ``timestamp[us]`` in another) and fail to combine.
+    ``force_microsecond_precision`` makes every batch use microseconds. It was
+    added in snowflake-connector-python 4.3.0, so pass it only when the
+    installed connector accepts it.
+    """
+    try:
+        params = inspect.signature(fetch).parameters
+    except (TypeError, ValueError):
+        return {}
+    if "force_microsecond_precision" in params:
+        return {"force_microsecond_precision": True}
+    return {}
+
+
 class Backend(
     SupportsTempTables,
     SQLBackend,
@@ -463,7 +483,7 @@ def to_pyarrow(
 
         sql = self.compile(expr, limit=limit, params=params, **kwargs)
         with self._safe_raw_sql(sql) as cur:
-            res = cur.fetch_arrow_all()
+            res = cur.fetch_arrow_all(**_arrow_fetch_kwargs(cur.fetch_arrow_all))
 
         ibis_schema = expr.as_table().schema()
         if res is None:
@@ -476,7 +496,8 @@ def to_pyarrow(
         return expr.__pyarrow_result__(res, data_mapper=SnowflakePyArrowData)
 
     def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
-        if (table := cursor.fetch_arrow_all()) is None:
+        fetch = cursor.fetch_arrow_all
+        if (table := fetch(**_arrow_fetch_kwargs(fetch))) is None:
             table = schema.to_pyarrow().empty_table()
         df = table.to_pandas(timestamp_as_object=True)
         df.columns = list(schema.names)
@@ -535,7 +556,9 @@ def _make_batch_iter(
                 t.rename_columns(target_schema.names)
                 .cast(target_schema)
                 .to_batches(max_chunksize=chunk_size)
-                for t in cur.fetch_arrow_batches()
+                for t in cur.fetch_arrow_batches(
+                    **_arrow_fetch_kwargs(cur.fetch_arrow_batches)
+                )
             )
 
     def get_schema(

ibis/backends/snowflake/tests/test_client.py

@@ -497,3 +497,14 @@ def test_fancy_column_names(con, column_name):
     t = con.sql(sql)
     assert t.columns == (column_name,)
     assert t[column_name].sum().execute() == value
+
+
+def test_arrow_fetch_kwargs_forces_microseconds_when_supported():
+    from ibis.backends.snowflake import _arrow_fetch_kwargs
+
+    def new_connector(force_microsecond_precision: bool = False): ...
+
+    def old_connector(): ...
+
+    assert _arrow_fetch_kwargs(new_connector) == {"force_microsecond_precision": True}
+    assert _arrow_fetch_kwargs(old_connector) == {}