| # 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. |
| # |
| |
| #cython: language_level=3 |
| from libc.string cimport strcpy |
| from libc.stdlib cimport malloc, free |
| import pandas as pd |
| from cpython.bytes cimport PyBytes_AsString |
| cimport numpy as cnp |
| import numpy as np |
| from .tsfile cimport * |
| |
| TIMESTAMP_STR = "Time" |
| TS_TYPE_INT32 = 1 << 8 |
| TS_TYPE_BOOLEAN = 1 << 9 |
| TS_TYPE_FLOAT = 1 << 10 |
| TS_TYPE_DOUBLE = 1 << 11 |
| TS_TYPE_INT64 = 1 << 12 |
| TS_TYPE_TEXT = 1 << 13 |
| |
| type_mapping = { |
| 'int32': TS_TYPE_INT32, |
| 'bool': TS_TYPE_BOOLEAN, |
| 'float32': TS_TYPE_FLOAT, |
| 'float64': TS_TYPE_DOUBLE, |
| 'int64': TS_TYPE_INT64 |
| } |
| |
| cdef class tsfile_reader: |
| |
| cdef CTsFileReader reader |
| cdef QueryDataRet ret |
| cdef int batch_size |
| cdef bint read_all_at_once |
| |
| |
| |
| def __init__(self, pathname, table_name, columns, start_time=None, end_time=None, batch_size=None): |
| self.open_reader(pathname) |
| self.query_data_ret(table_name, columns, start_time, end_time) |
| |
| if batch_size is not None: |
| self.batch_size = batch_size |
| self.read_all_at_once = False |
| else: |
| self.batch_size = 1024 |
| self.read_all_at_once = True |
| |
| cdef open_reader(self, pathname): |
| cdef ErrorCode err_code |
| err_code = 0 |
| self.reader = ts_reader_open(pathname.encode('utf-8'), &err_code) |
| if (err_code != 0): |
| raise Exception("Failed to open tsfile: %s, %s" %( pathname, err_code)) |
| |
| cdef query_data_ret(self, table_name, columns, start_time = None, end_time=None): |
| cdef bytes py_table_name |
| cdef char** c_columns |
| py_table_name = table_name.encode('utf-8') |
| c_table_name = PyBytes_AsString(py_table_name) |
| if isinstance(columns, str): |
| columns = [columns] |
| |
| c_columns = <char**>malloc(len(columns) * sizeof(char*)) |
| if not c_columns: |
| raise MemoryError("Failed to allocate memory for columns") |
| |
| for i in range(len(columns)): |
| c_columns[i] = <char*>malloc(len(columns[i]) + 1) |
| if not c_columns[i]: |
| for j in range(i): |
| free(c_columns[j]) |
| free(c_columns) |
| raise MemoryError("Failed to allocate memory for columns") |
| column_binary = columns[i].encode('utf-8') |
| column = PyBytes_AsString(column_binary) |
| strcpy(c_columns[i], column) |
| |
| # query data from tsfile |
| if start_time is not None or end_time is not None: |
| if start_time is None: |
| start_time = LLONG_MIN |
| if end_time is None: |
| end_time = LLONG_MAX |
| self.ret = ts_reader_begin_end(self.reader, c_table_name, c_columns, len(columns), start_time, end_time) |
| else: |
| self.ret = ts_reader_read(self.reader, table_name.encode('utf-8'), c_columns, len(columns)) |
| |
| for i in range(len(columns)): |
| free(c_columns[i]) |
| free(c_columns) |
| |
| |
| def read_tsfile(self): |
| # open tsfile to read |
| res = pd.DataFrame() |
| not_null_maps = [] |
| if self.read_all_at_once: |
| while True: |
| chunk, not_null_map = self.get_next_dataframe() |
| if chunk is not None: |
| res = pd.concat([res, chunk]) |
| not_null_maps.append(not_null_map) |
| else: |
| break |
| else: |
| res, not_null_map = self.get_next_dataframe() |
| not_null_maps.append(not_null_map) |
| |
| self.free_resources() |
| not_null_map_all = None |
| if (not_null_maps != []): |
| not_null_map_all = np.vstack(not_null_maps) |
| return res, not_null_map_all |
| |
| def __iter__(self): |
| return self |
| |
| def __next__(self): |
| res, not_null_map = self.get_next_dataframe() |
| if res is None: |
| raise StopIteration |
| return res, not_null_map |
| |
| def get_next_dataframe(self): |
| cdef: |
| DataResult* result |
| ColumnSchema* schema = NULL |
| cnp.ndarray[cnp.int64_t, ndim=1, mode='c'] np_array_i64 |
| cnp.ndarray[cnp.int32_t, ndim=1, mode='c'] np_array_i32 |
| cnp.ndarray[cnp.float32_t, ndim=1, mode='c'] np_array_float |
| cnp.ndarray[cnp.float64_t, ndim=1, mode='c'] np_array_double |
| cnp.ndarray[bint, ndim=1, mode='c'] np_array_bool |
| cnp.npy_intp length |
| bint has_null |
| bytes pystr |
| str py_string |
| |
| res = {} |
| column_order = [] |
| not_null_map = [] |
| |
| # Time column will be the first column |
| column_order.append(TIMESTAMP_STR) |
| |
| for i in range(self.ret.column_num): |
| pystr = self.ret.column_names[i] |
| py_string = pystr.decode('utf-8', 'ignore') |
| column_order.append(py_string) |
| res[py_string] = [] |
| |
| res[TIMESTAMP_STR] = [] |
| |
| if self.ret.data == NULL: |
| return None, None |
| |
| result = ts_next(self.ret, self.batch_size) |
| |
| # there is no data meet our requirement |
| if result.column_schema == NULL: |
| # free memory |
| if (destory_tablet(result) != 0): |
| raise Exception("Failed to destroy tablet") |
| return None, None |
| |
| # time column |
| length = result.cur_num + 1 |
| cdef cnp.ndarray[cnp.int64_t, ndim=1, mode='c'] data_array = \ |
| cnp.PyArray_SimpleNewFromData(1, &length, cnp.NPY_INT64, result.times) |
| res[TIMESTAMP_STR] = np.array(data_array, dtype = np.int64) |
| |
| for i in range(result.column_num): |
| |
| # column name |
| schema = result.column_schema[i] |
| pystr = schema.name |
| column_name = pystr.decode('utf-8') |
| |
| # column bitmap |
| is_not_null = np.empty(length, dtype = bool) |
| bool_ptr = <char*> result.bitmap[i] |
| has_null = False |
| for j in range(length): |
| is_not_null[j] = bool_ptr[j] != 0 |
| if bool_ptr[j] == 0 and ~has_null: |
| has_null = True |
| |
| |
| if schema.column_def == TS_TYPE_INT32: |
| np_array_i32 = cnp.PyArray_SimpleNewFromData(1, &length, cnp.NPY_INT32, result.value[i]) |
| arr = np.array(np_array_i32, dtype = np.int32) |
| |
| elif schema.column_def == TS_TYPE_BOOLEAN: |
| arr_bool_ = np.empty(length, dtype=np.bool_) |
| bool_ptr = <char*> result.value[i] |
| for j in range(length): |
| arr_bool_[j] = bool_ptr[j] != 0 |
| arr = np.array(arr_bool_, dtype = np.bool_) |
| |
| elif schema.column_def == TS_TYPE_FLOAT: |
| np_array_float = cnp.PyArray_SimpleNewFromData(1, &length, cnp.NPY_FLOAT32, result.value[i]) |
| arr = np.array(np_array_float, dtype = np.float32) |
| arr = np.where(is_not_null, arr, np.nan) |
| res[column_name]=arr |
| continue |
| |
| elif schema.column_def == TS_TYPE_DOUBLE: |
| np_array_double = cnp.PyArray_SimpleNewFromData(1, &length, cnp.NPY_FLOAT64, result.value[i]) |
| arr= np.array(np_array_double, dtype = np.float64) |
| arr = np.where(is_not_null, arr, np.nan) |
| res[column_name]=arr |
| continue |
| |
| elif schema.column_def == TS_TYPE_INT64: |
| np_array_i64 = cnp.PyArray_SimpleNewFromData(1, &length, cnp.NPY_INT64, result.value[i]) |
| arr = np.array(np_array_i64, dtype = np.int64) |
| else: |
| raise Exception("UnSupport column type") |
| |
| if has_null: |
| tmp_array = np.full(length, np.nan, np.float64) |
| tmp_array[is_not_null] = arr[is_not_null] |
| if schema.column_def == TS_TYPE_INT32: |
| arr = pd.Series(tmp_array).astype('Int32') |
| elif schema.column_def == TS_TYPE_BOOLEAN: |
| arr = pd.Series(tmp_array).astype(np.bool_) |
| elif schema.column_def == TS_TYPE_INT64: |
| arr = pd.Series(tmp_array).astype('Int64') |
| |
| res[column_name] = arr |
| not_null_map.append(is_not_null) |
| if (destory_tablet(result) != 0): |
| raise Exception("Failed to destroy tablet") |
| return pd.DataFrame(res, columns = column_order), not_null_map |
| |
| def __dealloc__(self): |
| self.free_resources() |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, exc_type, exc_value, traceback): |
| self.free_resources() |
| |
| cdef free_resources(self): |
| if self.reader: |
| if ts_reader_close(self.reader) != 0 : |
| raise Exception("Failed to close tsfile") |
| if self.ret: |
| if destory_query_dataret(self.ret) != 0: |
| raise Exception("Failed to free query data ret") |
| self.reader = NULL |
| self.ret = NULL |
| |
| cdef class tsfile_writer: |
| cdef CTsFileWriter writer |
| cdef TsFileRowData row_data |
| |
| def __init__(self, pathname): |
| self.open_writer(pathname) |
| |
| |
| cdef open_writer(self, pathname): |
| cdef ErrorCode err_code |
| err_code = 0 |
| self.writer = ts_writer_open(pathname.encode('utf-8'), &err_code) |
| if (err_code != 0): |
| raise Exception("Failed to open tsfile: %s, %s" %( pathname, err_code)) |
| |
| def resister_timeseries(self, table_name, column_name, data_type): |
| cdef char* c_columns |
| cdef bytes py_table_name |
| cdef ColumnSchema schema |
| cdef bytes encoded_column_name = column_name.encode('utf-8') |
| py_table_name = table_name.encode('utf-8') |
| c_table_name = PyBytes_AsString(py_table_name) |
| schema.name = encoded_column_name |
| schema.column_def = data_type |
| if tsfile_register_table_column(self.writer, c_table_name, &schema) != 0: |
| raise Exception("Failed to register timeseries") |
| cdef create_row_data(self, table_name, time, column_length): |
| self.row_data = create_tsfile_row(table_name.encode('utf-8'), time, column_length) |
| def write_into_row_data(self, column_name, value, type): |
| cdef char* c_column_name = PyBytes_AsString(column_name.encode('utf-8')) |
| if type == TS_TYPE_INT32: |
| insert_data_into_tsfile_row_int32(self.row_data, c_column_name, value) |
| elif type == TS_TYPE_BOOLEAN: |
| insert_data_into_tsfile_row_boolean(self.row_data, c_column_name, value) |
| elif type == TS_TYPE_FLOAT: |
| insert_data_into_tsfile_row_float(self.row_data, c_column_name, value) |
| elif type == TS_TYPE_DOUBLE: |
| insert_data_into_tsfile_row_double(self.row_data, c_column_name, value) |
| elif type == TS_TYPE_INT64: |
| insert_data_into_tsfile_row_int64(self.row_data, c_column_name, value) |
| else: |
| raise TypeError("Unknown column type") |
| def write_tsfile(self, table_name, df): |
| column_names = df.columns.tolist() |
| column_types = df.dtypes |
| column_ctypes = [] |
| for i in range(len(column_names)): |
| column_type = column_types[i].name |
| if column_type in type_mapping: |
| column_ctypes.append(type_mapping[column_type]) |
| else: |
| raise TypeError("Unknown column type") |
| |
| if (column_names[i] != TIMESTAMP_STR): |
| self.resister_timeseries(table_name, column_names[i], column_ctypes[i]) |
| |
| |
| for i in range(len(df)): |
| time = df.iloc[i][TIMESTAMP_STR] |
| self.create_row_data(table_name, time, len(column_names)) |
| for j in range(1, len(column_names)): |
| column_name = column_names[j] |
| column_value = df.iloc[i][column_name] |
| column_ctype = column_ctypes[j] |
| self.write_into_row_data(column_name, column_value, column_ctype) |
| if tsfile_write_row_data(self.writer, self.row_data) != 0: |
| raise Exception("Failed to write row data") |
| |
| if tsfile_flush_data(self.writer) != 0: |
| raise Exception("Failed to flush data") |
| self.row_data = NULL |
| self.free_resources() |
| def free_resources(self): |
| if self.writer != NULL: |
| if ts_writer_close(self.writer) != 0: |
| raise Exception("Failed to close tsfile") |
| self.writer = NULL |
| def __dealloc__(self): |
| self.free_resources() |
