perf(thrift): join refill pieces once instead of concatenating per read

If urllib3 returns short reads for a field larger than a chunk, bytes +=
recopied everything read so far on every read. Collect the pieces and join
them once. Adds a test that decodes a reply from short reads on both protocols.

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>

Author: vuanhphung

src/databricks/sql/auth/thrift_http_client.py

@@ -175,15 +175,17 @@ def cstringio_buf(self):
 
     def cstringio_refill(self, partialread, reqlen):
         # As in TBufferedTransport: read one more chunk, or all of a larger field.
-        data = partialread
+        parts, have = [partialread], len(partialread)
         if reqlen < self._READ_CHUNK_SIZE:
-            data += self.__resp.read(self._READ_CHUNK_SIZE)
-        while len(data) < reqlen:
-            chunk = self.__resp.read(reqlen - len(data))
+            parts.append(self.__resp.read(self._READ_CHUNK_SIZE))
+            have += len(parts[-1])
+        while have < reqlen:
+            chunk = self.__resp.read(reqlen - have)
             if not chunk:
                 raise EOFError()
-            data += chunk
-        self.__rbuf = BytesIO(data)
+            parts.append(chunk)
+            have += len(chunk)
+        self.__rbuf = BytesIO(b"".join(parts))
         return self.__rbuf
 
     def isOpen(self):

tests/unit/test_thrift_backend.py

@@ -503,6 +503,25 @@ def test_http_client_reply_is_decoded_by_thrift_c_extension(self):
         self.assertEqual(result.results.arrowBatches[0].batch, batch)
         protocol._fast_decode.assert_called()
 
+    def test_http_client_decodes_reply_from_short_urllib3_reads(self):
+        """urllib3 may return fewer bytes than asked for; a field spanning many
+        short reads must still be assembled correctly."""
+        small = [bytes([i % 256]) * 1000 for i in range(300)]
+        big = b"y" * (4 * 1024 * 1024)
+        reply = self._fetch_results_reply(small + [big])
+        for protocol_class in self._protocol_classes():
+            with self.subTest(protocol=protocol_class.__name__):
+                http_client, mock_resp = self._http_client_with_reply(b"")
+                stream = BytesIO(reply)
+                mock_resp.read.side_effect = lambda amt: stream.read(min(amt, 65536))
+                _, client = self._protocol_client(http_client, protocol_class)
+
+                result = client.FetchResults(self._fetch_results_req())
+
+                self.assertEqual(
+                    [b.batch for b in result.results.arrowBatches], small + [big]
+                )
+
     def test_http_client_truncated_reply_raises_eof(self):
         reply = self._fetch_results_reply([b"x" * 100])
         for protocol_class in self._protocol_classes():