Optimize Python session result decoding

iotdb-client/client-py/iotdb/utils/SessionDataSet.py

@@ -18,15 +18,14 @@
 import logging
 from typing import Optional
 
+import pandas as pd
 from iotdb.utils.Field import Field
+from iotdb.utils.iotdb_rpc_dataset import IoTDBRpcDataSet
 
 # for package
 from iotdb.utils.IoTDBConstants import TSDataType
-from iotdb.utils.iotdb_rpc_dataset import IoTDBRpcDataSet
 from iotdb.utils.RowRecord import RowRecord
 
-import pandas as pd
-
 logger = logging.getLogger("IoTDB")
 
 
@@ -109,14 +108,14 @@ def has_next(self):
         return self.iotdb_rpc_data_set.next()
 
     def next(self):
-        if not self.iotdb_rpc_data_set.has_cached_data_frame:
-            if not self.has_next():
-                return None
-        return self.construct_row_record_from_data_frame()
-
-    def construct_row_record_from_data_frame(self):
-        df = self.iotdb_rpc_data_set.data_frame
-        row = df.iloc[self.row_index].to_list()
+        if not self.has_next():
+            return None
+        return self._construct_row_record()
+
+    def _construct_row_record(self):
+        row = self.iotdb_rpc_data_set._pop_row()
+        if row is None:
+            return None
         if self.iotdb_rpc_data_set.ignore_timestamp:
             for field, value in zip(self.__field_list, row):
                 field.value = value
@@ -133,14 +132,11 @@ def construct_row_record_from_data_frame(self):
                 row[0],
                 self.__field_list,
             )
-        self.row_index += 1
-        if self.row_index == len(df):
-            self.row_index = 0
-            self.iotdb_rpc_data_set.has_cached_data_frame = False
-            self.iotdb_rpc_data_set.data_frame = None
-
         return row_record
 
+    def construct_row_record_from_data_frame(self):
+        return self._construct_row_record()
+
     def close_operation_handle(self):
         self.iotdb_rpc_data_set.close()
 

iotdb-client/client-py/iotdb/utils/iotdb_rpc_dataset.py

@@ -18,18 +18,18 @@
 
 # for package
 import logging
+from collections import deque
 from typing import Optional
 
 import numpy as np
 import pandas as pd
-from thrift.transport import TTransport
-
-from iotdb.thrift.rpc.IClientRPCService import TSFetchResultsReq, TSCloseOperationReq
+from iotdb.thrift.rpc.IClientRPCService import TSCloseOperationReq, TSFetchResultsReq
 from iotdb.tsfile.utils.date_utils import parse_int_to_date
 from iotdb.tsfile.utils.tsblock_serde import deserialize
 from iotdb.utils.exception import IoTDBConnectionException
 from iotdb.utils.IoTDBConstants import TSDataType
-from iotdb.utils.rpc_utils import verify_success, convert_to_timestamp
+from iotdb.utils.rpc_utils import convert_to_timestamp, verify_success
+from thrift.transport import TTransport
 
 logger = logging.getLogger("IoTDB")
 TIMESTAMP_STR = "Time"
@@ -122,11 +122,13 @@ def __init__(
         self.__zone_id = zone_id
         self.__time_precision = time_precision
         self.__df_buffer = None  # Buffer for streaming DataFrames
+        self.__row_buffer = deque()
 
     def close(self):
         if self.__is_closed:
             return
         self.__df_buffer = None  # Clean up streaming DataFrame buffer
+        self.__row_buffer.clear()
         if self.__client is not None:
             try:
                 status = self.__client.closeOperation(
@@ -148,16 +150,74 @@ def close(self):
             self.__client = None
 
     def next(self):
-        if not self.has_cached_data_frame:
-            self.construct_one_data_frame()
-        if self.has_cached_data_frame:
-            return True
-        if self.__empty_resultSet:
-            return False
-        if self.__more_data and self.fetch_results():
-            self.construct_one_data_frame()
-            return True
-        return False
+        return bool(self.__row_buffer) or self._fill_row_buffer()
+
+    def _pop_row(self):
+        if not self.next():
+            return None
+        return self.__row_buffer.popleft()
+
+    def _pop_rows(self, max_rows=None):
+        rows = []
+        while max_rows is None or len(rows) < max_rows:
+            if not self.__row_buffer and not self._fill_row_buffer():
+                break
+            remaining = (
+                len(self.__row_buffer)
+                if max_rows is None
+                else min(len(self.__row_buffer), max_rows - len(rows))
+            )
+            rows.extend(self.__row_buffer.popleft() for _ in range(remaining))
+        return rows
+
+    def _fill_row_buffer(self):
+        while not self.__row_buffer:
+            while self.__query_result is not None and self.__query_result_index < len(
+                self.__query_result
+            ):
+                block = self.__query_result[self.__query_result_index]
+                self.__query_result[self.__query_result_index] = None
+                self.__query_result_index += 1
+                self.__row_buffer.extend(self._deserialize_rows(block))
+                if self.__row_buffer:
+                    return True
+            if self.__empty_resultSet or not self.__more_data:
+                return False
+            if not self.fetch_results():
+                return False
+        return True
+
+    def _deserialize_rows(self, serialized_block):
+        time_array, column_arrays, null_indicators, row_count = deserialize(
+            memoryview(serialized_block)
+        )
+        columns = []
+        if not self.ignore_timestamp:
+            columns.append(iter(time_array.tolist()))
+        for location in self.__column_index_2_tsblock_column_index_list:
+            if location < 0:
+                continue
+            columns.append(
+                self._column_values(
+                    column_arrays[location],
+                    null_indicators[location],
+                    row_count,
+                    self.__data_type_for_tsblock_column[location],
+                )
+            )
+        return list(zip(*columns)) if columns else [tuple() for _ in range(row_count)]
+
+    @staticmethod
+    def _column_values(values, nulls, row_count, data_type):
+        source = values.tolist() if hasattr(values, "tolist") else list(values)
+        if nulls is None:
+            return iter(source)
+        if data_type == TSDataType.BOOLEAN and len(source) == row_count:
+            return (
+                None if nulls[index] else value for index, value in enumerate(source)
+            )
+        source_values = iter(source)
+        return (None if is_null else next(source_values) for is_null in nulls)
 
     def construct_one_data_frame(self):
         if self.has_cached_data_frame or self.__query_result is None:
@@ -233,7 +293,7 @@ def construct_one_data_frame(self):
             self.has_cached_data_frame = True
 
     def has_cached_result(self):
-        return self.has_cached_data_frame
+        return bool(self.__row_buffer) or self.has_cached_data_frame
 
     def _has_next_result_set(self):
         if (self.__query_result is not None) and (