Make Arrow Timestamp physical type be returned by new Thrift client.

Pass the `spark.thriftserver.arrowBasedRowSet.timestampAsString` config to make Thriftserver return actual Timestamp type.

Author: juliuszsompolski

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

@@ -17,33 +17,8 @@
 DEFAULT_RESULT_BUFFER_SIZE_BYTES = 10485760
 DEFAULT_ARRAY_SIZE = 100000
 
-_TIMESTAMP_PATTERN = re.compile(r'(\d+-\d+-\d+ \d+:\d+:\d+(\.\d{,6})?)')
-
-
-def _parse_timestamp(value):
-    if type(value) is datetime.datetime:
-        # The cmd exec server will return a datetime.datetime, so no further parsing is needed
-        return value
-    elif value:
-        match = _TIMESTAMP_PATTERN.match(value)
-        if match:
-            if match.group(2):
-                format = '%Y-%m-%d %H:%M:%S.%f'
-                # use the pattern to truncate the value
-                value = match.group()
-            else:
-                format = '%Y-%m-%d %H:%M:%S'
-            value = datetime.datetime.strptime(value, format)
-            return value
-        else:
-            raise Exception('Cannot convert "{}" into a datetime'.format(value))
-    else:
-        return None
-
-
 TYPES_CONVERTER = {
     "decimal": Decimal,
-    "timestamp": _parse_timestamp,
 }
 
 

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

@@ -155,7 +155,13 @@ def open_session(self, session_id=None):
             open_session_req = ttypes.TOpenSessionReq(
                 sessionId=handle_identifier,
                 client_protocol_i64=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
-                client_protocol=None)
+                client_protocol=None,
+                configuration={
+                    # We want to receive proper Timestamp arrow types.
+                    # We set it also in confOverlay in TExecuteStatementReq on a per query basic,
+                    # but it doesn't hurt to also set for the whole session.
+                    "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
+                })
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_protocol_version(response)
             return response.sessionHandle
@@ -271,7 +277,7 @@ def map_type(t_type_entry):
                     ttypes.TTypeId.FLOAT_TYPE: pyarrow.float32(),
                     ttypes.TTypeId.DOUBLE_TYPE: pyarrow.float64(),
                     ttypes.TTypeId.STRING_TYPE: pyarrow.string(),
-                    ttypes.TTypeId.TIMESTAMP_TYPE: pyarrow.string(),
+                    ttypes.TTypeId.TIMESTAMP_TYPE: pyarrow.timestamp('us', None),
                     ttypes.TTypeId.BINARY_TYPE: pyarrow.binary(),
                     ttypes.TTypeId.ARRAY_TYPE: pyarrow.string(),
                     ttypes.TTypeId.MAP_TYPE: pyarrow.string(),
@@ -383,7 +389,11 @@ def execute_command(self, operation, session_handle, max_rows, max_bytes, cursor
             getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
             canReadArrowResult=True,
             canDecompressLZ4Result=False,
-            canDownloadResult=False)
+            canDownloadResult=False,
+            confOverlay={
+                # We want to receive proper Timestamp arrow types.
+                "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
+            })
         resp = self.make_request(self._client.ExecuteStatement, req)
         return self._handle_execute_response(resp, cursor)
 

cmdexec/clients/python/tests/tests.py

@@ -344,21 +344,6 @@ def test_parse_type_converts_decimal(self):
                     self.assertEqual(type(res), type(None))
                     self.assertEqual(res, None)
 
-    def test_parse_type_converts_timestamp(self):
-        # Supported formats: %Y-%m-%d %H:%M:%S.%f and %Y-%m-%d %H:%M:%S
-        for input, output in [(None, None),
-                              ("2021-11-15 11:18:32.349293",
-                               datetime(2021, 11, 15, 11, 18, 32, 349293)),
-                              ("2021-11-15 11:18:32.3492931235",
-                               datetime(2021, 11, 15, 11, 18, 32, 349293)),
-                              ("2021-11-15 11:18:32.34", datetime(2021, 11, 15, 11, 18, 32,
-                                                                  340000)),
-                              ("2021-11-15 11:18:32", datetime(2021, 11, 15, 11, 18, 32))]:
-            with self.subTest(input=input, output=output):
-                res = client.ResultSet.parse_type("timestamp", input)
-                self.assertEqual(type(res), type(output))
-                self.assertEqual(res, output)
-
 
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])