| # 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. |
| # |
| from typing import Iterator, Union |
| from typing import Optional |
| |
| import numpy as np |
| import pandas as pd |
| from pandas.core.dtypes.common import is_integer_dtype, is_object_dtype |
| from pandas.core.interchange.dataframe_protocol import DataFrame |
| |
| from tsfile import ColumnSchema, TableSchema, ColumnCategory, TSDataType, TIME_COLUMN |
| from tsfile.exceptions import TableNotExistError, ColumnNotExistError |
| from tsfile.tsfile_reader import TsFileReaderPy |
| from tsfile.tsfile_table_writer import ( |
| TsFileTableWriter, |
| infer_object_column_type, |
| validate_dataframe_for_tsfile, |
| ) |
| |
| |
| def to_dataframe( |
| file_path: str, |
| table_name: Optional[str] = None, |
| column_names: Optional[list[str]] = None, |
| start_time: Optional[int] = None, |
| end_time: Optional[int] = None, |
| max_row_num: Optional[int] = None, |
| as_iterator: bool = False, |
| ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]: |
| """ |
| Read data from a TsFile and convert it into a Pandas DataFrame or |
| an iterator of DataFrames. |
| |
| This function supports both table-model and tree-model TsFiles. |
| Users can filter data by table name, column names, time range, |
| and maximum number of rows. |
| |
| Parameters |
| ---------- |
| file_path : str |
| Path to the TsFile to be read. |
| |
| table_name : Optional[str], default None |
| Name of the table to query in table-model TsFiles. |
| If None and the file is in table model, the first table |
| found in the schema will be used. |
| |
| column_names : Optional[list[str]], default None |
| List of column names to query. |
| - If None, all columns will be returned. |
| - Column existence will be validated in table-model TsFiles. |
| |
| start_time : Optional[int], default None |
| Start timestamp for the query. |
| If None, the minimum int64 value is used. |
| |
| end_time : Optional[int], default None |
| End timestamp for the query. |
| If None, the maximum int64 value is used. |
| |
| max_row_num : Optional[int], default None |
| Maximum number of rows to read. |
| - If None, all available rows will be returned. |
| - When `as_iterator` is False, the final DataFrame will be |
| truncated to this size if necessary. |
| |
| as_iterator : bool, default False |
| Whether to return an iterator of DataFrames instead of |
| a single concatenated DataFrame. |
| - True: returns an iterator yielding DataFrames in batches |
| - False: returns a single Pandas DataFrame |
| |
| Returns |
| ------- |
| Union[pandas.DataFrame, Iterator[pandas.DataFrame]] |
| - A Pandas DataFrame if `as_iterator` is False |
| - An iterator of Pandas DataFrames if `as_iterator` is True |
| |
| Raises |
| ------ |
| TableNotExistError |
| If the specified table name does not exist in a table-model TsFile. |
| |
| ColumnNotExistError |
| If any specified column does not exist in the table schema. |
| """ |
| |
| def _gen(is_iterator: bool) -> Iterator[pd.DataFrame]: |
| _table_name = table_name |
| _column_names = column_names |
| _start_time = start_time if start_time is not None else np.iinfo(np.int64).min |
| _end_time = end_time if end_time is not None else np.iinfo(np.int64).max |
| |
| ## Time column handling (table model): |
| ## 1. Request has no column list (query all): |
| ## 1.1 TsFile has a time column in schema: query only non-time columns; then rename |
| ## the first column of the returned DataFrame to the schema time column name. |
| ## 1.2 TsFile has no time column in schema: query as-is; first column is "time". |
| ## 2. Request has a column list but no time column: |
| ## 2.1 TsFile has a time column in schema: query with requested columns; rename the |
| ## first column to the schema time column name. |
| ## 2.2 TsFile has no time column in schema: first column stays "time"; no rename. |
| ## 3. Request has a column list including the time column: |
| ## 3.1 Query with requested columns (including time); do not rename the first column. |
| with TsFileReaderPy(file_path) as reader: |
| total_rows = 0 |
| table_schema = reader.get_all_table_schemas() |
| |
| is_tree_model = len(table_schema) == 0 |
| time_column = None |
| column_name_to_query = [] |
| no_field_query = True |
| if is_tree_model: |
| if _column_names is None: |
| print("columns name is None, return all columns") |
| # When querying tables in the tree, only measurements are allowed currently. |
| no_field_query = False |
| else: |
| _table_name = _table_name.lower() if _table_name else None |
| _column_names = ( |
| [column.lower() for column in _column_names] |
| if _column_names |
| else None |
| ) |
| if _table_name is None: |
| _table_name, table_schema = next(iter(table_schema.items())) |
| else: |
| _table_name = _table_name.lower() |
| if _table_name.lower() not in table_schema: |
| raise TableNotExistError(_table_name) |
| table_schema = table_schema[_table_name] |
| |
| column_names_in_file = [] |
| for column in table_schema.get_columns(): |
| if column.get_category() == ColumnCategory.TIME: |
| time_column = column.get_column_name() |
| else: |
| column_names_in_file.append(column.get_column_name()) |
| |
| if _column_names is not None: |
| for column in _column_names: |
| if column not in column_names_in_file and column != time_column: |
| raise ColumnNotExistError(column) |
| if ( |
| table_schema.get_column(column).get_category() |
| == ColumnCategory.FIELD |
| ): |
| no_field_query = False |
| if no_field_query: |
| if time_column is not None: |
| column_name_to_query.append(time_column) |
| column_name_to_query.extend(column_names_in_file) |
| else: |
| column_name_to_query = _column_names |
| else: |
| no_field_query = False |
| column_name_to_query = column_names_in_file |
| |
| if is_tree_model: |
| if _column_names is not None: |
| column_name_to_query = _column_names |
| query_result = reader.query_table_on_tree( |
| column_name_to_query, _start_time, _end_time |
| ) |
| else: |
| query_result = reader.query_table( |
| _table_name, column_name_to_query, _start_time, _end_time |
| ) |
| |
| with query_result as result: |
| while result.next(): |
| if max_row_num is None: |
| dataframe = result.read_data_frame() |
| elif is_iterator: |
| dataframe = result.read_data_frame(max_row_num) |
| else: |
| remaining_rows = max_row_num - total_rows |
| if remaining_rows <= 0: |
| break |
| dataframe = result.read_data_frame(remaining_rows) |
| if dataframe is None or dataframe.empty: |
| continue |
| total_rows += len(dataframe) |
| if time_column is not None: |
| if _column_names is None or time_column not in _column_names: |
| dataframe = dataframe.rename( |
| columns={dataframe.columns[0]: time_column} |
| ) |
| if no_field_query and _column_names is not None: |
| _column_names.insert(0, TIME_COLUMN) |
| dataframe = dataframe[_column_names] |
| yield dataframe |
| if ( |
| (not is_iterator) |
| and max_row_num is not None |
| and total_rows >= max_row_num |
| ): |
| break |
| |
| if as_iterator: |
| return _gen(True) |
| else: |
| df_list = list(_gen(False)) |
| if df_list: |
| df = pd.concat(df_list, ignore_index=True) |
| if max_row_num is not None and len(df) > max_row_num: |
| df = df.iloc[:max_row_num] |
| return df |
| else: |
| return pd.DataFrame() |
| |
| |
| def dataframe_to_tsfile( |
| dataframe: pd.DataFrame, |
| file_path: str, |
| table_name: Optional[str] = None, |
| time_column: Optional[str] = None, |
| tag_column: Optional[list[str]] = None, |
| ): |
| """ |
| Write a pandas DataFrame to a TsFile by inferring the table schema from the DataFrame. |
| |
| This function automatically infers the table schema based on the DataFrame's column |
| names and data types, then writes the data to a TsFile. |
| |
| Parameters |
| ---------- |
| dataframe : pd.DataFrame |
| The pandas DataFrame to write to TsFile. |
| - If a 'time' column (case-insensitive) exists, it will be used as the time column. |
| - Otherwise, the DataFrame index will be used as timestamps. |
| - All other columns will be treated as data columns. |
| |
| file_path : str |
| Path to the TsFile to write. Will be created if it doesn't exist. |
| |
| table_name : Optional[str], default None |
| Name of the table. If None, defaults to "default_table". |
| |
| time_column : Optional[str], default None |
| Name of the time column. If None, will look for a column named 'time' (case-insensitive), |
| or use the DataFrame index if no 'time' column is found. |
| |
| tag_column : Optional[list[str]], default None |
| List of column names to be treated as TAG columns. All other columns will be FIELD columns. |
| If None, all columns are treated as FIELD columns. |
| |
| Returns |
| ------- |
| None |
| |
| Raises |
| ------ |
| ValueError |
| If the DataFrame is empty or has no data columns. |
| """ |
| validate_dataframe_for_tsfile(dataframe) |
| df = dataframe.rename(columns=str.lower) |
| |
| if not table_name: |
| table_name = "default_table" |
| |
| if time_column is not None: |
| if time_column.lower() not in df.columns: |
| raise ValueError(f"Time column '{time_column}' not found in DataFrame") |
| if tag_column is not None: |
| for tag_col in tag_column: |
| if tag_col.lower() not in df.columns: |
| raise ValueError(f"Tag column '{tag_col}' not found in DataFrame") |
| tag_columns_lower = {t.lower() for t in (tag_column or [])} |
| |
| if time_column is not None: |
| time_col_name = time_column.lower() |
| elif "time" in df.columns: |
| time_col_name = "time" |
| else: |
| time_col_name = None |
| |
| if time_col_name is not None: |
| if not is_integer_dtype(df[time_col_name].dtype): |
| raise TypeError( |
| f"Time column '{time_col_name}' must be integer type (int64 or int), got {df[time_col_name].dtype}" |
| ) |
| |
| column_schemas = [] |
| if time_col_name is not None: |
| column_schemas.append( |
| ColumnSchema(time_col_name, TSDataType.TIMESTAMP, ColumnCategory.TIME) |
| ) |
| |
| for col in df.columns: |
| if col == time_col_name: |
| continue |
| col_dtype = df[col].dtype |
| if is_object_dtype(col_dtype): |
| ts_data_type = infer_object_column_type(df[col]) |
| else: |
| ts_data_type = TSDataType.from_pandas_datatype(col_dtype) |
| |
| category = ( |
| ColumnCategory.TAG if col in tag_columns_lower else ColumnCategory.FIELD |
| ) |
| column_schemas.append(ColumnSchema(col, ts_data_type, category)) |
| |
| data_columns = [ |
| s for s in column_schemas if s.get_category() != ColumnCategory.TIME |
| ] |
| if len(data_columns) == 0: |
| raise ValueError( |
| "DataFrame must have at least one data column besides the time column" |
| ) |
| |
| table_schema = TableSchema(table_name, column_schemas) |
| |
| with TsFileTableWriter(file_path, table_schema) as writer: |
| writer.write_dataframe(df) |
