| # 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. |
| # |
| """ |
| Tests for write_arrow_batch: write PyArrow RecordBatch/Table to tsfile |
| and verify correctness by reading back. |
| """ |
| |
| import os |
| from datetime import date |
| |
| import numpy as np |
| import pandas as pd |
| import pytest |
| |
| pa = pytest.importorskip("pyarrow", reason="pyarrow is not installed") |
| |
| from tsfile import ColumnCategory, ColumnSchema, TableSchema, TSDataType, TsFileReader |
| from tsfile import TsFileTableWriter |
| |
| |
| # --------------------------------------------------------------------------- |
| # Helpers |
| # --------------------------------------------------------------------------- |
| |
| def _make_schema(table_name, extra_cols): |
| """Build a TableSchema with a string TAG 'device' plus the given field cols.""" |
| return TableSchema( |
| table_name, |
| [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG)] + extra_cols, |
| ) |
| |
| |
| def _read_all_arrow(file_path, table_name, columns, start=0, end=10**18, batch_size=4096): |
| """Read all rows from file via read_arrow_batch and return as a pa.Table.""" |
| reader = TsFileReader(file_path) |
| rs = reader.query_table_batch( |
| table_name=table_name, |
| column_names=columns, |
| start_time=start, |
| end_time=end, |
| batch_size=batch_size, |
| ) |
| batches = [] |
| while True: |
| batch = rs.read_arrow_batch() |
| if batch is None: |
| break |
| batches.append(batch) |
| rs.close() |
| reader.close() |
| if not batches: |
| return pa.table({}) |
| return pa.concat_tables(batches) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Basic write + read-back |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_basic(): |
| """Write 1 000 rows via write_arrow_batch and verify count + values.""" |
| path = "test_write_arrow_basic.tsfile" |
| table_name = "t" |
| n = 1000 |
| |
| schema = _make_schema(table_name, [ |
| ColumnSchema("value1", TSDataType.INT64, ColumnCategory.FIELD), |
| ColumnSchema("value2", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ]) |
| |
| batch = pa.record_batch({ |
| "time": pa.array(np.arange(n, dtype="int64"), type=pa.timestamp("ns")), |
| "device": pa.array([f"d{i}" for i in range(n)], type=pa.string()), |
| "value1": pa.array(np.arange(n, dtype="int64"), type=pa.int64()), |
| "value2": pa.array(np.arange(n, dtype="float64") * 1.5, type=pa.float64()), |
| }) |
| |
| try: |
| if os.path.exists(path): |
| os.remove(path) |
| with TsFileTableWriter(path, schema) as w: |
| w.write_arrow_batch(batch) |
| |
| result = _read_all_arrow(path, table_name, ["device", "value1", "value2"]) |
| assert len(result) == n |
| |
| df = result.to_pandas().sort_values("time").reset_index(drop=True) |
| assert list(df["value1"]) == list(range(n)) |
| assert all(abs(df["value2"].iloc[i] - i * 1.5) < 1e-9 for i in range(n)) |
| finally: |
| if os.path.exists(path): |
| os.remove(path) |
| |
| |
| # --------------------------------------------------------------------------- |
| # pa.Table input |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_from_table(): |
| """write_arrow_batch should accept pa.Table (multi-chunk) as well.""" |
| path = "test_write_arrow_from_table.tsfile" |
| table_name = "t" |
| n = 500 |
| |
| schema = _make_schema(table_name, [ |
| ColumnSchema("v", TSDataType.INT32, ColumnCategory.FIELD), |
| ]) |
| |
| tbl = pa.table({ |
| "time": pa.array(np.arange(n, dtype="int64"), type=pa.timestamp("ns")), |
| "device": pa.array(["dev"] * n, type=pa.string()), |
| "v": pa.array(np.arange(n, dtype="int32"), type=pa.int32()), |
| }) |
| |
| try: |
| if os.path.exists(path): |
| os.remove(path) |
| with TsFileTableWriter(path, schema) as w: |
| w.write_arrow_batch(tbl) |
| |
| result = _read_all_arrow(path, table_name, ["device", "v"]) |
| assert len(result) == n |
| df = result.to_pandas().sort_values("time").reset_index(drop=True) |
| assert list(df["v"]) == list(range(n)) |
| finally: |
| if os.path.exists(path): |
| os.remove(path) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Multiple batches |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_multiple_batches(): |
| """Write several batches sequentially and verify the total row count.""" |
| path = "test_write_arrow_multi.tsfile" |
| table_name = "t" |
| rows_per_batch = 300 |
| num_batches = 4 |
| total = rows_per_batch * num_batches |
| |
| schema = _make_schema(table_name, [ |
| ColumnSchema("v", TSDataType.INT64, ColumnCategory.FIELD), |
| ]) |
| |
| try: |
| if os.path.exists(path): |
| os.remove(path) |
| with TsFileTableWriter(path, schema) as w: |
| for b in range(num_batches): |
| start_ts = b * rows_per_batch |
| batch = pa.record_batch({ |
| "time": pa.array( |
| np.arange(start_ts, start_ts + rows_per_batch, dtype="int64"), |
| type=pa.timestamp("ns")), |
| "device": pa.array(["dev"] * rows_per_batch, type=pa.string()), |
| "v": pa.array( |
| np.arange(start_ts, start_ts + rows_per_batch, dtype="int64"), |
| type=pa.int64()), |
| }) |
| w.write_arrow_batch(batch) |
| |
| result = _read_all_arrow(path, table_name, ["device", "v"]) |
| assert len(result) == total |
| finally: |
| if os.path.exists(path): |
| os.remove(path) |
| |
| |
| # --------------------------------------------------------------------------- |
| # All supported data types |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_all_datatypes(): |
| """Write every supported data type and verify values read back correctly.""" |
| path = "test_write_arrow_all_types.tsfile" |
| table_name = "t" |
| n = 200 |
| |
| schema = TableSchema(table_name, [ |
| ColumnSchema("tag", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("bool_col", TSDataType.BOOLEAN, ColumnCategory.FIELD), |
| ColumnSchema("int32_col", TSDataType.INT32, ColumnCategory.FIELD), |
| ColumnSchema("int64_col", TSDataType.INT64, ColumnCategory.FIELD), |
| ColumnSchema("float_col", TSDataType.FLOAT, ColumnCategory.FIELD), |
| ColumnSchema("double_col", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ColumnSchema("str_col", TSDataType.STRING, ColumnCategory.FIELD), |
| ColumnSchema("date_col", TSDataType.DATE, ColumnCategory.FIELD), |
| ]) |
| |
| dates_days = [ |
| (date(2025, 1, (i % 28) + 1) - date(1970, 1, 1)).days for i in range(n) |
| ] |
| |
| batch = pa.record_batch({ |
| "time": pa.array(np.arange(n, dtype="int64"), type=pa.timestamp("ns")), |
| "tag": pa.array([f"dev{i}" for i in range(n)], type=pa.string()), |
| "bool_col": pa.array([i % 2 == 0 for i in range(n)], type=pa.bool_()), |
| "int32_col": pa.array(np.arange(n, dtype="int32"), type=pa.int32()), |
| "int64_col": pa.array(np.arange(n, dtype="int64") * 10, type=pa.int64()), |
| "float_col": pa.array(np.arange(n, dtype="float32") * 0.5, type=pa.float32()), |
| "double_col": pa.array(np.arange(n, dtype="float64") * 1.1, type=pa.float64()), |
| "str_col": pa.array([f"s{i}" for i in range(n)], type=pa.string()), |
| "date_col": pa.array(dates_days, type=pa.date32()), |
| }) |
| |
| try: |
| if os.path.exists(path): |
| os.remove(path) |
| with TsFileTableWriter(path, schema) as w: |
| w.write_arrow_batch(batch) |
| |
| result = _read_all_arrow( |
| path, table_name, |
| ["tag", "bool_col", "int32_col", "int64_col", |
| "float_col", "double_col", "str_col", "date_col"], |
| ) |
| assert len(result) == n |
| df = result.to_pandas().sort_values("time").reset_index(drop=True) |
| |
| for col in ["tag", "bool_col", "int32_col", "int64_col", |
| "float_col", "double_col", "str_col", "date_col"]: |
| assert col in df.columns, f"Column '{col}' missing from result" |
| |
| assert list(df["int32_col"]) == list(range(n)) |
| assert list(df["int64_col"]) == [i * 10 for i in range(n)] |
| for i in range(n): |
| assert df["bool_col"].iloc[i] == (i % 2 == 0) |
| assert abs(df["double_col"].iloc[i] - i * 1.1) < 1e-9 |
| assert df["str_col"].iloc[i] == f"s{i}" |
| finally: |
| if os.path.exists(path): |
| os.remove(path) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Parity with write_dataframe |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_parity_with_dataframe(): |
| """Data written via write_arrow_batch must match data written via write_dataframe.""" |
| arrow_path = "test_write_arrow_parity_arrow.tsfile" |
| df_path = "test_write_arrow_parity_df.tsfile" |
| table_name = "t" |
| n = 500 |
| |
| schema_arrow = TableSchema(table_name, [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("v_i32", TSDataType.INT32, ColumnCategory.FIELD), |
| ColumnSchema("v_f64", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ColumnSchema("v_bool", TSDataType.BOOLEAN, ColumnCategory.FIELD), |
| ColumnSchema("v_str", TSDataType.STRING, ColumnCategory.FIELD), |
| ]) |
| schema_df = TableSchema(table_name, [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("v_i32", TSDataType.INT32, ColumnCategory.FIELD), |
| ColumnSchema("v_f64", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ColumnSchema("v_bool", TSDataType.BOOLEAN, ColumnCategory.FIELD), |
| ColumnSchema("v_str", TSDataType.STRING, ColumnCategory.FIELD), |
| ]) |
| |
| timestamps = np.arange(n, dtype="int64") |
| v_i32 = np.arange(n, dtype="int32") |
| v_f64 = np.arange(n, dtype="float64") * 2.5 |
| v_bool = np.array([i % 3 == 0 for i in range(n)]) |
| v_str = [f"row{i}" for i in range(n)] |
| device = ["dev"] * n |
| |
| batch = pa.record_batch({ |
| "time": pa.array(timestamps, type=pa.timestamp("ns")), |
| "device": pa.array(device, type=pa.string()), |
| "v_i32": pa.array(v_i32, type=pa.int32()), |
| "v_f64": pa.array(v_f64, type=pa.float64()), |
| "v_bool": pa.array(v_bool, type=pa.bool_()), |
| "v_str": pa.array(v_str, type=pa.string()), |
| }) |
| |
| dataframe = pd.DataFrame({ |
| "time": pd.Series(timestamps, dtype="int64"), |
| "device": device, |
| "v_i32": pd.Series(v_i32, dtype="int32"), |
| "v_f64": pd.Series(v_f64, dtype="float64"), |
| "v_bool": pd.Series(v_bool, dtype="bool"), |
| "v_str": v_str, |
| }) |
| |
| cols = ["device", "v_i32", "v_f64", "v_bool", "v_str"] |
| |
| try: |
| for p in (arrow_path, df_path): |
| if os.path.exists(p): |
| os.remove(p) |
| |
| with TsFileTableWriter(arrow_path, schema_arrow) as w: |
| w.write_arrow_batch(batch) |
| with TsFileTableWriter(df_path, schema_df) as w: |
| w.write_dataframe(dataframe) |
| |
| result_arrow = _read_all_arrow(arrow_path, table_name, cols).to_pandas() |
| result_df = _read_all_arrow(df_path, table_name, cols).to_pandas() |
| |
| result_arrow = result_arrow.sort_values("time").reset_index(drop=True) |
| result_df = result_df.sort_values("time").reset_index(drop=True) |
| |
| assert len(result_arrow) == len(result_df) == n |
| |
| assert list(result_arrow["v_i32"]) == list(result_df["v_i32"]) |
| assert list(result_arrow["v_str"]) == list(result_df["v_str"]) |
| assert list(result_arrow["v_bool"]) == list(result_df["v_bool"]) |
| for i in range(n): |
| assert abs(result_arrow["v_f64"].iloc[i] - result_df["v_f64"].iloc[i]) < 1e-9 |
| finally: |
| for p in (arrow_path, df_path): |
| if os.path.exists(p): |
| os.remove(p) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Large batch |
| # --------------------------------------------------------------------------- |
| |
| def test_write_arrow_large_batch(): |
| """Write a single large batch (100 k rows) and verify row count.""" |
| path = "test_write_arrow_large.tsfile" |
| table_name = "t" |
| n = 100_000 |
| |
| schema = _make_schema(table_name, [ |
| ColumnSchema("v", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ]) |
| |
| batch = pa.record_batch({ |
| "time": pa.array(np.arange(n, dtype="int64"), type=pa.timestamp("ns")), |
| "device": pa.array(["d"] * n, type=pa.string()), |
| "v": pa.array(np.random.rand(n), type=pa.float64()), |
| }) |
| |
| try: |
| if os.path.exists(path): |
| os.remove(path) |
| with TsFileTableWriter(path, schema) as w: |
| w.write_arrow_batch(batch) |
| |
| result = _read_all_arrow(path, table_name, ["device", "v"], batch_size=8192) |
| assert len(result) == n |
| finally: |
| if os.path.exists(path): |
| os.remove(path) |
| |
| |
| if __name__ == "__main__": |
| os.chdir(os.path.dirname(os.path.abspath(__file__))) |
| pytest.main([__file__, "-v", "-s"]) |
