fmt

Author: madhav-db

src/databricks/sql/client.py

@@ -1351,7 +1351,9 @@ def _convert_arrow_table(self, table):
         # Need to use nullable types, as otherwise type can change when there are missing values.
         # See https://arrow.apache.org/docs/python/pandas.html#nullable-types
         # NOTE: This api is epxerimental https://pandas.pydata.org/pandas-docs/stable/user_guide/integer_na.html
-        DEFAULT_DTYPE_MAPPING: Dict[pyarrow.DataType, pandas.api.extensions.ExtensionDtype] = {
+        DEFAULT_DTYPE_MAPPING: Dict[
+            pyarrow.DataType, pandas.api.extensions.ExtensionDtype
+        ] = {
             pyarrow.int8(): pandas.Int8Dtype(),
             pyarrow.int16(): pandas.Int16Dtype(),
             pyarrow.int32(): pandas.Int32Dtype(),
@@ -1365,7 +1367,10 @@ def _convert_arrow_table(self, table):
             pyarrow.float64(): pandas.Float64Dtype(),
             pyarrow.string(): pandas.StringDtype(),
         }
-        dtype_mapping = {**DEFAULT_DTYPE_MAPPING, **self.connection._arrow_pandas_type_override}
+        dtype_mapping = {
+            **DEFAULT_DTYPE_MAPPING,
+            **self.connection._arrow_pandas_type_override,
+        }
 
         to_pandas_kwargs: dict[str, Any] = {
             "types_mapper": dtype_mapping.get,

tests/unit/test_arrow_conversion.py

@@ -15,23 +15,31 @@ def mock_connection():
     conn.disable_pandas = False
     conn._arrow_pandas_type_override = {}
     conn._arrow_to_pandas_kwargs = {}
-    if not hasattr(conn, '_arrow_to_pandas_kwargs'):
+    if not hasattr(conn, "_arrow_to_pandas_kwargs"):
         conn._arrow_to_pandas_kwargs = {}
     return conn
 
+
 @pytest.fixture
 def mock_thrift_backend(sample_arrow_table):
     tb = MagicMock()
-    empty_arrays = [pyarrow.array([], type=field.type) for field in sample_arrow_table.schema]
-    empty_table = pyarrow.Table.from_arrays(empty_arrays, schema=sample_arrow_table.schema)
-    tb.fetch_results.return_value = (ArrowQueue(empty_table, 0) , False)
+    empty_arrays = [
+        pyarrow.array([], type=field.type) for field in sample_arrow_table.schema
+    ]
+    empty_table = pyarrow.Table.from_arrays(
+        empty_arrays, schema=sample_arrow_table.schema
+    )
+    tb.fetch_results.return_value = (ArrowQueue(empty_table, 0), False)
     return tb
 
+
 @pytest.fixture
 def mock_raw_execute_response():
     er = MagicMock(spec=ExecuteResponse)
-    er.description = [("col_int", "int", None, None, None, None, None),
-                      ("col_str", "string", None, None, None, None, None)]
+    er.description = [
+        ("col_int", "int", None, None, None, None, None),
+        ("col_str", "string", None, None, None, None, None),
+    ]
     er.arrow_schema_bytes = None
     er.arrow_queue = None
     er.has_more_rows = False
@@ -42,27 +50,33 @@ def mock_raw_execute_response():
     er.is_staging_operation = False
     return er
 
+
 @pytest.fixture
 def sample_arrow_table():
     data = [
         pyarrow.array([1, 2, 3], type=pyarrow.int32()),
-        pyarrow.array(["a", "b", "c"], type=pyarrow.string())
+        pyarrow.array(["a", "b", "c"], type=pyarrow.string()),
     ]
-    schema = pyarrow.schema([
-        ('col_int', pyarrow.int32()),
-        ('col_str', pyarrow.string())
-    ])
+    schema = pyarrow.schema(
+        [("col_int", pyarrow.int32()), ("col_str", pyarrow.string())]
+    )
     return pyarrow.Table.from_arrays(data, schema=schema)
 
 
-def test_convert_arrow_table_default(mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table):
-    mock_raw_execute_response.arrow_queue = ArrowQueue(sample_arrow_table, sample_arrow_table.num_rows)
+def test_convert_arrow_table_default(
+    mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table
+):
+    mock_raw_execute_response.arrow_queue = ArrowQueue(
+        sample_arrow_table, sample_arrow_table.num_rows
+    )
     rs = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
     result_one = rs.fetchone()
     assert isinstance(result_one, Row)
     assert result_one.col_int == 1
     assert result_one.col_str == "a"
-    mock_raw_execute_response.arrow_queue = ArrowQueue(sample_arrow_table, sample_arrow_table.num_rows)
+    mock_raw_execute_response.arrow_queue = ArrowQueue(
+        sample_arrow_table, sample_arrow_table.num_rows
+    )
     rs = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
     result_all = rs.fetchall()
     assert len(result_all) == 3
@@ -71,9 +85,13 @@ def test_convert_arrow_table_default(mock_connection, mock_thrift_backend, mock_
     assert result_all[1].col_str == "b"
 
 
-def test_convert_arrow_table_disable_pandas(mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table):
+def test_convert_arrow_table_disable_pandas(
+    mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table
+):
     mock_connection.disable_pandas = True
-    mock_raw_execute_response.arrow_queue = ArrowQueue(sample_arrow_table, sample_arrow_table.num_rows)
+    mock_raw_execute_response.arrow_queue = ArrowQueue(
+        sample_arrow_table, sample_arrow_table.num_rows
+    )
     rs = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
     result = rs.fetchall()
     assert len(result) == 3
@@ -84,9 +102,15 @@ def test_convert_arrow_table_disable_pandas(mock_connection, mock_thrift_backend
     assert isinstance(sample_arrow_table.column(1)[0].as_py(), str)
 
 
-def test_convert_arrow_table_type_override(mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table):
-    mock_connection._arrow_pandas_type_override = {pyarrow.int32(): pandas.Float64Dtype()}
-    mock_raw_execute_response.arrow_queue = ArrowQueue(sample_arrow_table, sample_arrow_table.num_rows)
+def test_convert_arrow_table_type_override(
+    mock_connection, mock_thrift_backend, mock_raw_execute_response, sample_arrow_table
+):
+    mock_connection._arrow_pandas_type_override = {
+        pyarrow.int32(): pandas.Float64Dtype()
+    }
+    mock_raw_execute_response.arrow_queue = ArrowQueue(
+        sample_arrow_table, sample_arrow_table.num_rows
+    )
     rs = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
     result = rs.fetchall()
     assert len(result) == 3
@@ -95,34 +119,44 @@ def test_convert_arrow_table_type_override(mock_connection, mock_thrift_backend,
     assert result[0].col_str == "a"
 
 
-def test_convert_arrow_table_to_pandas_kwargs(mock_connection, mock_thrift_backend, mock_raw_execute_response):
+def test_convert_arrow_table_to_pandas_kwargs(
+    mock_connection, mock_thrift_backend, mock_raw_execute_response
+):
     dt_obj = datetime.datetime(2021, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc)
-    ts_array = pyarrow.array([dt_obj], type=pyarrow.timestamp('us', tz='UTC'))
-    ts_schema = pyarrow.schema([('col_ts', pyarrow.timestamp('us', tz='UTC'))])
+    ts_array = pyarrow.array([dt_obj], type=pyarrow.timestamp("us", tz="UTC"))
+    ts_schema = pyarrow.schema([("col_ts", pyarrow.timestamp("us", tz="UTC"))])
     ts_table = pyarrow.Table.from_arrays([ts_array], schema=ts_schema)
 
-    mock_raw_execute_response.description = [("col_ts", "timestamp", None, None, None, None, None)]
+    mock_raw_execute_response.description = [
+        ("col_ts", "timestamp", None, None, None, None, None)
+    ]
     mock_raw_execute_response.arrow_queue = ArrowQueue(ts_table, ts_table.num_rows)
 
     # Scenario 1: timestamp_as_object = True. Observed as datetime.datetime in Row.
     mock_connection._arrow_to_pandas_kwargs = {"timestamp_as_object": True}
-    rs_ts_true = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
+    rs_ts_true = ResultSet(
+        mock_connection, mock_raw_execute_response, mock_thrift_backend
+    )
     result_true = rs_ts_true.fetchall()
     assert len(result_true) == 1
     assert isinstance(result_true[0].col_ts, datetime.datetime)
 
     # Scenario 2: timestamp_as_object = False. Observed as pandas.Timestamp in Row for this input.
     mock_raw_execute_response.arrow_queue = ArrowQueue(ts_table, ts_table.num_rows)
     mock_connection._arrow_to_pandas_kwargs = {"timestamp_as_object": False}
-    rs_ts_false = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
+    rs_ts_false = ResultSet(
+        mock_connection, mock_raw_execute_response, mock_thrift_backend
+    )
     result_false = rs_ts_false.fetchall()
     assert len(result_false) == 1
     assert isinstance(result_false[0].col_ts, pandas.Timestamp)
 
     # Scenario 3: no override. Observed as datetime.datetime in Row since timestamp_as_object is True by default.
     mock_raw_execute_response.arrow_queue = ArrowQueue(ts_table, ts_table.num_rows)
     mock_connection._arrow_to_pandas_kwargs = {}
-    rs_ts_true = ResultSet(mock_connection, mock_raw_execute_response, mock_thrift_backend)
+    rs_ts_true = ResultSet(
+        mock_connection, mock_raw_execute_response, mock_thrift_backend
+    )
     result_true = rs_ts_true.fetchall()
     assert len(result_true) == 1
     assert isinstance(result_true[0].col_ts, datetime.datetime)