tests: address review feedback on test_pandas_compatibility

- Refactor _make_result_set to use ThriftResultSet via normal constructor
  with mocked args (matches test_client.py pattern), instead of bypassing
  __init__ via object.__new__
- Add coverage for additional Arrow types: decimal128, date32/date64,
  timestamp, binary, large_string, list_, struct, map_

Author: moomindani

tests/unit/test_pandas_compatibility.py

@@ -2,10 +2,13 @@
 
 These tests verify that _convert_arrow_table correctly converts Arrow tables
 to Row objects using pandas as an intermediary, covering various data types
-including nullable integers, floats, booleans, and strings.
+including nullable integers, floats, booleans, strings, decimal, dates,
+timestamps, binary, and nested types.
 """
 
+import datetime
 import unittest
+from decimal import Decimal
 from unittest.mock import Mock
 
 import pandas
@@ -16,43 +19,38 @@
 except ImportError:
     pa = None
 
-from databricks.sql.result_set import ResultSet
-from databricks.sql.types import Row
+from databricks.sql.result_set import ThriftResultSet
 
 
-class _ConcreteResultSet(ResultSet):
-    """Minimal concrete subclass of ResultSet for testing _convert_arrow_table."""
+def _make_result_set(description, disable_pandas=False):
+    """Create a ThriftResultSet with mocked dependencies for testing _convert_arrow_table.
 
-    def fetchone(self):
-        pass
+    Mirrors the construction pattern used in test_client.py: pass mocked
+    connection/execute_response/thrift_client to the normal constructor.
+    """
+    mock_connection = Mock()
+    mock_connection.disable_pandas = disable_pandas
 
-    def fetchmany(self, size):
-        pass
+    mock_execute_response = Mock()
+    mock_execute_response.description = description
+    # t_row_set defaults to None, so result_format being None keeps the queue-build branch off.
+    mock_execute_response.result_format = None
 
-    def fetchall(self):
-        pass
+    mock_backend = Mock()
+    # _fill_results_buffer() expects (results, has_more_rows, result_links_count).
+    mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-    def fetchmany_arrow(self, size):
-        pass
-
-    def fetchall_arrow(self):
-        pass
+    return ThriftResultSet(
+        connection=mock_connection,
+        execute_response=mock_execute_response,
+        thrift_client=mock_backend,
+    )
 
 
 @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
 class TestConvertArrowTablePandasCompat(unittest.TestCase):
     """Test _convert_arrow_table with various Arrow types under current pandas version."""
 
-    def _make_result_set(self, description):
-        """Create a minimal ResultSet instance for testing _convert_arrow_table."""
-        mock_connection = Mock()
-        mock_connection.disable_pandas = False
-
-        rs = object.__new__(_ConcreteResultSet)
-        rs.connection = mock_connection
-        rs.description = description
-        return rs
-
     def test_integer_types(self):
         table = pa.table(
             {
@@ -69,7 +67,7 @@ def test_integer_types(self):
             ("int64_col", "bigint", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 3)
@@ -94,7 +92,7 @@ def test_unsigned_integer_types(self):
             ("uint64_col", "bigint", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 2)
@@ -115,7 +113,7 @@ def test_float_types(self):
             ("float64_col", "double", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 3)
@@ -134,7 +132,7 @@ def test_boolean_type(self):
             ("bool_col", "boolean", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 3)
@@ -158,14 +156,186 @@ def test_string_type(self):
             ("str_col", "string", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 3)
         self.assertEqual(rows[0].str_col, "hello")
         self.assertEqual(rows[1].str_col, "world")
         self.assertIsNone(rows[2].str_col)
 
+    def test_large_string_type(self):
+        """large_string is not in dtype_mapping → default Arrow→pandas conversion."""
+        table = pa.table(
+            {
+                "lstr_col": pa.array(["foo", "bar", None], type=pa.large_string()),
+            }
+        )
+        description = [
+            ("lstr_col", "string", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].lstr_col, "foo")
+        self.assertEqual(rows[1].lstr_col, "bar")
+        self.assertIsNone(rows[2].lstr_col)
+
+    def test_binary_type(self):
+        table = pa.table(
+            {
+                "bin_col": pa.array([b"hello", None, b"world"], type=pa.binary()),
+            }
+        )
+        description = [
+            ("bin_col", "binary", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].bin_col, b"hello")
+        self.assertIsNone(rows[1].bin_col)
+        self.assertEqual(rows[2].bin_col, b"world")
+
+    def test_decimal_type(self):
+        table = pa.table(
+            {
+                "dec_col": pa.array(
+                    [Decimal("1.23"), None, Decimal("99.99")],
+                    type=pa.decimal128(5, 2),
+                ),
+            }
+        )
+        description = [
+            ("dec_col", "decimal", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].dec_col, Decimal("1.23"))
+        self.assertIsNone(rows[1].dec_col)
+        self.assertEqual(rows[2].dec_col, Decimal("99.99"))
+
+    def test_date_types(self):
+        """date32 and date64 → datetime.date objects via date_as_object=True."""
+        table = pa.table(
+            {
+                "date32_col": pa.array(
+                    [datetime.date(2024, 1, 1), None, datetime.date(2026, 5, 19)],
+                    type=pa.date32(),
+                ),
+                "date64_col": pa.array(
+                    [None, datetime.date(2024, 12, 31), datetime.date(2026, 1, 1)],
+                    type=pa.date64(),
+                ),
+            }
+        )
+        description = [
+            ("date32_col", "date", None, None, None, None, None),
+            ("date64_col", "date", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].date32_col, datetime.date(2024, 1, 1))
+        self.assertIsNone(rows[0].date64_col)
+        self.assertIsNone(rows[1].date32_col)
+        self.assertEqual(rows[1].date64_col, datetime.date(2024, 12, 31))
+        self.assertEqual(rows[2].date32_col, datetime.date(2026, 5, 19))
+        self.assertEqual(rows[2].date64_col, datetime.date(2026, 1, 1))
+
+    def test_timestamp_type(self):
+        """timestamp → datetime.datetime objects via timestamp_as_object=True."""
+        ts1 = datetime.datetime(2024, 1, 1, 12, 30, 45)
+        ts2 = datetime.datetime(2026, 5, 19, 9, 15, 0)
+        table = pa.table(
+            {
+                "ts_col": pa.array([ts1, None, ts2], type=pa.timestamp("us")),
+            }
+        )
+        description = [
+            ("ts_col", "timestamp", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].ts_col, ts1)
+        self.assertIsNone(rows[1].ts_col)
+        self.assertEqual(rows[2].ts_col, ts2)
+
+    def test_list_type(self):
+        table = pa.table(
+            {
+                "list_col": pa.array(
+                    [[1, 2, 3], None, [4, 5]],
+                    type=pa.list_(pa.int64()),
+                ),
+            }
+        )
+        description = [
+            ("list_col", "array", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(list(rows[0].list_col), [1, 2, 3])
+        self.assertIsNone(rows[1].list_col)
+        self.assertEqual(list(rows[2].list_col), [4, 5])
+
+    def test_struct_type(self):
+        table = pa.table(
+            {
+                "struct_col": pa.array(
+                    [{"x": 1, "y": "a"}, None, {"x": 3, "y": "c"}],
+                    type=pa.struct([("x", pa.int64()), ("y", pa.string())]),
+                ),
+            }
+        )
+        description = [
+            ("struct_col", "struct", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(rows[0].struct_col, {"x": 1, "y": "a"})
+        self.assertIsNone(rows[1].struct_col)
+        self.assertEqual(rows[2].struct_col, {"x": 3, "y": "c"})
+
+    def test_map_type(self):
+        table = pa.table(
+            {
+                "map_col": pa.array(
+                    [[("k1", 1), ("k2", 2)], None, [("k3", 3)]],
+                    type=pa.map_(pa.string(), pa.int64()),
+                ),
+            }
+        )
+        description = [
+            ("map_col", "map", None, None, None, None, None),
+        ]
+
+        rs = _make_result_set(description)
+        rows = rs._convert_arrow_table(table)
+
+        self.assertEqual(len(rows), 3)
+        self.assertEqual(list(rows[0].map_col), [("k1", 1), ("k2", 2)])
+        self.assertIsNone(rows[1].map_col)
+        self.assertEqual(list(rows[2].map_col), [("k3", 3)])
+
     def test_mixed_types(self):
         """Test a table with a mix of types, similar to real query results."""
         table = pa.table(
@@ -183,7 +353,7 @@ def test_mixed_types(self):
             ("active", "boolean", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 3)
@@ -220,7 +390,7 @@ def test_duplicate_column_names(self):
             ("col", "int", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 2)
@@ -238,7 +408,7 @@ def test_empty_table(self):
             ("col", "int", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 0)
@@ -255,7 +425,7 @@ def test_all_nulls(self):
             ("str_col", "string", None, None, None, None, None),
         ]
 
-        rs = self._make_result_set(description)
+        rs = _make_result_set(description)
         rows = rs._convert_arrow_table(table)
 
         self.assertEqual(len(rows), 2)
@@ -266,15 +436,11 @@ def test_all_nulls(self):
 
     def test_disable_pandas_path(self):
         """Verify the non-pandas code path still works."""
-        mock_connection = Mock()
-        mock_connection.disable_pandas = True
-
-        rs = object.__new__(_ConcreteResultSet)
-        rs.connection = mock_connection
-        rs.description = [
+        description = [
             ("id", "bigint", None, None, None, None, None),
             ("name", "string", None, None, None, None, None),
         ]
+        rs = _make_result_set(description, disable_pandas=True)
 
         table = pa.table(
             {