Add structured type support to PySQL

This PR adds support for complex types (timestamps, decimals, structs & arrays) being returned as native Arrow types in PySQL.

Author: NiallEgan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -72,6 +72,15 @@ def __init__(self,
         # _disable_pandas
         #  In case the deserialisation through pandas causes any issues, it can be disabled with
         #  this flag.
+        # _use_arrow_native_complex_types
+        # DBR will return native Arrow types for structs, arrays and maps instead of Arrow strings
+        # (True by default)
+        # _use_arrow_native_decimals
+        # Databricks runtime will return native Arrow types for decimals instead of Arrow strings
+        # (True by default)
+        # _use_arrow_native_timestamps
+        # Databricks runtime will return native Arrow types for timestamps instead of Arrow strings
+        # (True by default)
 
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)

cmdexec/clients/python/src/databricks/sql/thrift_backend.py

@@ -87,6 +87,9 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
             raise ValueError("No valid connection settings.")
 
         self._initialize_retry_args(kwargs)
+        self._use_arrow_native_complex_types = kwargs.get("_use_arrow_native_complex_types", True)
+        self._use_arrow_native_decimals = kwargs.get("_use_arrow_native_decimals", True)
+        self._use_arrow_native_timestamps = kwargs.get("_use_arrow_native_timestamps", True)
 
         # Configure tls context
         ssl_context = create_default_context(cafile=kwargs.get("_tls_trusted_ca_file"))
@@ -602,6 +605,13 @@ def _check_direct_results_for_error(t_spark_direct_results):
     def execute_command(self, operation, session_handle, max_rows, max_bytes, cursor):
         assert (session_handle is not None)
 
+        spark_arrow_types = ttypes.TSparkArrowTypes(
+            timestampAsArrow=self._use_arrow_native_timestamps,
+            decimalAsArrow=self._use_arrow_native_decimals,
+            complexTypesAsArrow=self._use_arrow_native_complex_types,
+            # TODO: The current Arrow type used for intervals can not be deserialised in PyArrow
+            # DBR should be changed to use month_day_nano_interval
+            intervalTypesAsArrow=False)
         req = ttypes.TExecuteStatementReq(
             sessionHandle=session_handle,
             statement=operation,
@@ -613,7 +623,8 @@ def execute_command(self, operation, session_handle, max_rows, max_bytes, cursor
             confOverlay={
                 # We want to receive proper Timestamp arrow types.
                 "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
-            })
+            },
+            useArrowNativeTypes=spark_arrow_types)
         resp = self.make_request(self._client.ExecuteStatement, req)
         return self._handle_execute_response(resp, cursor)
 

cmdexec/clients/python/tests/test_thrift_backend.py

@@ -1251,6 +1251,35 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
         self.assertIn("Setting initial namespace not supported by the DBR version",
                       str(cm.exception))
 
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.ThriftBackend._handle_execute_response")
+    def test_execute_command_sets_complex_type_fields_correctly(self, mock_handle_execute_response,
+                                                                tcli_service_class):
+        tcli_service_instance = tcli_service_class.return_value
+        # Iterate through each possible combination of native types (True, False and unset)
+        for (complex, timestamp, decimals) in itertools.product(
+            [True, False, None], [True, False, None], [True, False, None]):
+            complex_arg_types = {}
+            if complex is not None:
+                complex_arg_types["_use_arrow_native_complex_types"] = complex
+            if timestamp is not None:
+                complex_arg_types["_use_arrow_native_timestamps"] = timestamp
+            if decimals is not None:
+                complex_arg_types["_use_arrow_native_decimals"] = decimals
+
+            thrift_backend = ThriftBackend("foobar", 443, "path", [], **complex_arg_types)
+            thrift_backend.execute_command(Mock(), Mock(), 100, 100, Mock())
+
+            t_execute_statement_req = tcli_service_instance.ExecuteStatement.call_args[0][0]
+            # If the value is unset, the native type should default to True
+            self.assertEqual(t_execute_statement_req.useArrowNativeTypes.timestampAsArrow,
+                             complex_arg_types.get("_use_arrow_native_timestamps", True))
+            self.assertEqual(t_execute_statement_req.useArrowNativeTypes.decimalAsArrow,
+                             complex_arg_types.get("_use_arrow_native_decimals", True))
+            self.assertEqual(t_execute_statement_req.useArrowNativeTypes.complexTypesAsArrow,
+                             complex_arg_types.get("_use_arrow_native_complex_types", True))
+            self.assertFalse(t_execute_statement_req.useArrowNativeTypes.intervalTypesAsArrow)
+
 
 if __name__ == '__main__':
     unittest.main()