| # 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 os |
| from datetime import date |
| |
| import numpy as np |
| import pandas as pd |
| import pytest |
| |
| from tsfile import to_dataframe, TsFileReader, ColumnCategory, TIME_COLUMN |
| from tsfile.utils import dataframe_to_tsfile |
| |
| |
| def convert_to_nullable_types(df): |
| df = df.copy() |
| for col in df.columns: |
| dtype = df[col].dtype |
| if dtype == "int64": |
| df[col] = df[col].astype("Int64") |
| elif dtype == "int32": |
| df[col] = df[col].astype("Int32") |
| elif dtype == "float64": |
| df[col] = df[col].astype("Float64") |
| elif dtype == "float32": |
| df[col] = df[col].astype("Float32") |
| elif dtype == "bool": |
| df[col] = df[col].astype("boolean") |
| elif pd.api.types.is_object_dtype(df[col]): |
| non_null = df[col].dropna() |
| if len(non_null) and non_null.map(lambda x: isinstance(x, str)).all(): |
| df[col] = df[col].astype("string") |
| elif pd.api.types.is_string_dtype(df[col]): |
| df[col] = df[col].astype("string") |
| return df |
| |
| |
| def test_dataframe_to_tsfile_basic(): |
| tsfile_path = "test_dataframe_to_tsfile_basic.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "time": [i for i in range(100)], |
| "device": [f"device{i}" for i in range(100)], |
| "value": [i * 1.5 for i in range(100)], |
| "value2": [i * 10 for i in range(100)], |
| } |
| ) |
| |
| dataframe_to_tsfile(df, tsfile_path, table_name="test_table") |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = df_read.sort_values("time").reset_index(drop=True) |
| df_sorted = convert_to_nullable_types( |
| df.sort_values("time").reset_index(drop=True) |
| ) |
| df_read = convert_to_nullable_types(df_read) |
| |
| assert df_read.shape == (100, 4) |
| assert df_read["time"].equals(df_sorted["time"]) |
| assert df_read["device"].equals(df_sorted["device"]) |
| assert df_read["value"].equals(df_sorted["value"]) |
| assert df_read["value2"].equals(df_sorted["value2"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_default_table_name(): |
| tsfile_path = "test_dataframe_to_tsfile_default.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame({"time": [0, 1], "value": [1.0, 2.0]}) |
| dataframe_to_tsfile(df, tsfile_path) |
| |
| df_read = to_dataframe(tsfile_path, table_name="default_table") |
| assert len(df_read) == 2 |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_with_index(): |
| tsfile_path = "test_dataframe_to_tsfile_index.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "device": [f"device{i}" for i in range(30)], |
| "value": [i * 2.0 for i in range(30)], |
| } |
| ) |
| df.index = [i * 100 for i in range(30)] |
| dataframe_to_tsfile(df, tsfile_path, table_name="test_table") |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = df_read.sort_values("time").reset_index(drop=True) |
| df_read = convert_to_nullable_types(df_read) |
| time_expected = pd.Series(df.index.values, dtype="Int64") |
| assert df_read.shape == (30, 3) |
| assert df_read["time"].equals(time_expected) |
| |
| with TsFileReader(tsfile_path) as reader: |
| table_schema = reader.get_table_schema("test_table") |
| device_col = table_schema.get_column("device") |
| assert device_col is not None |
| assert device_col.get_category() == ColumnCategory.FIELD |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_custom_time_column(): |
| tsfile_path = "test_dataframe_to_tsfile_custom_time.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "timestamp": [i for i in range(30)], |
| "device": [f"device{i}" for i in range(30)], |
| "value": [i * 3.0 for i in range(30)], |
| } |
| ) |
| |
| dataframe_to_tsfile( |
| df, tsfile_path, table_name="test_table", time_column="timestamp" |
| ) |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = df_read.sort_values("timestamp").reset_index(drop=True) |
| df_sorted = convert_to_nullable_types( |
| df.sort_values("timestamp").reset_index(drop=True) |
| ) |
| df_read = convert_to_nullable_types(df_read) |
| |
| assert df_read.shape == (30, 3) |
| assert df_read["timestamp"].equals(df_sorted["timestamp"]) |
| assert df_read["device"].equals(df_sorted["device"]) |
| assert df_read["value"].equals(df_sorted["value"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_case_insensitive_time(): |
| tsfile_path = "test_dataframe_to_tsfile_case_time.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| {"Time": [i for i in range(20)], "value": [i * 2.0 for i in range(20)]} |
| ) |
| |
| dataframe_to_tsfile(df, tsfile_path, table_name="test_table") |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = convert_to_nullable_types(df_read) |
| assert df_read.shape == (20, 2) |
| assert df_read["time"].equals(pd.Series([i for i in range(20)], dtype="Int64")) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_with_tag_columns(): |
| tsfile_path = "test_dataframe_to_tsfile_tags.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "time": [i for i in range(20)], |
| "device": [f"device{i}" for i in range(20)], |
| "location": [f"loc{i % 5}" for i in range(20)], |
| "value": [i * 1.5 for i in range(20)], |
| } |
| ) |
| |
| dataframe_to_tsfile( |
| df, tsfile_path, table_name="test_table", tag_column=["device", "location"] |
| ) |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True) |
| df_sorted = convert_to_nullable_types( |
| df.sort_values("time").reset_index(drop=True) |
| ) |
| df_read = convert_to_nullable_types(df_read) |
| |
| assert df_read.shape == (20, 4) |
| assert df_read["device"].equals(df_sorted["device"]) |
| assert df_read["location"].equals(df_sorted["location"]) |
| assert df_read["value"].equals(df_sorted["value"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_tag_time_unsorted(): |
| tsfile_path = "test_dataframe_to_tsfile_tag_time_unsorted.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "time": [30, 10, 20, 50, 40, 15, 25, 35, 5, 45], |
| "device": [ |
| "device1", |
| "device1", |
| "device1", |
| "device2", |
| "device2", |
| "device1", |
| "device1", |
| "device2", |
| "device1", |
| "device2", |
| ], |
| "value": [i * 1.5 for i in range(10)], |
| } |
| ) |
| |
| dataframe_to_tsfile( |
| df, tsfile_path, table_name="test_table", tag_column=["device"] |
| ) |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_expected = df.sort_values(by=["device", "time"]).reset_index(drop=True) |
| df_expected = convert_to_nullable_types(df_expected) |
| df_read = convert_to_nullable_types(df_read) |
| |
| assert df_read.shape == (10, 3) |
| assert df_read["device"].equals(df_expected["device"]) |
| assert df_read[TIME_COLUMN].equals(df_expected["time"]) |
| assert df_read["value"].equals(df_expected["value"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_all_datatypes(): |
| tsfile_path = "test_dataframe_to_tsfile_all_types.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "time": [i for i in range(50)], |
| "bool_col": [i % 2 == 0 for i in range(50)], |
| "int32_col": pd.Series([i for i in range(50)], dtype="int32"), |
| "int64_col": [i * 10 for i in range(50)], |
| "float_col": pd.Series([i * 1.5 for i in range(50)], dtype="float32"), |
| "double_col": [i * 2.5 for i in range(50)], |
| "string_col": [f"str{i}" for i in range(50)], |
| "blob_col": [f"blob{i}".encode("utf-8") for i in range(50)], |
| "text_col": [f"text{i}" for i in range(50)], |
| "date_col": [date(2025, i % 11 + 1, i % 20 + 1) for i in range(50)], |
| "timestamp_col": [i for i in range(50)], |
| } |
| ) |
| |
| dataframe_to_tsfile(df, tsfile_path, table_name="test_table") |
| |
| df_read = to_dataframe(tsfile_path, table_name="test_table") |
| df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True) |
| df_sorted = convert_to_nullable_types( |
| df.sort_values("time").reset_index(drop=True) |
| ) |
| df_read = convert_to_nullable_types(df_read) |
| |
| assert df_read.shape == (50, 11) |
| assert df_read["bool_col"].equals(df_sorted["bool_col"]) |
| assert df_read["int32_col"].equals(df_sorted["int32_col"]) |
| assert df_read["int64_col"].equals(df_sorted["int64_col"]) |
| assert np.allclose(df_read["float_col"], df_sorted["float_col"]) |
| assert np.allclose(df_read["double_col"], df_sorted["double_col"]) |
| assert df_read["string_col"].equals(df_sorted["string_col"]) |
| assert df_read["text_col"].equals(df_sorted["text_col"]) |
| assert df_read["date_col"].equals(df_sorted["date_col"]) |
| assert df_read["timestamp_col"].equals(df_sorted["timestamp_col"]) |
| for i in range(50): |
| assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i] |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_empty_dataframe(): |
| tsfile_path = "test_dataframe_to_tsfile_empty.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame() |
| |
| with pytest.raises(ValueError, match="DataFrame cannot be None or empty"): |
| dataframe_to_tsfile(df, tsfile_path) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_no_data_columns(): |
| tsfile_path = "test_dataframe_to_tsfile_no_data.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame({"time": [i for i in range(10)]}) |
| |
| with pytest.raises( |
| ValueError, match="DataFrame must have at least one data column" |
| ): |
| dataframe_to_tsfile(df, tsfile_path) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_only_time_column_raises(): |
| """Only time column (no FIELD/TAG columns) must raise ValueError.""" |
| tsfile_path = "test_only_time_column.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| df = pd.DataFrame({"time": [1, 2, 3]}) |
| with pytest.raises( |
| ValueError, match="at least one data column besides the time column" |
| ): |
| dataframe_to_tsfile(df, tsfile_path) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_time_column_not_found(): |
| tsfile_path = "test_dataframe_to_tsfile_time_err.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame({"time": [0, 1], "value": [1.0, 2.0]}) |
| with pytest.raises(ValueError, match="Time column 'timestamp' not found"): |
| dataframe_to_tsfile(df, tsfile_path, time_column="timestamp") |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_invalid_time_column(): |
| tsfile_path = "test_dataframe_to_tsfile_invalid_time.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| {"timestamp": [i for i in range(10)], "value": [i * 1.0 for i in range(10)]} |
| ) |
| |
| with pytest.raises(ValueError, match="Time column 'time' not found"): |
| dataframe_to_tsfile(df, tsfile_path, time_column="time") |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_non_integer_time_column(): |
| tsfile_path = "test_dataframe_to_tsfile_non_int_time.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| { |
| "time": [f"time{i}" for i in range(10)], |
| "value": [i * 1.0 for i in range(10)], |
| } |
| ) |
| |
| with pytest.raises(TypeError, match="must be integer type"): |
| dataframe_to_tsfile(df, tsfile_path) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_tag_column_not_found(): |
| tsfile_path = "test_dataframe_to_tsfile_tag_err.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame({"time": [0, 1], "device": ["a", "b"], "value": [1.0, 2.0]}) |
| with pytest.raises(ValueError, match="Tag column 'invalid' not found"): |
| dataframe_to_tsfile(df, tsfile_path, tag_column=["invalid"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_dataframe_to_tsfile_invalid_tag_column(): |
| tsfile_path = "test_dataframe_to_tsfile_invalid_tag.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| df = pd.DataFrame( |
| {"time": [i for i in range(10)], "value": [i * 1.0 for i in range(10)]} |
| ) |
| |
| with pytest.raises(ValueError, match="Tag column 'invalid' not found"): |
| dataframe_to_tsfile(df, tsfile_path, tag_column=["invalid"]) |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
