Skip to main content

IoTDB AI Toolkit - High-performance time series data toolkit for IoTDB

Project description

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_namestart_timeend_timecount
  • 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_timeend_timecount

方法:

  • __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

📮 联系


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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

iotdb_ai-0.1.3b1.tar.gz (50.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

iotdb_ai-0.1.3b1-py3-none-any.whl (21.9 kB view details)

Uploaded Python 3

File details

Details for the file iotdb_ai-0.1.3b1.tar.gz.

File metadata

  • Download URL: iotdb_ai-0.1.3b1.tar.gz
  • Upload date:
  • Size: 50.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.9

File hashes

Hashes for iotdb_ai-0.1.3b1.tar.gz
Algorithm Hash digest
SHA256 7d8c58b4e6cf3d5205ee1b0d9e78d41946475d746e58ee4f949a763ab73e5153
MD5 8231aac0b369774d205da0d6cc791e44
BLAKE2b-256 2a9e4619b817f91dcf46728f1fbb8e5bdc727d8f56b145ab2aabdc3ea7479099

See more details on using hashes here.

File details

Details for the file iotdb_ai-0.1.3b1-py3-none-any.whl.

File metadata

  • Download URL: iotdb_ai-0.1.3b1-py3-none-any.whl
  • Upload date:
  • Size: 21.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.9

File hashes

Hashes for iotdb_ai-0.1.3b1-py3-none-any.whl
Algorithm Hash digest
SHA256 684d1b18e74e40c529bff6d1f491e543eafb6086359b73ba2ccb327cd9297c3d
MD5 5ccc2120f6c7a981dae59901f0c485fc
BLAKE2b-256 ffdaa7c2f53667fbb98c3ad10c880da6d5f340c890afc4c69802dd33b88c0ada

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page