perf(factor): fast_ops v2 去over化——rolling.over(vt_symbol)在10M行分组滚动分钟级/算子(py-spy实锤ts_mean),改符号内shift展开统一模式(ts_mean/std/min/max/sum/corr/cov),等价性测试对齐原版+1M行速度护栏 [vps]

This commit is contained in:
2026-08-25 22:15:40 +08:00
parent a2fe51b4b0
commit 7bbba554f8
2 changed files with 412 additions and 5 deletions
+184 -3
View File
@@ -28,6 +28,13 @@ from sanguo_factor.fast_ops import (
fast_ts_rsquare,
fast_ts_resi,
fast_ts_quantile,
fast_ts_mean,
fast_ts_std,
fast_ts_sum,
fast_ts_min,
fast_ts_max,
fast_ts_corr_v2,
fast_ts_cov_v2,
)
@@ -282,6 +289,161 @@ class TestFastTsQuantile:
assert_dataproxy_equal(original, fast, rtol=1e-9)
class TestFastTsMean:
"""Test fast_ts_mean equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values = np.random.randn(100) * 10 + 50
feature = create_dataproxy(values)
original = ts_function.ts_mean(feature, window=10)
fast = fast_ts_mean(feature, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-9)
def test_with_nan(self):
"""Handle NaN correctly."""
values = np.array([1.0, np.nan, 3.0, np.nan, 5.0] * 20)
feature = create_dataproxy(values)
original = ts_function.ts_mean(feature, window=5)
fast = fast_ts_mean(feature, window=5)
assert_dataproxy_equal(original, fast, check_nan=False)
class TestFastTsStd:
"""Test fast_ts_std equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values = np.random.randn(100) * 10 + 50
feature = create_dataproxy(values)
original = ts_function.ts_std(feature, window=10)
fast = fast_ts_std(feature, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-7)
def test_constant_values(self):
"""Constant values should have std = 0."""
values = np.array([5.0] * 100)
feature = create_dataproxy(values)
result = fast_ts_std(feature, window=10)
df = result.df
std_values = df.filter(pl.col("data").is_not_null())["data"].to_numpy()
assert np.all(std_values < 1e-10), "Std of constant values should be ~0"
class TestFastTsSum:
"""Test fast_ts_sum equivalence."""
def test_random_data(self):
"""Exact match on random data (ts_sum has no min_samples, so partial windows are NaN)."""
np.random.seed(42)
values = np.random.randn(100) * 10 + 50
feature = create_dataproxy(values)
original = ts_function.ts_sum(feature, window=10)
fast = fast_ts_sum(feature, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-9)
def test_partial_windows_nan(self):
"""ts_sum returns NaN for partial windows (no min_samples=1)."""
values = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0] * 10)
feature = create_dataproxy(values)
result = fast_ts_sum(feature, window=5)
# First 4 values should be NaN (partial window)
df = result.df
first_four = df["data"].to_numpy()[:4]
assert all(np.isnan(first_four)), "Partial windows should be NaN"
class TestFastTsMin:
"""Test fast_ts_min equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values = np.random.randn(100) * 10 + 50
feature = create_dataproxy(values)
original = ts_function.ts_min(feature, window=10)
fast = fast_ts_min(feature, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-9)
class TestFastTsMax:
"""Test fast_ts_max equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values = np.random.randn(100) * 10 + 50
feature = create_dataproxy(values)
original = ts_function.ts_max(feature, window=10)
fast = fast_ts_max(feature, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-9)
class TestFastTsCorrV2:
"""Test fast_ts_corr_v2 (over-free) equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values1 = np.random.randn(100) * 10 + 50
values2 = np.random.randn(100) * 10 + 50
feature1 = create_dataproxy(values1)
feature2 = create_dataproxy(values2)
original = ts_function.ts_corr(feature1, feature2, window=10)
fast = fast_ts_corr_v2(feature1, feature2, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-7)
def test_perfect_correlation(self):
"""Perfect correlation should be exactly 1.0."""
values = np.random.randn(100) * 10 + 50
feature1 = create_dataproxy(values)
feature2 = create_dataproxy(values * 2 + 10)
result = fast_ts_corr_v2(feature1, feature2, window=10)
df = result.df
assert (df.filter(pl.col("data").is_not_null())["data"] - 1.0).abs().max() < 1e-6
class TestFastTsCovV2:
"""Test fast_ts_cov_v2 (over-free) equivalence."""
def test_random_data(self):
"""Exact match on random data."""
np.random.seed(42)
values1 = np.random.randn(100) * 10 + 50
values2 = np.random.randn(100) * 10 + 50
feature1 = create_dataproxy(values1)
feature2 = create_dataproxy(values2)
original = ts_function.ts_cov(feature1, feature2, window=10)
fast = fast_ts_cov_v2(feature1, feature2, window=10)
assert_dataproxy_equal(original, fast, rtol=1e-7)
class TestRegistrationAndIntegration:
"""Test registration and vnpy integration."""
@@ -291,7 +453,7 @@ class TestRegistrationAndIntegration:
overrides2 = register_fast_ops()
assert overrides1 == overrides2
assert len(overrides1) == 8
assert len(overrides1) == 13
def test_expression_override(self):
"""Test that registered functions override vnpy defaults."""
@@ -325,8 +487,9 @@ class TestRegistrationAndIntegration:
overrides = register_fast_ops()
expected = {
"ts_rank", "ts_corr", "ts_cov",
"ts_decay_linear", "ts_slope", "ts_rsquare", "ts_resi", "ts_quantile"
"ts_mean", "ts_std", "ts_sum", "ts_min", "ts_max",
"ts_corr", "ts_cov",
"ts_rank", "ts_decay_linear", "ts_slope", "ts_rsquare", "ts_resi", "ts_quantile"
}
assert set(overrides) == expected
@@ -390,6 +553,24 @@ class TestEdgeCases:
# Should complete in under 10 seconds for 200k rows
assert elapsed < 10, f"Performance regression: {elapsed:.2f}s > 10s"
def test_ts_mean_overfree_speed_guard(self):
"""Speed guard for over-free ts_mean (critical hot path)."""
import time
n = 1_000_000 # 1M rows (1/10 of full 10M dataset)
df = pl.DataFrame({
"vt_symbol": ["A"] * (n // 2) + ["B"] * (n // 2),
"datetime": list(range(n // 2)) * 2,
"data": [((i * 37) % 997) / 997 for i in range(n // 2)] * 2
})
feature = DataProxy(df)
t0 = time.time()
fast_ts_mean(feature, 20)
elapsed = time.time() - t0
# .over() version would take minutes; target is <10 seconds
assert elapsed < 10, f"Performance regression: {elapsed:.2f}s > 10s"
if __name__ == "__main__":
pytest.main([__file__, "-v"])