Skip to main content

IoTDB AI Toolkit - High-performance time series machine learning toolkit

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.3.tar.gz (56.9 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.3-py3-none-any.whl (21.1 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for iotdb_ai-0.1.3.tar.gz
Algorithm Hash digest
SHA256 4bdf27091b8fe8c4e8b9e8729113c9cea7f761d1e4ca290b368fc68266606fb5
MD5 5d76915468642df0325164ebb05e5434
BLAKE2b-256 c71cbb14c7cba7c0a2ea25112643241709e898fd3923dd9e443ae6de212ad177

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for iotdb_ai-0.1.3-py3-none-any.whl
Algorithm Hash digest
SHA256 c4e491a1436a61eb8b316fdf4a5c666830f262806cb903870850114a6699be95
MD5 e8c08b8655728f1bdb453b1e963a4f9d
BLAKE2b-256 524d5727dfa4d0d0fb0b54d46fbb2f0e7df5d9800d8a907a311d45cdc4c62684

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