liaozhaorun/ProStock
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: HDF5迁移至DuckDB存储

Browse Source 
- 新增DuckDB Storage与ThreadSafeStorage实现
- 新增db_manager模块支持增量同步策略
- DataLoader与Sync模块适配DuckDB
- 补充迁移相关文档与测试
- 修复README文档链接
This commit is contained in:
liaozhaorun
2026-02-23 00:07:21 +08:00
parent 0a16129548
commit e58b39970c
14 changed files with 2265 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
tests/factors/test_data_loader.py
View File

@@ -1,14 +1,16 @@
"""测试数据加载器 - DataLoader
测试需求（来自 factor_implementation_plan.md）：
- 测试从单个 H5 文件加载数据
- 测试从多个 H5 文件加载并合并
- 测试从 DuckDB 加载数据
- 测试从多个查询加载并合并
- 测试列选择（只加载需要的列）
- 测试缓存机制（第二次加载更快）
- 测试 clear_cache() 清空缓存
- 测试按 date_range 过滤
- 测试文件不存在时抛出 FileNotFoundError
- 测试表不存在时的处理
- 测试列不存在时抛出 KeyError
使用 3 个月的真实数据进行测试 (2024年1月-3月)
"""
import pytest
@@ -22,6 +24,10 @@ from src.factors import DataSpec, DataLoader
class TestDataLoaderBasic:
"""测试 DataLoader 基本功能"""
# 测试数据时间范围：3个月
TEST_START_DATE = "20240101"
TEST_END_DATE = "20240331"
@pytest.fixture
def loader(self):
"""创建 DataLoader 实例"""
@@ -34,7 +40,7 @@ class TestDataLoaderBasic:
assert loader._cache == {}
def test_load_single_source(self, loader):
"""测试从单个 H5 文件加载数据"""
"""测试从 DuckDB 加载数据"""
specs = [
DataSpec(
source="daily",
@@ -43,7 +49,8 @@ class TestDataLoaderBasic:
)
]
df = loader.load(specs)
# 使用 3 个月日期范围限制数据量
df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
assert isinstance(df, pl.DataFrame)
assert len(df) > 0
@@ -51,10 +58,29 @@ class TestDataLoaderBasic:
assert "trade_date" in df.columns
assert "close" in df.columns
def test_load_multiple_sources(self, loader):
"""测试从多个 H5 文件加载并合并"""
# 注意：这里假设只有一个 daily.h5 文件
# 如果有多个文件，可以测试合并逻辑
def test_load_with_date_range(self, loader):
"""测试加载特定日期范围（3个月）"""
specs = [
DataSpec(
source="daily",
columns=["ts_code", "trade_date", "close", "open", "high", "low"],
lookback_days=1,
)
]
df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
assert isinstance(df, pl.DataFrame)
assert len(df) > 0
# 验证日期范围
if len(df) > 0:
dates = df["trade_date"].to_list()
assert all(self.TEST_START_DATE <= d <= self.TEST_END_DATE for d in dates)
print(f"[TEST] Loaded {len(df)} rows from {min(dates)} to {max(dates)}")
def test_load_multiple_specs(self, loader):
"""测试从多个 DataSpec 加载并合并"""
specs = [
DataSpec(
source="daily",
@@ -68,7 +94,7 @@ class TestDataLoaderBasic:
),
]
df = loader.load(specs)
df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
assert isinstance(df, pl.DataFrame)
assert len(df) > 0
@@ -92,13 +118,13 @@ class TestDataLoaderBasic:
)
]
df = loader.load(specs)
df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
# 只应该有 3 列
assert set(df.columns) == {"ts_code", "trade_date", "close"}
def test_date_range_filter(self, loader):
"""测试按 date_range 过滤"""
"""测试按 date_range 过滤 - 使用3个月数据的不同子集"""
specs = [
DataSpec(
source="daily",
@@ -107,11 +133,13 @@ class TestDataLoaderBasic:
)
]
# 先加载所有数据
df_all = loader.load(specs)
# 加载完整的3个月数据
df_all = loader.load(
specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE)
)
total_rows = len(df_all)
# 清空缓存，重新加载特定日期范围
# 清空缓存，重新加载1个月数据
loader.clear_cache()
df_filtered = loader.load(specs, date_range=("20240101", "20240131"))
@@ -127,6 +155,9 @@ class TestDataLoaderBasic:
class TestDataLoaderCache:
"""测试 DataLoader 缓存机制"""
TEST_START_DATE = "20240101"
TEST_END_DATE = "20240331"
@pytest.fixture
def loader(self):
"""创建 DataLoader 实例"""
@@ -143,7 +174,7 @@ class TestDataLoaderCache:
]
# 第一次加载
loader.load(specs)
loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
# 检查缓存
assert len(loader._cache) > 0
@@ -162,20 +193,20 @@ class TestDataLoaderCache:
# 第一次加载
start = time.time()
df1 = loader.load(specs)
df1 = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
time1 = time.time() - start
# 第二次加载（应该使用缓存）
start = time.time()
df2 = loader.load(specs)
df2 = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
time2 = time.time() - start
# 数据应该相同
assert df1.shape == df2.shape
# 第二次应该更快（至少快 50%）
# 注意：如果数据量很小，这个测试可能不稳定
# assert time2 < time1 * 0.5
# 第二次应该更快
print(f"[TEST] First load: {time1:.3f}s, cached load: {time2:.3f}s")
assert time2 < time1, "Cached load should be faster"
def test_clear_cache(self, loader):
"""测试 clear_cache() 清空缓存"""
@@ -188,7 +219,7 @@ class TestDataLoaderCache:
]
# 加载数据
loader.load(specs)
loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
assert len(loader._cache) > 0
# 清空缓存
@@ -210,7 +241,7 @@ class TestDataLoaderCache:
assert info_before["entries"] == 0
# 加载后
loader.load(specs)
loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
info_after = loader.get_cache_info()
assert info_after["entries"] > 0
assert info_after["total_rows"] > 0
@@ -219,18 +250,19 @@ class TestDataLoaderCache:
class TestDataLoaderErrors:
"""测试 DataLoader 错误处理"""
def test_file_not_found(self):
"""测试文件不存在时抛出 FileNotFoundError"""
loader = DataLoader(data_dir="nonexistent_dir")
def test_table_not_exists(self):
"""测试表不存在时的处理"""
loader = DataLoader(data_dir="data")
specs = [
DataSpec(
source="daily",
source="nonexistent_table",
columns=["ts_code", "trade_date", "close"],
lookback_days=1,
)
]
with pytest.raises(FileNotFoundError):
# 应该返回空 DataFrame 或抛出异常
with pytest.raises(Exception):
loader.load(specs)
def test_column_not_found(self):
@@ -246,3 +278,7 @@ class TestDataLoaderErrors:
with pytest.raises(KeyError, match="nonexistent_column"):
loader.load(specs)
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])