SNOW-1279205: add capability to debug arrow data (#1918)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -11,6 +11,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 - v3.11.0(TBD)
 
   - Added support for `token_file_path` connection parameter to read an OAuth token from a file when connecting to Snowflake.
+  - Added support for `debug_arrow_chunk` connection parameter to allow debugging raw arrow data in case of arrow data parsing failure.
 
 - v3.10.0(April 29,2024)
 

src/snowflake/connector/connection.py

@@ -281,6 +281,10 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # Disable console login and fall back to getting SSO URL from GS
+    "debug_arrow_chunk": (
+        False,
+        bool,
+    ),  # log raw arrow chunk for debugging purpuse in case there is malformed arrow data
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")

src/snowflake/connector/result_batch.py

@@ -664,11 +664,21 @@ def _create_iter(
     ) -> Iterator[dict | Exception] | Iterator[tuple | Exception] | Iterator[Table]:
         """Create an iterator for the ResultBatch. Used by get_arrow_iter."""
         if self._local:
-            return self._from_data(self._data, iter_unit)
+            try:
+                return self._from_data(self._data, iter_unit)
+            except Exception:
+                if connection and getattr(connection, "_debug_arrow_chunk", False):
+                    logger.debug(f"arrow data can not be parsed: {self._data}")
+                raise
         response = self._download(connection=connection)
         logger.debug(f"started loading result batch id: {self.id}")
         with TimerContextManager() as load_metric:
-            loaded_data = self._load(response, iter_unit)
+            try:
+                loaded_data = self._load(response, iter_unit)
+            except Exception:
+                if connection and getattr(connection, "_debug_arrow_chunk", False):
+                    logger.debug(f"arrow data can not be parsed: {response}")
+                raise
         logger.debug(f"finished loading result batch id: {self.id}")
         self._metrics[DownloadMetrics.load.value] = load_metric.get_timing_millis()
         return loaded_data

test/integ/test_arrow_result.py

@@ -5,8 +5,10 @@
 
 from __future__ import annotations
 
+import base64
 import itertools
 import json
+import logging
 import os
 import random
 import re
@@ -16,7 +18,7 @@
 import pytest
 
 import snowflake.connector
-from snowflake.connector.errors import ProgrammingError
+from snowflake.connector.errors import OperationalError, ProgrammingError
 
 try:
     from snowflake.connector.util_text import random_string
@@ -1103,6 +1105,20 @@ def iterate_over_test_chunk(
                     assert str(arrow_res[0]) == expected[i]
 
 
+@pytest.mark.parametrize("debug_arrow_chunk", [True, False])
+def test_arrow_bad_data(conn_cnx, caplog, debug_arrow_chunk):
+    with caplog.at_level(logging.DEBUG):
+        with conn_cnx(
+            debug_arrow_chunk=debug_arrow_chunk
+        ) as arrow_cnx, arrow_cnx.cursor() as cursor:
+            cursor.execute("select 1")
+            cursor._result_set.batches[0]._data = base64.b64encode(b"wrong_data")
+            with pytest.raises(OperationalError):
+                cursor.fetchone()
+    expr = bool("arrow data can not be parsed" in caplog.text)
+    assert expr if debug_arrow_chunk else not expr
+
+
 def init(conn_cnx, table, column, values):
     with conn_cnx() as json_cnx:
         cursor_json = json_cnx.cursor()