| # 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 datetime import date |
| |
| import numpy as np |
| import pandas as pd |
| import pytest |
| from pandas.core.dtypes.common import is_integer_dtype |
| |
| from tsfile import ColumnSchema, TableSchema, TSEncoding |
| from tsfile import Compressor |
| from tsfile import TSDataType |
| from tsfile import Tablet, RowRecord, Field |
| from tsfile import TimeseriesSchema |
| from tsfile import TsFileTableWriter |
| from tsfile import TsFileWriter, TsFileReader, ColumnCategory |
| from tsfile import to_dataframe |
| from tsfile.exceptions import TableNotExistError, ColumnNotExistError, NotSupportedError |
| |
| |
| def test_row_record_write_and_read(): |
| try: |
| if os.path.exists("record_write_and_read.tsfile"): |
| os.remove("record_write_and_read.tsfile") |
| writer = TsFileWriter("record_write_and_read.tsfile") |
| writer.register_timeseries("root.device1", TimeseriesSchema("level1", TSDataType.INT64)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level2", TSDataType.DOUBLE)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level3", TSDataType.INT32)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level4", TSDataType.STRING)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level5", TSDataType.TEXT)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level6", TSDataType.BLOB)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level7", TSDataType.DATE)) |
| writer.register_timeseries("root.device1", TimeseriesSchema("level8", TSDataType.TIMESTAMP)) |
| |
| max_row_num = 10 |
| |
| for i in range(max_row_num): |
| row = RowRecord("root.device1", i, |
| [Field("level1", i + 1, TSDataType.INT64), |
| Field("level2", i * 1.1, TSDataType.DOUBLE), |
| Field("level3", i * 2, TSDataType.INT32), |
| Field("level4", f"string_value_{i}", TSDataType.STRING), |
| Field("level5", f"text_value_{i}", TSDataType.TEXT), |
| Field("level6", f"blob_data_{i}".encode('utf-8'), TSDataType.BLOB), |
| Field("level7", date(2025, 1, i % 20 + 1), TSDataType.DATE), |
| Field("level8", i, TSDataType.TIMESTAMP)]) |
| writer.write_row_record(row) |
| |
| writer.close() |
| |
| reader = TsFileReader("record_write_and_read.tsfile") |
| result = reader.query_timeseries( |
| "root.device1", |
| ["level1", "level2", "level3", "level4", "level5", "level6", "level7", "level8"], |
| 0, |
| 100, |
| ) |
| |
| assert len(reader.get_active_query_result()) == 1 |
| |
| for row_num in range(max_row_num): |
| assert result.next() |
| assert result.get_value_by_index(1) == row_num |
| assert result.get_value_by_index(2) == row_num + 1 |
| assert result.get_value_by_index(3) == pytest.approx(row_num * 1.1) |
| assert result.get_value_by_index(4) == row_num * 2 |
| assert result.get_value_by_index(5) == f"string_value_{row_num}" |
| assert result.get_value_by_index(6) == f"text_value_{row_num}" |
| assert result.get_value_by_index(7) == f"blob_data_{row_num}" |
| assert result.get_value_by_index(8) == date(2025, 1, row_num % 20 + 1) |
| assert result.get_value_by_index(9) == row_num |
| |
| assert not result.next() |
| assert len(reader.get_active_query_result()) == 1 |
| result.close() |
| print(reader.get_active_query_result()) |
| assert len(reader.get_active_query_result()) == 0 |
| reader.close() |
| |
| |
| |
| finally: |
| if os.path.exists("record_write_and_read.tsfile"): |
| os.remove("record_write_and_read.tsfile") |
| |
| |
| def test_tree_query_to_dataframe_variants(): |
| file_path = "tree_query_to_dataframe.tsfile" |
| device_ids = [ |
| "root.db1.t1", |
| "root.db2.t1", |
| "root.db3.t2.t3", |
| "root.db3.t3", |
| "device", |
| "device.ln", |
| "device2.ln1.tmp", |
| "device3.ln2.tmp.v1.v2", |
| "device3.ln2.tmp.v1.v3", |
| ] |
| device_path_map = [ |
| "root.db1.t1.null.null", |
| "root.db2.t1.null.null", |
| "root.db3.t2.t3.null", |
| "root.db3.t3.null.null", |
| "device.null.null.null.null", |
| "device.ln.null.null.null", |
| "device2.ln1.tmp.null.null", |
| "device3.ln2.tmp.v1.v2", |
| "device3.ln2.tmp.v1.v3", |
| ] |
| measurement_ids1 = ["temperature", "hudi", "level"] |
| measurement_ids2 = ["level", "vol"] |
| rows_per_device = 2 |
| expected_values = {} |
| all_measurements = set() |
| |
| def _is_null(value): |
| return value is None or pd.isna(value) |
| |
| def _extract_device(row, path_columns): |
| parts = [] |
| for col in path_columns: |
| value = row[col] |
| if not _is_null(value): |
| parts.append(str(value)) |
| else: |
| parts.append("null") |
| return ".".join(parts) |
| |
| try: |
| writer = TsFileWriter(file_path) |
| for idx, device_id in enumerate(device_ids): |
| measurements = measurement_ids1 if idx % 2 == 0 else measurement_ids2 |
| all_measurements.update(measurements) |
| for measurement in measurements: |
| writer.register_timeseries( |
| device_id, TimeseriesSchema(measurement, TSDataType.INT32) |
| ) |
| for ts in range(rows_per_device): |
| fields = [] |
| measurement_snapshot = {} |
| for m_idx, measurement in enumerate(measurements): |
| value = idx * 100 + ts * 10 + m_idx |
| fields.append(Field(measurement, value, TSDataType.INT32)) |
| measurement_snapshot[measurement] = value |
| writer.write_row_record(RowRecord(device_id, ts, fields)) |
| expected_values[(device_path_map[idx], ts)] = measurement_snapshot |
| writer.close() |
| |
| df_all = to_dataframe(file_path, start_time=0, end_time=rows_per_device) |
| total_rows = len(device_ids) * rows_per_device |
| assert df_all.shape[0] == total_rows |
| for measurement in all_measurements: |
| assert measurement in df_all.columns |
| assert "time" in df_all.columns |
| path_columns = sorted( |
| [col for col in df_all.columns if col.startswith("col_")], |
| key=lambda name: int(name.split("_")[1]), |
| ) |
| assert len(path_columns) > 0 |
| |
| for _, row in df_all.iterrows(): |
| device = _extract_device(row, path_columns) |
| timestamp = int(row["time"]) |
| assert (device, timestamp) in expected_values |
| expected_row = expected_values[(device, timestamp)] |
| for measurement in all_measurements: |
| value = row.get(measurement) |
| if measurement in expected_row: |
| assert value == expected_row[measurement] |
| else: |
| assert _is_null(value) |
| assert device in device_path_map |
| |
| requested_columns = ["level", "temperature"] |
| df_subset = to_dataframe( |
| file_path, column_names=requested_columns, start_time=0, end_time=rows_per_device |
| ) |
| for column in requested_columns: |
| assert column in df_subset.columns |
| for measurement in all_measurements: |
| if measurement not in requested_columns: |
| assert measurement not in df_subset.columns |
| for _, row in df_subset.iterrows(): |
| device = _extract_device(row, path_columns) |
| timestamp = int(row["time"]) |
| expected_row = expected_values[(device, timestamp)] |
| for measurement in requested_columns: |
| value = row.get(measurement) |
| if measurement in expected_row: |
| assert value == expected_row[measurement] |
| else: |
| assert _is_null(value) |
| assert device in device_path_map |
| df_limited = to_dataframe( |
| file_path, column_names=["level"], max_row_num=5, start_time=0, end_time=rows_per_device |
| ) |
| assert df_limited.shape[0] == 5 |
| assert "level" in df_limited.columns |
| |
| iterator = to_dataframe( |
| file_path, |
| column_names=["level", "temperature"], |
| max_row_num=3, |
| start_time=0, |
| end_time=rows_per_device, |
| as_iterator=True, |
| ) |
| iter_rows = 0 |
| for batch in iterator: |
| assert isinstance(batch, pd.DataFrame) |
| assert set(batch.columns).issuperset({"time", "level"}) |
| iter_rows += len(batch) |
| assert iter_rows == 18 |
| |
| iterator = to_dataframe( |
| file_path, |
| column_names=["level", "temperature"], |
| max_row_num=3, |
| start_time=0, |
| end_time=0, |
| as_iterator=True, |
| ) |
| iter_rows = 0 |
| for batch in iterator: |
| assert isinstance(batch, pd.DataFrame) |
| assert set(batch.columns).issuperset({"time", "level"}) |
| iter_rows += len(batch) |
| assert iter_rows == 9 |
| |
| with pytest.raises(ColumnNotExistError): |
| to_dataframe(file_path, column_names=["level", "not_exists"]) |
| finally: |
| if os.path.exists(file_path): |
| os.remove(file_path) |
| |
| |
| def test_get_all_timeseries_schemas(): |
| file_path = "get_all_timeseries_schema.tsfile" |
| device_ids = [ |
| "root.db1.t1", |
| "root.db2.t1", |
| "root.db3.t2.t3", |
| "root.db3.t3", |
| "device", |
| "device.ln", |
| "device2.ln1.tmp", |
| "device3.ln2.tmp.v1.v2", |
| "device3.ln2.tmp.v1.v3", |
| ] |
| measurement_ids1 = ["temperature", "hudi", "level"] |
| measurement_ids2 = ["level", "vol"] |
| rows_per_device = 2 |
| |
| try: |
| writer = TsFileWriter(file_path) |
| for idx, device_id in enumerate(device_ids): |
| measurements = measurement_ids1 if idx % 2 == 0 else measurement_ids2 |
| for measurement in measurements: |
| writer.register_timeseries( |
| device_id, TimeseriesSchema(measurement, TSDataType.INT32) |
| ) |
| for ts in range(rows_per_device): |
| fields = [] |
| for measurement in measurements: |
| fields.append( |
| Field( |
| measurement, |
| idx * 100 + ts * 10 + len(fields), |
| TSDataType.INT32, |
| ) |
| ) |
| writer.write_row_record(RowRecord(device_id, ts, fields)) |
| writer.close() |
| |
| reader = TsFileReader(file_path) |
| device_schema_map = reader.get_all_timeseries_schemas() |
| expected_devices = {device_id.lower() for device_id in device_ids} |
| assert set(device_schema_map.keys()) == expected_devices |
| print(device_schema_map) |
| |
| for idx, device_id in enumerate(device_ids): |
| measurements = measurement_ids1 if idx % 2 == 0 else measurement_ids2 |
| normalized_device = device_id.lower() |
| assert normalized_device in device_schema_map |
| device_schema = device_schema_map[normalized_device] |
| assert device_schema.get_device_name() == normalized_device |
| timeseries_list = device_schema.get_timeseries_list() |
| assert len(timeseries_list) == len(measurements) |
| actual_measurements = { |
| ts_schema.get_timeseries_name() for ts_schema in timeseries_list |
| } |
| assert actual_measurements == {m.lower() for m in measurements} |
| for ts_schema in timeseries_list: |
| assert ts_schema.get_data_type() == TSDataType.INT32 |
| reader.close() |
| finally: |
| if os.path.exists(file_path): |
| os.remove(file_path) |
| |
| |
| def test_tablet_write_and_read(): |
| try: |
| if os.path.exists("tablet_write_and_read.tsfile"): |
| os.remove("tablet_write_and_read.tsfile") |
| writer = TsFileWriter("tablet_write_and_read.tsfile") |
| measurement_num = 30 |
| for i in range(measurement_num): |
| writer.register_timeseries("root.device1", TimeseriesSchema('level' + str(i), TSDataType.INT64)) |
| |
| max_row_num = 10000 |
| tablet_row_num = 1000 |
| tablet_num = 0 |
| for i in range(max_row_num // tablet_row_num): |
| tablet = Tablet([f'level{j}' for j in range(measurement_num)], |
| [TSDataType.INT64 for _ in range(measurement_num)], tablet_row_num) |
| tablet.set_table_name("root.device1") |
| for row in range(tablet_row_num): |
| tablet.add_timestamp(row, row + tablet_num * tablet_row_num) |
| for col in range(measurement_num): |
| tablet.add_value_by_index(col, row, row + tablet_num * tablet_row_num) |
| writer.write_tablet(tablet) |
| tablet_num += 1 |
| |
| writer.close() |
| |
| reader = TsFileReader("tablet_write_and_read.tsfile") |
| result = reader.query_timeseries("root.device1", ["level0"], 0, 1000000) |
| row_num = 0 |
| print(result.get_result_column_info()) |
| while result.next(): |
| assert result.is_null_by_index(1) == False |
| assert result.get_value_by_index(1) == row_num |
| assert result.get_value_by_name("level0") == row_num |
| assert result.get_value_by_index(2) == row_num |
| row_num = row_num + 1 |
| |
| assert row_num == max_row_num |
| reader.close() |
| with pytest.raises(Exception): |
| result.next() |
| |
| finally: |
| if os.path.exists("tablet_write_and_read.tsfile"): |
| os.remove("tablet_write_and_read.tsfile") |
| |
| |
| def test_table_writer_and_reader(): |
| table = TableSchema("test_table", |
| [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)]) |
| try: |
| if os.path.exists("table_write.tsfile"): |
| os.remove("table_write.tsfile") |
| with TsFileTableWriter("table_write.tsfile", table) as writer: |
| tablet = Tablet(["device", "value"], |
| [TSDataType.STRING, TSDataType.DOUBLE], 100) |
| for i in range(100): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("device", i, "device" + str(i)) |
| tablet.add_value_by_index(1, i, i * 100.0) |
| writer.write_table(tablet) |
| |
| with TsFileReader("table_write.tsfile") as reader: |
| with reader.query_table("test_table", ["device", "value"], |
| 0, 10) as result: |
| cur_line = 0 |
| while result.next(): |
| cur_time = result.get_value_by_name("time") |
| assert result.get_value_by_name("device") == "device" + str(cur_time) |
| assert result.is_null_by_name("device") == False |
| assert result.is_null_by_name("value") == False |
| assert result.is_null_by_index(1) == False |
| assert result.is_null_by_index(2) == False |
| assert result.is_null_by_index(3) == False |
| assert result.get_value_by_name("value") == cur_time * 100.0 |
| cur_line = cur_line + 1 |
| assert cur_line == 11 |
| with reader.query_table("test_table", ["device", "value"], |
| 0, 100) as result: |
| line_num = 0 |
| print("dataframe") |
| while result.next(): |
| data_frame = result.read_data_frame(max_row_num=30) |
| if 100 - line_num >= 30: |
| assert data_frame.shape == (30, 3) |
| else: |
| assert data_frame.shape == (100 - line_num, 3) |
| line_num += len(data_frame) |
| |
| schemas = reader.get_all_table_schemas() |
| assert len(schemas) == 1 |
| assert schemas["test_table"] is not None |
| tableSchema = schemas["test_table"] |
| assert tableSchema.get_table_name() == "test_table" |
| print(tableSchema) |
| assert tableSchema.__repr__() == ("TableSchema(test_table, [ColumnSchema(device," |
| " STRING, TAG), ColumnSchema(value, DOUBLE, FIELD)])") |
| finally: |
| if os.path.exists("table_write.tsfile"): |
| os.remove("table_write.tsfile") |
| |
| |
| def test_query_result_detach_from_reader(): |
| try: |
| ## Prepare data |
| writer = TsFileWriter("query_result_detach_from_reader.tsfile") |
| timeseries = TimeseriesSchema("level1", TSDataType.INT64) |
| writer.register_timeseries("root.device1", timeseries) |
| max_row_num = 1000 |
| for i in range(max_row_num): |
| row = RowRecord("root.device1", i, |
| [Field("level1", i, TSDataType.INT64)]) |
| writer.write_row_record(row) |
| |
| writer.close() |
| |
| reader = TsFileReader("query_result_detach_from_reader.tsfile") |
| result1 = reader.query_timeseries("root.device1", ["level1"], 0, 100) |
| assert 1 == len(reader.get_active_query_result()) |
| result2 = reader.query_timeseries("root.device1", ["level1"], 20, 100) |
| assert 2 == len(reader.get_active_query_result()) |
| result1.close() |
| assert 1 == len(reader.get_active_query_result()) |
| reader.close() |
| with pytest.raises(Exception): |
| result1.next() |
| with pytest.raises(Exception): |
| result2.next() |
| finally: |
| if os.path.exists("query_result_detach_from_reader.tsfile"): |
| os.remove("query_result_detach_from_reader.tsfile") |
| |
| |
| def test_lower_case_name(): |
| if os.path.exists("lower_case_name.tsfile"): |
| os.remove("lower_case_name.tsfile") |
| table = TableSchema("tEst_Table", |
| [ColumnSchema("Device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("vAlue", TSDataType.DOUBLE, ColumnCategory.FIELD)]) |
| with TsFileTableWriter("lower_case_name.tsfile", table) as writer: |
| tablet = Tablet(["device", "VALUE"], [TSDataType.STRING, TSDataType.DOUBLE]) |
| for i in range(100): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("device", i, "device" + str(i)) |
| tablet.add_value_by_name("valuE", i, i * 1.1) |
| |
| writer.write_table(tablet) |
| |
| with TsFileReader("lower_case_name.tsfile") as reader: |
| result = reader.query_table("test_Table", ["DEvice", "value"], 0, 100) |
| while result.next(): |
| print(result.get_value_by_name("DEVICE")) |
| data_frame = result.read_data_frame(max_row_num=130) |
| assert data_frame.shape == (100, 3) |
| assert data_frame["value"].sum() == 5445.0 |
| |
| |
| def test_tsfile_config(): |
| from tsfile import get_tsfile_config, set_tsfile_config |
| |
| config = get_tsfile_config() |
| |
| table = TableSchema("tEst_Table", |
| [ColumnSchema("Device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("vAlue", TSDataType.DOUBLE, ColumnCategory.FIELD)]) |
| if os.path.exists("test1.tsfile"): |
| os.remove("test1.tsfile") |
| with TsFileTableWriter("test1.tsfile", table) as writer: |
| tablet = Tablet(["device", "VALUE"], [TSDataType.STRING, TSDataType.DOUBLE]) |
| for i in range(100): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("device", i, "device" + str(i)) |
| tablet.add_value_by_name("valuE", i, i * 1.1) |
| |
| writer.write_table(tablet) |
| |
| config_normal = get_tsfile_config() |
| print(config_normal) |
| assert config_normal["chunk_group_size_threshold_"] == 128 * 1024 * 1024 |
| |
| os.remove("test1.tsfile") |
| with TsFileTableWriter("test1.tsfile", table, 100 * 100) as writer: |
| tablet = Tablet(["device", "VALUE"], [TSDataType.STRING, TSDataType.DOUBLE]) |
| for i in range(100): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("device", i, "device" + str(i)) |
| tablet.add_value_by_name("valuE", i, i * 1.1) |
| |
| writer.write_table(tablet) |
| config_modified = get_tsfile_config() |
| assert config_normal != config_modified |
| assert config_modified["chunk_group_size_threshold_"] == 100 * 100 |
| set_tsfile_config({'chunk_group_size_threshold_': 100 * 20}) |
| assert get_tsfile_config()["chunk_group_size_threshold_"] == 100 * 20 |
| with pytest.raises(TypeError): |
| set_tsfile_config({"time_compress_type_": TSDataType.DOUBLE}) |
| with pytest.raises(TypeError): |
| set_tsfile_config({'chunk_group_size_threshold_': -1 * 100 * 20}) |
| |
| set_tsfile_config({'float_encoding_type_': TSEncoding.PLAIN}) |
| assert get_tsfile_config()["float_encoding_type_"] == TSEncoding.PLAIN |
| |
| with pytest.raises(TypeError): |
| set_tsfile_config({"float_encoding_type_": -1 * 100 * 20}) |
| with pytest.raises(NotSupportedError): |
| set_tsfile_config({"float_encoding_type_": TSEncoding.BITMAP}) |
| with pytest.raises(NotSupportedError): |
| set_tsfile_config({"time_compress_type_": Compressor.PAA}) |
| |
| |
| def test_tsfile_to_df(): |
| table = TableSchema("test_table", |
| [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ColumnSchema("value2", TSDataType.INT64, ColumnCategory.FIELD)]) |
| try: |
| with TsFileTableWriter("table_write_to_df.tsfile", table) as writer: |
| tablet = Tablet(["device", "value", "value2"], |
| [TSDataType.STRING, TSDataType.DOUBLE, TSDataType.INT64], 4097) |
| for i in range(4097): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("device", i, "device" + str(i)) |
| tablet.add_value_by_index(1, i, i * 100.0) |
| tablet.add_value_by_index(2, i, i * 100) |
| writer.write_table(tablet) |
| df1 = to_dataframe("table_write_to_df.tsfile") |
| assert df1.shape == (4097, 4) |
| assert df1["value2"].sum() == 100 * (1 + 4096) / 2 * 4096 |
| assert is_integer_dtype(df1["time"]) |
| assert df1["value"].dtype == np.float64 |
| assert is_integer_dtype(df1["value2"]) |
| df2 = to_dataframe("table_write_to_df.tsfile", column_names=["device", "value2"]) |
| assert df2.shape == (4097, 3) |
| assert df1["value2"].equals(df2["value2"]) |
| df3 = to_dataframe("table_write_to_df.tsfile", column_names=["device", "value"], max_row_num=8000) |
| assert df3.shape == (4097, 3) |
| with pytest.raises(TableNotExistError): |
| to_dataframe("table_write_to_df.tsfile", "test_tb") |
| with pytest.raises(ColumnNotExistError): |
| to_dataframe("table_write_to_df.tsfile", "test_table", ["device1"]) |
| finally: |
| os.remove("table_write_to_df.tsfile") |
| |
| |
| def test_tree_all_datatype_query_to_dataframe_variants(): |
| tsfile_path = "record_write_and_read.tsfile" |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| writer = TsFileWriter(tsfile_path) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL1", TSDataType.INT64) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL2", TSDataType.DOUBLE) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL3", TSDataType.INT32) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL4", TSDataType.STRING) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL5", TSDataType.TEXT) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL6", TSDataType.BLOB) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL7", TSDataType.DATE) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL8", TSDataType.TIMESTAMP) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL9", TSDataType.BOOLEAN) |
| ) |
| writer.register_timeseries( |
| "root.Device1", TimeseriesSchema("LeveL10", TSDataType.FLOAT) |
| ) |
| |
| max_row_num = 100 |
| |
| for i in range(max_row_num): |
| row = RowRecord( |
| "root.Device1", |
| i - int(max_row_num / 2), |
| [ |
| Field("LeveL1", i * 1, TSDataType.INT64), |
| Field("LeveL2", i * 2.2, TSDataType.DOUBLE), |
| Field("LeveL3", i * 3, TSDataType.INT32), |
| Field("LeveL4", f"string_value_{i}", TSDataType.STRING), |
| Field("LeveL5", f"text_value_{i}", TSDataType.TEXT), |
| Field("LeveL6", f"blob_data_{i}".encode('utf-8'), TSDataType.BLOB), |
| Field("LeveL7", date(2025, 1, i % 20 + 1), TSDataType.DATE), |
| Field("LeveL8", i * 8, TSDataType.TIMESTAMP), |
| Field("LeveL9", i % 2 == 0, TSDataType.BOOLEAN), |
| Field("LeveL10", i * 10.1, TSDataType.FLOAT), |
| ], |
| ) |
| writer.write_row_record(row) |
| |
| writer.close() |
| |
| df1_1 = to_dataframe(tsfile_path) |
| assert df1_1.shape[0] == max_row_num |
| for i in range(max_row_num): |
| assert df1_1.iloc[i, 0] == i - int(max_row_num / 2) |
| assert df1_1.iloc[i, 1] == "root" |
| assert df1_1.iloc[i, 2] == "Device1" |
| |
| df2_1 = to_dataframe(tsfile_path, column_names=["LeveL1"]) |
| for i in range(max_row_num): |
| assert df2_1.iloc[i, 3] == np.int64(i * 1) |
| df2_2 = to_dataframe(tsfile_path, column_names=["LeveL2"]) |
| for i in range(max_row_num): |
| assert df2_2.iloc[i, 3] == np.float64(i * 2.2) |
| df2_3 = to_dataframe(tsfile_path, column_names=["LeveL3"]) |
| for i in range(max_row_num): |
| assert df2_3.iloc[i, 3] == np.int32(i * 3) |
| df2_4 = to_dataframe(tsfile_path, column_names=["LeveL4"]) |
| for i in range(max_row_num): |
| assert df2_4.iloc[i, 3] == f"string_value_{i}" |
| df2_5 = to_dataframe(tsfile_path, column_names=["LeveL5"]) |
| for i in range(max_row_num): |
| assert df2_5.iloc[i, 3] == f"text_value_{i}" |
| df2_6 = to_dataframe(tsfile_path, column_names=["LeveL6"]) |
| for i in range(max_row_num): |
| assert df2_6.iloc[i, 3] == f"blob_data_{i}".encode('utf-8') |
| df2_7 = to_dataframe(tsfile_path, column_names=["LeveL7"]) |
| for i in range(max_row_num): |
| assert df2_7.iloc[i, 3] == date(2025, 1, i % 20 + 1) |
| df2_8 = to_dataframe(tsfile_path, column_names=["LeveL8"]) |
| for i in range(max_row_num): |
| assert df2_8.iloc[i, 3] == np.int64(i * 8) |
| df2_9 = to_dataframe(tsfile_path, column_names=["LeveL9"]) |
| for i in range(max_row_num): |
| assert df2_9.iloc[i, 3] == (i % 2 == 0) |
| df2_10 = to_dataframe(tsfile_path, column_names=["LeveL10"]) |
| for i in range(max_row_num): |
| assert df2_10.iloc[i, 3] == np.float32(i * 10.1) |
| df2_11 = to_dataframe(tsfile_path, column_names=["LeveL9"]) |
| for i in range(max_row_num): |
| assert df2_11.iloc[i, 3] == (i % 2 == 0) |
| df2_12 = to_dataframe( |
| tsfile_path, |
| column_names=[ |
| "LeveL1", |
| "LeveL2", |
| "LeveL3", |
| "LeveL4", |
| "LeveL5", |
| "LeveL6", |
| "LeveL7", |
| "LeveL8", |
| "LeveL9", |
| "LeveL10", |
| ], |
| ) |
| for i in range(max_row_num): |
| assert df2_12.iloc[i, 3] == np.int64(i * 1) |
| assert df2_12.iloc[i, 4] == np.float64(i * 2.2) |
| assert df2_12.iloc[i, 5] == np.int32(i * 3) |
| assert df2_12.iloc[i, 6] == f"string_value_{i}" |
| assert df2_12.iloc[i, 7] == f"text_value_{i}" |
| assert df2_12.iloc[i, 8] == f"blob_data_{i}".encode('utf-8') |
| assert df2_12.iloc[i, 9] == date(2025, 1, i % 20 + 1) |
| assert df2_12.iloc[i, 10] == np.int64(i * 8) |
| assert df2_12.iloc[i, 11] == (i % 2 == 0) |
| assert df2_12.iloc[i, 12] == np.float32(i * 10.1) |
| |
| df3_1 = to_dataframe(tsfile_path, start_time=10) |
| assert df3_1.shape[0] == 40 |
| df3_2 = to_dataframe(tsfile_path, start_time=-10) |
| assert df3_2.shape[0] == 60 |
| df3_3 = to_dataframe(tsfile_path, end_time=5) |
| assert df3_3.shape[0] == 56 |
| df3_4 = to_dataframe(tsfile_path, end_time=-5) |
| assert df3_4.shape[0] == 46 |
| df3_5 = to_dataframe(tsfile_path, start_time=5, end_time=5) |
| assert df3_5.shape[0] == 1 |
| df3_6 = to_dataframe(tsfile_path, start_time=-5, end_time=-5) |
| assert df3_6.shape[0] == 1 |
| df3_7 = to_dataframe(tsfile_path, start_time=10, end_time=-10) |
| assert df3_7.shape[0] == 0 |
| df3_8 = to_dataframe(tsfile_path, start_time=-10, end_time=10) |
| assert df3_8.shape[0] == 21 |
| df3_8 = to_dataframe(tsfile_path, start_time=-50, end_time=50) |
| assert df3_8.shape[0] == max_row_num |
| |
| df4_1 = to_dataframe(tsfile_path, max_row_num=1) |
| assert df4_1.shape[0] == 1 |
| df4_2 = to_dataframe(tsfile_path, max_row_num=10) |
| assert df4_2.shape[0] == 10 |
| df4_3 = to_dataframe(tsfile_path, max_row_num=100) |
| assert df4_3.shape[0] == 100 |
| df4_4 = to_dataframe(tsfile_path, max_row_num=1000) |
| assert df4_4.shape[0] == 100 |
| df4_5 = to_dataframe(tsfile_path, max_row_num=0) |
| assert df4_5.shape[0] == 0 |
| df4_6 = to_dataframe(tsfile_path, max_row_num=-10) |
| assert df4_6.shape[0] == 0 |
| |
| for df5_1 in to_dataframe(tsfile_path, max_row_num=10, as_iterator=True): |
| assert df5_1.shape[0] == 10 |
| for df5_2 in to_dataframe(tsfile_path, max_row_num=-10, as_iterator=True): |
| assert df5_2.shape[0] == 1 |
| for df5_3 in to_dataframe(tsfile_path, max_row_num=1000, as_iterator=True): |
| assert df5_3.shape[0] == 100 |
| for df5_4 in to_dataframe(tsfile_path, max_row_num=3, as_iterator=True): |
| if df5_4.iloc[0, 0] <= 48: |
| assert df5_4.shape[0] == 3 |
| else: |
| assert df5_4.shape[0] == 1 |
| |
| row_num = 0 |
| for df6_1 in to_dataframe( |
| tsfile_path, |
| column_names=["LeveL1", "LeveL2"], |
| start_time=-50, |
| end_time=10, |
| max_row_num=1, |
| as_iterator=True, |
| ): |
| assert df6_1.shape[0] == 1 |
| assert df6_1.iloc[0, 0] == -50 + row_num |
| assert df6_1.iloc[0, 3] == row_num |
| assert df6_1.iloc[0, 4] == row_num * 2.2 |
| row_num += 1 |
| |
| df7_1 = to_dataframe(tsfile_path, table_name="test") |
| assert df7_1.shape[0] == max_row_num |
| assert df7_1.iloc[0, 0] == -int(max_row_num / 2) |
| |
| try: |
| to_dataframe(tsfile_path, column_names=["non_existent_column"]) |
| except ColumnNotExistError: |
| pass |
| |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| def test_table_all_datatype_query_to_dataframe_variants(): |
| tsfile_path = "test_table.tsfile" |
| table = TableSchema( |
| "test_table", |
| [ |
| ColumnSchema("Device1", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("Device2", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("Value1", TSDataType.BOOLEAN, ColumnCategory.FIELD), |
| ColumnSchema("Value2", TSDataType.INT32, ColumnCategory.FIELD), |
| ColumnSchema("Value3", TSDataType.INT64, ColumnCategory.FIELD), |
| ColumnSchema("Value4", TSDataType.FLOAT, ColumnCategory.FIELD), |
| ColumnSchema("Value5", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ColumnSchema("Value6", TSDataType.TEXT, ColumnCategory.FIELD), |
| ColumnSchema("Value7", TSDataType.STRING, ColumnCategory.FIELD), |
| ColumnSchema("Value8", TSDataType.BLOB, ColumnCategory.FIELD), |
| ColumnSchema("Value9", TSDataType.TIMESTAMP, ColumnCategory.FIELD), |
| ColumnSchema("Value10", TSDataType.DATE, ColumnCategory.FIELD), |
| ], |
| ) |
| dateSet = set() |
| try: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| max_row_num = 100 |
| with TsFileTableWriter(tsfile_path, table) as writer: |
| tablet = Tablet( |
| [ |
| "Device1", |
| "Device2", |
| "Value1", |
| "Value2", |
| "Value3", |
| "Value4", |
| "Value5", |
| "Value6", |
| "Value7", |
| "Value8", |
| "Value9", |
| "Value10", |
| ], |
| [ |
| TSDataType.STRING, |
| TSDataType.STRING, |
| TSDataType.BOOLEAN, |
| TSDataType.INT32, |
| TSDataType.INT64, |
| TSDataType.FLOAT, |
| TSDataType.DOUBLE, |
| TSDataType.TEXT, |
| TSDataType.STRING, |
| TSDataType.BLOB, |
| TSDataType.TIMESTAMP, |
| TSDataType.DATE, |
| ], |
| max_row_num, |
| ) |
| for i in range(max_row_num): |
| tablet.add_timestamp(i, i) |
| tablet.add_value_by_name("Device1", i, "d1_" + str(i)) |
| tablet.add_value_by_name("Device2", i, "d2_" + str(i)) |
| tablet.add_value_by_name("Value1", i, i % 2 == 0) |
| tablet.add_value_by_name("Value2", i, i * 3) |
| tablet.add_value_by_name("Value3", i, i * 4) |
| tablet.add_value_by_name("Value4", i, i * 5.5) |
| tablet.add_value_by_name("Value5", i, i * 6.6) |
| tablet.add_value_by_name("Value6", i, f"string_value_{i}") |
| tablet.add_value_by_name("Value7", i, f"text_value_{i}") |
| tablet.add_value_by_name("Value8", i, f"blob_data_{i}".encode('utf-8')) |
| tablet.add_value_by_name("Value9", i, i * 9) |
| tablet.add_value_by_name("Value10", i, date(2025, 1, i % 20 + 1)) |
| dateSet.add(date(2025, 1, i % 20 + 1)) |
| writer.write_table(tablet) |
| |
| df1_1 = to_dataframe(tsfile_path) |
| assert df1_1.shape[0] == max_row_num |
| for i in range(max_row_num): |
| assert df1_1.iloc[i, 1] == "d1_" + str(df1_1.iloc[i, 0]) |
| assert df1_1.iloc[i, 2] == "d2_" + str(df1_1.iloc[i, 0]) |
| |
| df2_1 = to_dataframe(tsfile_path, column_names=["Value1"]) |
| for i in range(max_row_num): |
| assert df2_1.iloc[i, 1] == np.bool_(df2_1.iloc[i, 0] % 2 == 0) |
| df2_2 = to_dataframe(tsfile_path, column_names=["Value2"]) |
| for i in range(max_row_num): |
| assert df2_2.iloc[i, 1] == np.int32(df2_2.iloc[i, 0] * 3) |
| df2_3 = to_dataframe(tsfile_path, column_names=["Value3"]) |
| for i in range(max_row_num): |
| assert df2_3.iloc[i, 1] == np.int64(df2_3.iloc[i, 0] * 4) |
| df2_4 = to_dataframe(tsfile_path, column_names=["Value4"]) |
| for i in range(max_row_num): |
| assert df2_4.iloc[i, 1] == np.float32(df2_4.iloc[i, 0] * 5.5) |
| df2_5 = to_dataframe(tsfile_path, column_names=["Value5"]) |
| for i in range(max_row_num): |
| assert df2_5.iloc[i, 1] == np.float64(df2_5.iloc[i, 0] * 6.6) |
| df2_6 = to_dataframe(tsfile_path, column_names=["Value6"]) |
| for i in range(max_row_num): |
| assert df2_6.iloc[i, 1] == f"string_value_{df2_6.iloc[i, 0]}" |
| df2_7 = to_dataframe(tsfile_path, column_names=["Value7"]) |
| for i in range(max_row_num): |
| assert df2_7.iloc[i, 1] == f"text_value_{df2_7.iloc[i, 0]}" |
| df2_8 = to_dataframe(tsfile_path, column_names=["Value8"]) |
| for i in range(max_row_num): |
| assert df2_8.iloc[i, 1] == f"blob_data_{df2_8.iloc[i, 0]}".encode('utf-8') |
| df2_9 = to_dataframe(tsfile_path, column_names=["Value9"]) |
| for i in range(max_row_num): |
| assert df2_9.iloc[i, 1] == np.int64(df2_9.iloc[i, 0] * 9) |
| df2_10 = to_dataframe(tsfile_path, column_names=["Value10"]) |
| for i in range(max_row_num): |
| assert df2_10.iloc[i, 1] in dateSet |
| df2_11 = to_dataframe(tsfile_path, column_names=["Device1", "Value1"]) |
| for i in range(max_row_num): |
| assert df2_11.iloc[i, 1] == "d1_" + str(df2_11.iloc[i, 0]) |
| assert df2_11.iloc[i, 2] == np.bool_(df2_11.iloc[i, 0] % 2 == 0) |
| df2_12 = to_dataframe( |
| tsfile_path, |
| column_names=[ |
| "Device1", |
| "Device2", |
| "Value1", |
| "Value2", |
| "Value3", |
| "Value4", |
| "Value5", |
| "Value6", |
| "Value7", |
| "Value8", |
| "Value9", |
| "Value10", |
| ], |
| ) |
| for i in range(max_row_num): |
| assert df2_12.iloc[i, 1] == "d1_" + str(df2_12.iloc[i, 0]) |
| assert df2_12.iloc[i, 2] == "d2_" + str(df2_12.iloc[i, 0]) |
| assert df2_12.iloc[i, 3] == np.bool_(df2_12.iloc[i, 0] % 2 == 0) |
| assert df2_12.iloc[i, 4] == np.int32(df2_12.iloc[i, 0] * 3) |
| assert df2_12.iloc[i, 5] == np.int64(df2_12.iloc[i, 0] * 4) |
| assert df2_12.iloc[i, 6] == np.float32(df2_12.iloc[i, 0] * 5.5) |
| assert df2_12.iloc[i, 7] == np.float64(df2_12.iloc[i, 0] * 6.6) |
| assert df2_12.iloc[i, 8] == f"string_value_{df2_12.iloc[i, 0]}" |
| assert df2_12.iloc[i, 9] == f"text_value_{df2_12.iloc[i, 0]}" |
| assert df2_12.iloc[i, 10] == f"blob_data_{df2_12.iloc[i, 0]}".encode( |
| "utf-8" |
| ) |
| assert df2_12.iloc[i, 11] == np.int64(df2_12.iloc[i, 0] * 9) |
| assert df2_12.iloc[i, 12] == date(2025, 1, df2_12.iloc[i, 0] % 20 + 1) |
| df2_13 = to_dataframe( |
| tsfile_path, column_names=["Device1", "Device2", "Value1"] |
| ) |
| for i in range(max_row_num): |
| assert df2_13.iloc[i, 1] == "d1_" + str(df2_13.iloc[i, 0]) |
| assert df2_13.iloc[i, 2] == "d2_" + str(df2_13.iloc[i, 0]) |
| assert df2_13.iloc[i, 3] == np.bool_(df2_13.iloc[i, 0] % 2 == 0) |
| |
| df3_1 = to_dataframe(tsfile_path, table_name="test_table") |
| assert df3_1.shape[0] == max_row_num |
| assert df3_1.iloc[0, 0] == 0 |
| df3_2 = to_dataframe(tsfile_path, table_name="TEST_TABLE") |
| assert df3_2.shape[0] == max_row_num |
| assert df3_2.iloc[0, 0] == 0 |
| |
| df4_1 = to_dataframe(tsfile_path, start_time=10) |
| assert df4_1.shape[0] == 90 |
| df4_2 = to_dataframe(tsfile_path, start_time=-10) |
| assert df4_2.shape[0] == max_row_num |
| df4_3 = to_dataframe(tsfile_path, end_time=5) |
| assert df4_3.shape[0] == 6 |
| df4_4 = to_dataframe(tsfile_path, end_time=-5) |
| assert df4_4.shape[0] == 0 |
| df4_5 = to_dataframe(tsfile_path, start_time=5, end_time=5) |
| assert df4_5.shape[0] == 1 |
| df4_6 = to_dataframe(tsfile_path, start_time=-5, end_time=-5) |
| assert df4_6.shape[0] == 0 |
| df4_7 = to_dataframe(tsfile_path, start_time=10, end_time=-10) |
| assert df4_7.shape[0] == 0 |
| df4_8 = to_dataframe(tsfile_path, start_time=-10, end_time=10) |
| assert df4_8.shape[0] == 11 |
| df4_8 = to_dataframe(tsfile_path, start_time=-50, end_time=50) |
| assert df4_8.shape[0] == 51 |
| |
| df5_1 = to_dataframe(tsfile_path, max_row_num=1) |
| assert df5_1.shape[0] == 1 |
| df5_2 = to_dataframe(tsfile_path, max_row_num=50) |
| assert df5_2.shape[0] == 50 |
| df5_3 = to_dataframe(tsfile_path, max_row_num=100) |
| assert df5_3.shape[0] == 100 |
| df5_4 = to_dataframe(tsfile_path, max_row_num=1000) |
| assert df5_4.shape[0] == 100 |
| df5_5 = to_dataframe(tsfile_path, max_row_num=0) |
| assert df5_5.shape[0] == 0 |
| df5_6 = to_dataframe(tsfile_path, max_row_num=-10) |
| assert df5_6.shape[0] == 0 |
| |
| for df6_1 in to_dataframe(tsfile_path, max_row_num=20, as_iterator=True): |
| assert df6_1.shape[0] == 20 |
| for df6_2 in to_dataframe(tsfile_path, max_row_num=1000, as_iterator=True): |
| assert df6_2.shape[0] == 100 |
| |
| for df7_1 in to_dataframe( |
| tsfile_path, |
| table_name="test_table", |
| column_names=["Device1", "Value1"], |
| start_time=21, |
| end_time=50, |
| max_row_num=10, |
| as_iterator=True, |
| ): |
| assert df7_1.shape[0] == 10 |
| for i in range(30): |
| assert df2_11.iloc[i, 1] == "d1_" + str(df2_11.iloc[i, 0]) |
| assert df2_11.iloc[i, 2] == np.bool_(df2_11.iloc[i, 0] % 2 == 0) |
| |
| try: |
| to_dataframe(tsfile_path, table_name="non_existent_table") |
| except TableNotExistError as e: |
| assert e.args[0] == "[non_existent_table] Requested table does not exist" |
| |
| try: |
| to_dataframe(tsfile_path, column_names=["non_existent_column"]) |
| except ColumnNotExistError as e: |
| assert e.args[0] == "[non_existent_column] Column does not exist" |
| |
| finally: |
| if os.path.exists(tsfile_path): |
| os.remove(tsfile_path) |
| |
| |
| import os |
| |
| if __name__ == "__main__": |
| os.chdir(os.path.dirname(os.path.abspath(__file__))) |
| pytest.main([ |
| "test_write_and_read.py::test_row_record_write_and_read", |
| "-s", "-v" |
| ]) |
