apache-iotdb 2.0.2__py3-none-any.whl → 2.0.4.dev0__py3-none-any.whl

iotdb/tsfile/utils/tsblock_serde.py

@@ -0,0 +1,266 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import numpy as np
+
+
+# Serialized tsBlock:
+#    +-------------+---------------+---------+------------+-----------+----------+
+#    | val col cnt | val col types | pos cnt | encodings  | time col  | val col  |
+#    +-------------+---------------+---------+------------+-----------+----------+
+#    | int32       | list[byte]    | int32   | list[byte] |  bytes    | byte     |
+#    +-------------+---------------+---------+------------+-----------+----------+
+
+
+def deserialize(buffer):
+    value_column_count, buffer = read_int_from_buffer(buffer)
+    data_types, buffer = read_column_types(buffer, value_column_count)
+
+    position_count, buffer = read_int_from_buffer(buffer)
+    column_encodings, buffer = read_column_encoding(buffer, value_column_count + 1)
+
+    time_column_values, _, buffer = read_column(
+        column_encodings[0], buffer, 2, position_count
+    )
+    column_values = []
+    null_indicators = []
+    for i in range(value_column_count):
+        column_value, null_indicator, buffer = read_column(
+            column_encodings[i + 1], buffer, data_types[i], position_count
+        )
+        column_values.append(column_value)
+        null_indicators.append(null_indicator)
+
+    return time_column_values, column_values, null_indicators, position_count
+
+
+# General Methods
+
+
+def read_int_from_buffer(buffer):
+    res = np.frombuffer(buffer, dtype=">i4", count=1)
+    buffer = buffer[4:]
+    return res[0], buffer
+
+
+def read_byte_from_buffer(buffer):
+    return read_from_buffer(buffer, 1)
+
+
+def read_from_buffer(buffer, size):
+    res = buffer[:size]
+    new_buffer = buffer[size:]
+    return res, new_buffer
+
+
+# Read ColumnType
+
+
+def read_column_types(buffer, value_column_count):
+    data_types = np.frombuffer(buffer, dtype=np.uint8, count=value_column_count)
+    new_buffer = buffer[value_column_count:]
+    if not np.all(np.isin(data_types, (0, 1, 2, 3, 4, 5))):
+        raise Exception("Invalid data type encountered: " + str(data_types))
+    return data_types, new_buffer
+
+
+# Read ColumnEncodings
+
+
+def read_column_encoding(buffer, size):
+    encodings = np.frombuffer(buffer, dtype=np.uint8, count=size)
+    new_buffer = buffer[size:]
+    return encodings, new_buffer
+
+
+# Read Column
+
+
+def deserialize_null_indicators(buffer, size):
+    may_have_null = buffer[0]
+    buffer = buffer[1:]
+    if may_have_null != 0:
+        return deserialize_from_boolean_array(buffer, size)
+    return None, buffer
+
+
+# Serialized data layout:
+#    +---------------+-----------------+-------------+
+#    | may have null | null indicators |   values    |
+#    +---------------+-----------------+-------------+
+#    | byte          | list[byte]      | list[int64] |
+#    +---------------+-----------------+-------------+
+
+
+def read_int64_column(buffer, data_type, position_count):
+    null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+    if null_indicators is None:
+        size = position_count
+    else:
+        size = np.count_nonzero(~null_indicators)
+
+    if data_type == 2:
+        dtype = ">i8"
+    elif data_type == 4:
+        dtype = ">f8"
+    else:
+        raise Exception("Invalid data type: " + str(data_type))
+    values = np.frombuffer(buffer, dtype, count=size)
+    buffer = buffer[size * 8 :]
+    return values, null_indicators, buffer
+
+
+# Serialized data layout:
+#    +---------------+-----------------+-------------+
+#    | may have null | null indicators |   values    |
+#    +---------------+-----------------+-------------+
+#    | byte          | list[byte]      | list[int32] |
+#    +---------------+-----------------+-------------+
+
+
+def read_int32_column(buffer, data_type, position_count):
+    null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+    if null_indicators is None:
+        size = position_count
+    else:
+        size = np.count_nonzero(~null_indicators)
+
+    if data_type == 1:
+        dtype = ">i4"
+    elif data_type == 3:
+        dtype = ">f4"
+    else:
+        raise Exception("Invalid data type: " + str(data_type))
+    values = np.frombuffer(buffer, dtype, count=size)
+    buffer = buffer[size * 4 :]
+    return values, null_indicators, buffer
+
+
+# Serialized data layout:
+#    +---------------+-----------------+-------------+
+#    | may have null | null indicators |   values    |
+#    +---------------+-----------------+-------------+
+#    | byte          | list[byte]      | list[byte] |
+#    +---------------+-----------------+-------------+
+
+
+def read_byte_column(buffer, data_type, position_count):
+    if data_type != 0:
+        raise Exception("Invalid data type: " + data_type)
+    null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+    res, buffer = deserialize_from_boolean_array(buffer, position_count)
+    return res, null_indicators, buffer
+
+
+def deserialize_from_boolean_array(buffer, size):
+    num_bytes = (size + 7) // 8
+    packed_boolean_array, buffer = read_from_buffer(buffer, num_bytes)
+    arr = np.frombuffer(packed_boolean_array, dtype=np.uint8)
+    output = np.unpackbits(arr)[:size].astype(bool)
+    return output, buffer
+
+
+# Serialized data layout:
+#    +---------------+-----------------+-------------+
+#    | may have null | null indicators |   values    |
+#    +---------------+-----------------+-------------+
+#    | byte          | list[byte]      | list[entry] |
+#    +---------------+-----------------+-------------+
+#
+# Each entry is represented as:
+#    +---------------+-------+
+#    | value length  | value |
+#    +---------------+-------+
+#    | int32         | bytes |
+#    +---------------+-------+
+
+
+def read_binary_column(buffer, data_type, position_count):
+    if data_type != 5:
+        raise Exception("Invalid data type: " + data_type)
+    null_indicators, buffer = deserialize_null_indicators(buffer, position_count)
+
+    if null_indicators is None:
+        size = position_count
+    else:
+        size = np.count_nonzero(~null_indicators)
+    values = np.empty(size, dtype=object)
+    for i in range(size):
+        length, buffer = read_int_from_buffer(buffer)
+        res, buffer = read_from_buffer(buffer, length)
+        values[i] = res.tobytes()
+    return values, null_indicators, buffer
+
+
+# Serialized data layout:
+#    +-----------+-------------------------+
+#    | encoding  | serialized inner column |
+#    +-----------+-------------------------+
+#    | byte      | list[byte]              |
+#    +-----------+-------------------------+
+
+
+def read_run_length_column(buffer, data_type, position_count):
+    encoding, buffer = read_byte_from_buffer(buffer)
+    column, null_indicators, buffer = read_column(encoding[0], buffer, data_type, 1)
+    return (
+        repeat(column, data_type, position_count),
+        (
+            None
+            if null_indicators is None
+            else np.full(position_count, null_indicators[0])
+        ),
+        buffer,
+    )
+
+
+def repeat(column, data_type, position_count):
+    if data_type in (0, 5):
+        if column.size == 1:
+            return np.full(
+                position_count, column[0], dtype=(bool if data_type == 0 else object)
+            )
+        else:
+            return np.array(column * position_count, dtype=object)
+    else:
+        res = bytearray()
+        for _ in range(position_count):
+            res.extend(column if isinstance(column, bytes) else bytes(column))
+        return bytes(res)
+
+
+def read_dictionary_column(buffer, data_type, position_count):
+    raise Exception("dictionary column not implemented")
+
+
+ENCODING_FUNC_MAP = {
+    0: read_byte_column,
+    1: read_int32_column,
+    2: read_int64_column,
+    3: read_binary_column,
+    4: read_run_length_column,
+    5: read_dictionary_column,
+}
+
+
+def read_column(encoding, buffer, data_type, position_count):
+    try:
+        func = ENCODING_FUNC_MAP[encoding]
+    except KeyError:
+        raise Exception("Unsupported encoding: " + str(encoding))
+    return func(buffer, data_type, position_count)

iotdb/utils/Field.py

@@ -18,18 +18,21 @@
 
 # for package
 from iotdb.utils.IoTDBConstants import TSDataType
-from iotdb.tsfile.utils.DateUtils import parse_int_to_date
+from iotdb.tsfile.utils.date_utils import parse_int_to_date
+from iotdb.utils.rpc_utils import convert_to_timestamp, isoformat
 import numpy as np
 import pandas as pd
 
 
 class Field(object):
-    def __init__(self, data_type, value=None):
+    def __init__(self, data_type, value=None, timezone=None, precision=None):
         """
         :param data_type: TSDataType
         """
         self.__data_type = data_type
         self.value = value
+        self.__timezone = timezone
+        self.__precision = precision
 
     @staticmethod
     def copy(field):
@@ -157,6 +160,17 @@ class Field(object):
             return None
         return self.value
 
+    def get_timestamp_value(self):
+        if self.__data_type is None:
+            raise Exception("Null Field Exception!")
+        if (
+            self.__data_type != TSDataType.TIMESTAMP
+            or self.value is None
+            or self.value is pd.NA
+        ):
+            return None
+        return convert_to_timestamp(self.value, self.__precision, self.__timezone)
+
     def get_date_value(self):
         if self.__data_type is None:
             raise Exception("Null Field Exception!")
@@ -172,11 +186,18 @@ class Field(object):
         if self.__data_type is None or self.value is None or self.value is pd.NA:
             return "None"
         # TEXT, STRING
-        elif self.__data_type == 5 or self.__data_type == 11:
+        if self.__data_type == 5 or self.__data_type == 11:
             return self.value.decode("utf-8")
         # BLOB
         elif self.__data_type == 10:
             return str(hex(int.from_bytes(self.value, byteorder="big")))
+        # TIMESTAMP
+        elif self.__data_type == 8:
+            return isoformat(
+                convert_to_timestamp(self.value, self.__precision, self.__timezone),
+                self.__precision,
+            )
+        # Others
         else:
             return str(self.get_object_value(self.__data_type))
 
@@ -193,15 +214,19 @@ class Field(object):
             return bool(self.value)
         elif data_type == 1:
             return np.int32(self.value)
-        elif data_type == 2 or data_type == 8:
+        elif data_type == 2:
             return np.int64(self.value)
         elif data_type == 3:
             return np.float32(self.value)
         elif data_type == 4:
             return np.float64(self.value)
+        elif data_type == 8:
+            return convert_to_timestamp(self.value, self.__precision, self.__timezone)
         elif data_type == 9:
             return parse_int_to_date(self.value)
-        elif data_type == 5 or data_type == 10 or data_type == 11:
+        elif data_type == 5 or data_type == 11:
+            return self.value.decode("utf-8")
+        elif data_type == 10:
             return self.value
         else:
             raise RuntimeError("Unsupported data type:" + str(data_type))

iotdb/utils/NumpyTablet.py

@@ -22,7 +22,7 @@ import numpy as np
 from numpy import ndarray
 from typing import List
 
-from iotdb.tsfile.utils.DateUtils import parse_date_to_int
+from iotdb.tsfile.utils.date_utils import parse_date_to_int
 from iotdb.utils.IoTDBConstants import TSDataType
 from iotdb.utils.BitMap import BitMap
 from iotdb.utils.Tablet import ColumnType

iotdb/utils/SessionDataSet.py

@@ -21,7 +21,7 @@ from iotdb.utils.Field import Field
 
 # for package
 from iotdb.utils.IoTDBConstants import TSDataType
-from iotdb.utils.IoTDBRpcDataSet import IoTDBRpcDataSet
+from iotdb.utils.iotdb_rpc_dataset import IoTDBRpcDataSet
 from iotdb.utils.RowRecord import RowRecord
 
 import pandas as pd
@@ -37,40 +37,51 @@ class SessionDataSet(object):
         column_type_list,
         column_name_index,
         query_id,
-        client,
         statement_id,
+        client,
         session_id,
-        query_data_set,
+        query_result,
         ignore_timestamp,
+        time_out,
+        more_data,
+        fetch_size,
+        zone_id,
+        time_precision,
+        column_index_2_tsblock_column_index_list,
     ):
         self.iotdb_rpc_data_set = IoTDBRpcDataSet(
             sql,
             column_name_list,
             column_type_list,
-            column_name_index,
             ignore_timestamp,
+            more_data,
             query_id,
             client,
             statement_id,
             session_id,
-            query_data_set,
-            5000,
-        )
-        self.column_size = self.iotdb_rpc_data_set.column_size
-        self.is_ignore_timestamp = self.iotdb_rpc_data_set.ignore_timestamp
-        self.column_names = tuple(self.iotdb_rpc_data_set.get_column_names())
-        self.column_ordinal_dict = self.iotdb_rpc_data_set.column_ordinal_dict
-        self.column_type_deduplicated_list = tuple(
-            self.iotdb_rpc_data_set.column_type_deduplicated_list
+            query_result,
+            fetch_size,
+            time_out,
+            zone_id,
+            time_precision,
+            column_index_2_tsblock_column_index_list,
         )
-        if self.is_ignore_timestamp:
+        if ignore_timestamp:
             self.__field_list = [
-                Field(data_type)
+                (
+                    Field(data_type, timezone=zone_id, precision=time_precision)
+                    if data_type == 8
+                    else Field(data_type)
+                )
                 for data_type in self.iotdb_rpc_data_set.get_column_types()
             ]
         else:
             self.__field_list = [
-                Field(data_type)
+                (
+                    Field(data_type, timezone=zone_id, precision=time_precision)
+                    if data_type == 8
+                    else Field(data_type)
+                )
                 for data_type in self.iotdb_rpc_data_set.get_column_types()[1:]
             ]
         self.row_index = 0
@@ -105,14 +116,22 @@ class SessionDataSet(object):
     def construct_row_record_from_data_frame(self):
         df = self.iotdb_rpc_data_set.data_frame
         row = df.iloc[self.row_index].to_list()
-        row_values = row[1:]
-        for i, value in enumerate(row_values):
-            self.__field_list[i].value = value
+        if self.iotdb_rpc_data_set.ignore_timestamp:
+            for field, value in zip(self.__field_list, row):
+                field.value = value
+
+            row_record = RowRecord(
+                0,
+                self.__field_list,
+            )
+        else:
+            for field, value in zip(self.__field_list, row[1:]):
+                field.value = value
 
-        row_record = RowRecord(
-            row[0],
-            self.__field_list,
-        )
+            row_record = RowRecord(
+                row[0],
+                self.__field_list,
+            )
         self.row_index += 1
         if self.row_index == len(df):
             self.row_index = 0
@@ -147,7 +166,7 @@ def get_typed_point(field: Field, none_value=None):
         TSDataType.INT32: lambda f: f.get_int_value(),
         TSDataType.DOUBLE: lambda f: f.get_double_value(),
         TSDataType.INT64: lambda f: f.get_long_value(),
-        TSDataType.TIMESTAMP: lambda f: f.get_long_value(),
+        TSDataType.TIMESTAMP: lambda f: f.get_timestamp_value(),
         TSDataType.STRING: lambda f: f.get_string_value(),
         TSDataType.DATE: lambda f: f.get_date_value(),
         TSDataType.BLOB: lambda f: f.get_binary_value(),

iotdb/utils/Tablet.py

@@ -20,7 +20,7 @@ import struct
 from enum import unique, IntEnum
 from typing import List, Union
 
-from iotdb.tsfile.utils.DateUtils import parse_date_to_int
+from iotdb.tsfile.utils.date_utils import parse_date_to_int
 from iotdb.utils.BitMap import BitMap
 from iotdb.utils.IoTDBConstants import TSDataType
 

iotdb/utils/{IoTDBConnectionException.py → exception.py}

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+from iotdb.thrift.common.ttypes import TEndPoint, TSStatus
 
 
 class IoTDBConnectionException(Exception):
@@ -27,3 +28,22 @@ class IoTDBConnectionException(Exception):
             super().__init__(message, cause)
         else:
             super().__init__()
+
+
+class StatementExecutionException(Exception):
+    def __init__(self, status: TSStatus = None, message=None):
+        if status is not None:
+            super().__init__(f"{status.code}: {status.message}")
+        elif message is not None:
+            super().__init__(message)
+        else:
+            super().__init__()
+
+
+class RedirectException(Exception):
+    def __init__(self, redirect_info):
+        Exception.__init__(self)
+        if isinstance(redirect_info, TEndPoint):
+            self.redirect_node = redirect_info
+        else:
+            self.device_to_endpoint = redirect_info