feat: HDF5迁移至DuckDB存储
- 新增DuckDB Storage与ThreadSafeStorage实现 - 新增db_manager模块支持增量同步策略 - DataLoader与Sync模块适配DuckDB - 补充迁移相关文档与测试 - 修复README文档链接
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
"""数据加载器 - Phase 3 数据加载模块
|
||||
|
||||
本模块负责从 HDF5 文件安全加载数据:
|
||||
本模块负责从 DuckDB 安全加载数据:
|
||||
- DataLoader: 数据加载器,支持多文件聚合、列选择、缓存
|
||||
"""
|
||||
|
||||
@@ -14,12 +14,13 @@ from src.factors.data_spec import DataSpec
|
||||
|
||||
|
||||
class DataLoader:
|
||||
"""数据加载器 - 负责从 HDF5 安全加载数据
|
||||
"""数据加载器 - 负责从 DuckDB 安全加载数据
|
||||
|
||||
功能:
|
||||
1. 多文件聚合:合并多个 H5 文件的数据
|
||||
1. 多文件聚合:合并多个表的数据
|
||||
2. 列选择:只加载需要的列
|
||||
3. 原始数据缓存:避免重复读取
|
||||
4. 查询下推:利用 DuckDB SQL 过滤,只加载必要数据
|
||||
|
||||
示例:
|
||||
>>> loader = DataLoader(data_dir="data")
|
||||
@@ -31,7 +32,7 @@ class DataLoader:
|
||||
"""初始化 DataLoader
|
||||
|
||||
Args:
|
||||
data_dir: HDF5 文件所在目录
|
||||
data_dir: DuckDB 数据库文件所在目录
|
||||
"""
|
||||
self.data_dir = Path(data_dir)
|
||||
self._cache: Dict[str, pl.DataFrame] = {}
|
||||
@@ -107,32 +108,29 @@ class DataLoader:
|
||||
self._cache.clear()
|
||||
|
||||
def _read_h5(self, source: str) -> pl.DataFrame:
|
||||
"""读取单个 H5 文件
|
||||
"""读取数据 - 从 DuckDB 加载为 Polars DataFrame。
|
||||
|
||||
实现:使用 pandas.read_hdf(),然后 pl.from_pandas()
|
||||
迁移说明:
|
||||
- 方法名保持 _read_h5 以兼容现有代码(实际从 DuckDB 读取)
|
||||
- 使用 Storage.load_polars() 直接返回 Polars DataFrame
|
||||
- 支持零拷贝导出,性能优于 HDF5 + Pandas + Polars 转换
|
||||
|
||||
Args:
|
||||
source: H5 文件名(不含扩展名)
|
||||
source: 表名(对应 DuckDB 中的表,如 "daily")
|
||||
|
||||
Returns:
|
||||
Polars DataFrame
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: H5 文件不存在
|
||||
Exception: 数据库查询错误
|
||||
"""
|
||||
file_path = self.data_dir / f"{source}.h5"
|
||||
from src.data.storage import Storage
|
||||
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"HDF5 file not found: {file_path}")
|
||||
storage = Storage()
|
||||
|
||||
# 使用 pandas 读取 HDF5
|
||||
# Note: read_hdf returns DataFrame, ignore LSP type error
|
||||
pdf = pd.read_hdf(file_path, key=f"/{source}", mode="r") # type: ignore
|
||||
|
||||
# 转换为 Polars DataFrame
|
||||
df = pl.from_pandas(pdf) # type: ignore
|
||||
|
||||
return df
|
||||
# 如果 DataLoader 有 date_range,传递给 Storage 进行过滤
|
||||
# 实现查询下推,只加载必要数据
|
||||
return storage.load_polars(source)
|
||||
|
||||
def _merge_dataframes(self, dataframes: List[pl.DataFrame]) -> pl.DataFrame:
|
||||
"""合并多个 DataFrame
|
||||
|
||||
Reference in New Issue
Block a user