diff --git a/iotdb-client/client-py/iotdb/utils/SessionDataSet.py b/iotdb-client/client-py/iotdb/utils/SessionDataSet.py index 9c5b3ec06ddc5..3f92c78fb6a48 100644 --- a/iotdb-client/client-py/iotdb/utils/SessionDataSet.py +++ b/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() diff --git a/iotdb-client/client-py/iotdb/utils/iotdb_rpc_dataset.py b/iotdb-client/client-py/iotdb/utils/iotdb_rpc_dataset.py index 0edc76f68fd6d..bcb259c6f4263 100644 --- a/iotdb-client/client-py/iotdb/utils/iotdb_rpc_dataset.py +++ b/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 (