IoTDB AI Toolkit

高性能时间序列机器学习工具包，为 IoTDB 和时间序列分析提供强大的 AI/ML 支持。

🚀 特性

TsFileDataset - 高性能 PyTorch 数据集

直接从 TsFile 文件读取数据用于深度学习模型训练，无需运行 IoTDB 服务器。

核心优势：

• 🔥 高性能：O(log n) 索引查找，使用官方 tsfile 包优化读取
• 💾 内存高效：支持懒加载和预加载两种模式
• 🎯 简洁易用：最少的参数配置，开箱即用
• 🔧 独立部署：不依赖 IoTDB 服务器运行
• 📦 多文件支持：自动处理多个 TsFile 文件
• ✨ 官方支持：使用 Apache TsFile 官方 Python 包，兼容性最佳

TsFileDataFrame - 懒加载时序数据分析视图

对多个 TsFile 文件提供统一的 DataFrame 式访问接口，支持数组式索引和时间对齐查询。

核心优势：

• 📂 多文件统一视图：自动合并多个 TsFile 中的同名序列，透明处理跨文件数据
• ⚡ 懒加载：仅在查询时读取数据，大幅减少内存占用
• 🔗 多文件并行加载：多文件场景下自动并行读取元数据，加速初始化
• 🎯 灵活索引：支持按序号、名称、切片、列表等多种方式访问序列
• 📊 时间对齐查询：通过 .loc 实现多序列时间对齐，自动填充 NaN
• 🖨️ 友好展示：内置格式化打印，支持截断展示和全量展示

📦 安装

# 从源码安装
git clone https://github.com/your-repo/IoTDB-AI.git
cd IoTDB-AI
pip install -e .

# 或者使用 pip（如果已发布）
pip install iotdb-ai

依赖说明：

• torch>=1.7.0 - PyTorch 深度学习框架
• numpy>=1.19.0 - 数值计算库
• tsfile>=2.0.0 - Apache TsFile 官方 Python 包（自动安装）

🎯 快速开始

基本使用

from iotdb_ai import TsFileDataset
from torch.utils.data import DataLoader

# 创建数据集
dataset = TsFileDataset(
    tsfile_paths="data/sensor_data.tsfile",  # 单个文件或文件列表
    seq_len=96,      # 输入序列长度（lookback window）
    pred_len=24,     # 预测长度（forecast horizon）
    window_stride=1  # 滑动窗口步长（默认1）
)

# 使用 PyTorch DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# 训练模型
for input_seq, target_seq in dataloader:
    # input_seq: (batch_size, seq_len)
    # target_seq: (batch_size, pred_len)
    ...

多文件支持

# 同时处理多个 TsFile 文件
dataset = TsFileDataset(
    tsfile_paths=[
        "data/2024-01.tsfile",
        "data/2024-02.tsfile",
        "data/2024-03.tsfile",
    ],
    seq_len=168,     # 一周数据
    pred_len=24,     # 预测未来一天
)

预加载模式（小数据集）

# 内存充足时，一次性加载所有数据获得最佳性能
dataset = TsFileDataset(
    tsfile_paths="data/small_dataset.tsfile",
    seq_len=96,
    pred_len=24,
    preload=True  # 启用预加载
)

📊 性能优化

推荐配置

根据数据集大小选择合适的配置：

数据集大小	推荐配置	说明
< 1GB	preload=True	全部加载到内存，极速访问
1-10GB	num_workers=4-8	多进程懒加载，平衡性能和内存
> 10GB	window_stride=8+	增大步长减少样本数

DataLoader 优化

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,           # 多进程并行加载
    pin_memory=True,         # 加速 GPU 传输
    persistent_workers=True  # 保持工作进程
)

📚 API 文档

TsFileDataset

TsFileDataset(
    tsfile_paths: Union[str, List[str]],  # TsFile 文件路径
    seq_len: int,                         # 输入序列长度
    pred_len: int,                        # 预测序列长度
    window_stride: int = 1,               # 滑动窗口步长
    preload: bool = False,                # 是否预加载到内存
)

方法：

• __len__(): 返回数据集样本总数
• __getitem__(index): 获取指定索引的样本
• get_series_count(): 返回时间序列数量
• get_info(): 返回数据集详细信息

返回值：

• (input_seq, target_seq) 元组
  • input_seq: torch.Tensor, shape (seq_len,)
  • target_seq: torch.Tensor, shape (pred_len,)

TsFileDataFrame

TsFileDataFrame(
    paths: Union[str, List[str]],  # TsFile 文件路径、路径列表或目录路径
    show_progress: bool = True,    # 是否显示加载进度
)

属性：

• timeseries -> List[dict]: 所有序列的元信息列表，每项包含 ts_name、start_time、end_time、count
• loc -> _LocIndexer: 时间对齐查询索引器

方法：

• __len__(): 返回序列总数
• __getitem__(key): 按索引(int)、名称(str)、切片(slice)或列表(list)访问序列
• list_timeseries(path_prefix=""): 按前缀过滤序列名称
• close(): 关闭所有底层文件句柄
• 支持 with 上下文管理器

Timeseries

由 TsFileDataFrame[...] 返回的单条时间序列对象。

属性：

• name -> str: 序列完整路径
• timestamps -> np.ndarray: 毫秒时间戳数组
• stats -> dict: 包含 start_time、end_time、count

方法：

• __len__(): 返回数据点数量
• __getitem__(key): 行索引访问，支持 ts[20]（单值）和 ts[20:100]（切片，返回 np.ndarray）

AlignedTimeseries

由 .loc[...] 返回的时间对齐查询结果。

属性：

• timestamps -> np.ndarray: 对齐后的毫秒时间戳
• values -> np.ndarray: shape (rows, cols) 的值矩阵，缺失值为 NaN
• series_names -> List[str]: 列对应的序列名称
• shape -> tuple: values 的形状

方法：

• __len__(): 返回行数
• __getitem__(key): 索引访问 values，支持 result[i] 和 result[i, j]
• show(max_rows=None): 格式化打印表格，可指定最大行数

TsFileDataFrame 使用

基本使用

from iotdb_ai import TsFileDataFrame

# 从单个文件创建
df = TsFileDataFrame("data/sensor_data.tsfile")

# 从多个文件创建（同名序列自动合并）
df = TsFileDataFrame([
    "data/2024-01.tsfile",
    "data/2024-02.tsfile",
])

# 从目录创建（自动扫描目录下所有 .tsfile 文件）
df = TsFileDataFrame("data/")

# 查看概览
print(df)
# TsFileDataFrame(5 time series, 2 files)
#      ts_name            start_time              end_time  count
#  0   weather.Beijing.temperature   2024-01-01 00:00:00   2024-06-30 23:00:00   4344
#  1   weather.Beijing.humidity      2024-01-01 00:00:00   2024-06-30 23:00:00   4344
#  ...

访问单条时间序列

# 按索引访问
ts = df[0]              # 第一条序列
ts = df[-1]             # 最后一条序列

# 按名称访问
ts = df['weather.Beijing.humidity']

# 查看序列信息
print(ts)               # Timeseries('weather.Beijing.humidity', count=4344, ...)
print(len(ts))          # 4344
print(ts.stats)         # {'start_time': ..., 'end_time': ..., 'count': 4344}

# 行索引访问数据
val = ts[0]             # 第一个值
vals = ts[100:200]      # 第 100~199 行的值（返回 np.ndarray）

批量访问序列

# 切片
series_list = df[1:3]       # 返回 List[Timeseries]

# 列表索引
series_list = df[[0, 2, 4]] # 返回 List[Timeseries]

# 列出满足前缀的序列名
names = df.list_timeseries("weather.Beijing")

时间对齐查询 (.loc)

# 按时间范围 + 序列名查询，返回 AlignedTimeseries
result = df.loc[
    1704067200000:1704153600000,       # 时间范围（毫秒时间戳）
    ['weather.Beijing.temperature',
     'weather.Beijing.humidity']
]

# 也支持按序列索引号
result = df.loc[1704067200000:1704153600000, [0, 1]]

print(result)
# AlignedTimeseries(24 rows, 2 series)
#            timestamp  weather.Beijing.temperature  weather.Beijing.humidity
#  2024-01-01 00:00:00                        -2.50                    45.00
#  2024-01-01 01:00:00                        -3.10                    47.20
#  ...

# 访问底层数据
result.timestamps   # np.ndarray, 毫秒时间戳
result.values       # np.ndarray, shape (rows, cols)
result.series_names # ['weather.Beijing.temperature', 'weather.Beijing.humidity']
result[0]           # 第一行所有列的值
result[0, 1]        # 第一行第二列的值

# 控制显示行数
result.show()       # 显示全部行
result.show(50)     # 显示前 50 行

上下文管理器

with TsFileDataFrame("data/sensor_data.tsfile") as df:
    ts = df[0]
    vals = ts[:100]
# 自动关闭文件句柄

🎓 示例

查看 examples/ 目录获取更多示例：

• tsfile_basic_usage.py - 基本使用方法
• 更多示例即将推出...

🏗️ 架构设计

核心组件

1. TsFileDataset (iotdb_ai/tsfile/dataset.py)

   • 主数据集类，实现 PyTorch Dataset 接口
   • 高效的索引结构（二分查找 + 前缀和）
   • 支持懒加载和预加载模式

2. TsFileDataFrame (iotdb_ai/tsfile/dataframe.py)

   • 多文件统一视图，DataFrame 式访问接口
   • 时间对齐查询（.loc），自动填充 NaN
   • 跨文件同名序列自动合并

3. TsFileReader (iotdb_ai/tsfile/reader.py)

   • 底层 TsFile 读取器
   • 双模式：快速解析 + IoTDB 官方库回退
   • 优化的批量读取

性能特性

• O(log n) 查找：使用二分查找快速定位时间序列
• 批量读取：一次读取完整序列，减少 I/O 次数
• 多进程支持：与 PyTorch DataLoader 无缝集成
• 内存优化：按需加载，可选预加载

🔧 开发

安装开发依赖

pip install -e ".[dev]"

运行测试

pytest tests/ -v

代码格式化

black iotdb_ai/
flake8 iotdb_ai/

📝 未来规划

• 多变量时间序列支持
• 时间范围过滤
• 序列名称/模式过滤
• 数据增强功能
• 更多预训练模型
• 分布式训练支持

🤝 贡献

欢迎提交 Issue 和 Pull Request！

📄 License

Apache License 2.0

📮 联系

• 项目主页: https://iotdb.apache.org/
• 问题反馈: GitHub Issues

---

注意： 这是一个早期版本，API 可能会有变化。欢迎反馈和建议！