From 45fc062627ecdbc615037929f6aa6ec1a17a50e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=AE=97=E5=B9=B3?= Date: Thu, 11 Jun 2026 21:14:01 +0800 Subject: [PATCH] feat: add --value-col param to inject for multi-column CSV support --- ts_anomaly_td/cli.py | 2 +- ts_anomaly_td/io_csv.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/ts_anomaly_td/cli.py b/ts_anomaly_td/cli.py index 60567a3..47cdbe5 100644 --- a/ts_anomaly_td/cli.py +++ b/ts_anomaly_td/cli.py @@ -17,7 +17,7 @@ def cmd_inject(args): """inject 子命令：读 CSV → 建表 → 批量插入。""" logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") - ts, vals, labels = read_csv(args.csv) + ts, vals, labels = read_csv(args.csv, value_col=args.value_col) rows = list(zip(ts, vals, labels)) logger.info("read %d rows from %s", len(rows), args.csv) diff --git a/ts_anomaly_td/io_csv.py b/ts_anomaly_td/io_csv.py index c91308f..819e55b 100644 --- a/ts_anomaly_td/io_csv.py +++ b/ts_anomaly_td/io_csv.py @@ -4,11 +4,12 @@ from pathlib import Path import pandas as pd -def read_csv(path: str | Path) -> tuple[list[int], list[float], list[int]]: +def read_csv(path: str | Path, value_col: str = "value") -> tuple[list[int], list[float], list[int]]: """读 timestamp,value,label CSV 文件。 Args: path: CSV 文件路径 + value_col: 数值列名 (default: "value") Returns: (timestamps_ms, values, labels) 三元组 @@ -22,13 +23,13 @@ def read_csv(path: str | Path) -> tuple[list[int], list[float], list[int]]: """ df = pd.read_csv(path) - if "timestamp" not in df.columns or "value" not in df.columns: - raise ValueError(f"CSV 缺少 timestamp 或 value 列: {path}") + if "timestamp" not in df.columns or value_col not in df.columns: + raise ValueError(f"CSV 缺少 timestamp 或 {value_col} 列: {path}") timestamps = pd.to_datetime(df["timestamp"]) timestamps_ms = (timestamps.astype("int64") // 10**6).tolist() - values = df["value"].astype(float).tolist() + values = df[value_col].astype(float).tolist() labels = df["label"].fillna(0).astype(int).tolist() if "label" in df.columns else [0] * len(df)