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]

sanguo_factor/fast_ops.py

@@ -288,6 +288,220 @@ def fast_ts_quantile(feature: DataProxy, window: int, quantile: float) -> DataPr
     return DataProxy(df)
 
 
+def fast_ts_mean(feature: DataProxy, window: int) -> DataProxy:
+    """Mean over rolling window with shift expansion (over-free).
+
+    Uses: mean = Σ shift_k / count_nonnull(shift_k) for k in 0..window-1
+    Replicates: rolling_map(lambda s: np.nanmean(s), window, min_samples=1).over("vt_symbol")
+    """
+    current = pl.col("data")
+
+    # Create shifted versions with symbol boundary protection
+    shifts = [current.shift(k).over("vt_symbol") for k in range(window)]
+
+    # Sum all shifts, then count non-null values
+    sum_expr = pl.sum_horizontal(shifts)
+    count_expr = pl.sum_horizontal([pl.when(s.is_not_null()).then(1).otherwise(0) for s in shifts])
+
+    # Mean = sum / count (handles partial windows automatically)
+    mean_expr = sum_expr / count_expr
+
+    df = feature.df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        mean_expr.alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_std(feature: DataProxy, window: int) -> DataProxy:
+    """Standard deviation over rolling window with shift expansion (over-free).
+
+    Uses: std = sqrt(E[x²] - E[x]²) with shift-based mean computation
+    Replicates: rolling_map(lambda s: np.nanstd(s, ddof=0), window, min_samples=1).over("vt_symbol")
+    """
+    current = pl.col("data")
+
+    # Create shifted versions with symbol boundary protection
+    shifts = [current.shift(k).over("vt_symbol") for k in range(window)]
+
+    # Count non-null values
+    count_expr = pl.sum_horizontal([pl.when(s.is_not_null()).then(1).otherwise(0) for s in shifts])
+
+    # Mean
+    sum_expr = pl.sum_horizontal(shifts)
+    mean_expr = sum_expr / count_expr
+
+    # E[x²] = sum(shift²) / count
+    sum_sq_expr = pl.sum_horizontal([s.pow(2) for s in shifts])
+    mean_sq_expr = sum_sq_expr / count_expr
+
+    # std = sqrt(E[x²] - E[x]²)
+    std_expr = (mean_sq_expr - mean_expr.pow(2)).sqrt()
+
+    df = feature.df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        std_expr.alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_sum(feature: DataProxy, window: int) -> DataProxy:
+    """Sum over rolling window with shift expansion (over-free).
+
+    Uses: sum = Σ shift_k for k in 0..window-1
+    Replicates: rolling_sum(window).over("vt_symbol") (no min_samples, so partial windows are NaN)
+    """
+    current = pl.col("data")
+
+    # Create shifted versions and sum them
+    shifts = [current.shift(k).over("vt_symbol") for k in range(window)]
+    sum_expr = pl.sum_horizontal(shifts)
+
+    # Apply boundary masking - null for partial windows (first window-1 rows per symbol)
+    boundary_mask = pl.int_range(pl.len()).over("vt_symbol") >= (window - 1)
+    result = pl.when(boundary_mask).then(sum_expr).otherwise(None)
+
+    df = feature.df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        result.alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_min(feature: DataProxy, window: int) -> DataProxy:
+    """Minimum over rolling window with shift expansion (over-free).
+
+    Uses: min = min(shift_0, shift_1, ..., shift_{w-1})
+    Replicates: rolling_min(window, min_samples=1).over("vt_symbol")
+    """
+    current = pl.col("data")
+
+    # Create shifted versions and find minimum
+    shifts = [current.shift(k).over("vt_symbol") for k in range(window)]
+    min_expr = pl.min_horizontal(shifts)
+
+    df = feature.df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        min_expr.alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_max(feature: DataProxy, window: int) -> DataProxy:
+    """Maximum over rolling window with shift expansion (over-free).
+
+    Uses: max = max(shift_0, shift_1, ..., shift_{w-1})
+    Replicates: rolling_max(window, min_samples=1).over("vt_symbol")
+    """
+    current = pl.col("data")
+
+    # Create shifted versions and find maximum
+    shifts = [current.shift(k).over("vt_symbol") for k in range(window)]
+    max_expr = pl.max_horizontal(shifts)
+
+    df = feature.df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        max_expr.alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_corr_v2(feature1: DataProxy, feature2: DataProxy, window: int) -> DataProxy:
+    """Correlation with shift expansion (over-free).
+
+    Uses: corr = (E[xy] - E[x]E[y]) / (std_x * std_y)
+    All statistics computed via shift-based mean/std
+    """
+    df_merged = feature1.df.join(feature2.df, on=["datetime", "vt_symbol"])
+
+    x = pl.col("data")
+    y = pl.col("data_right")
+
+    # Create shifted versions for both series
+    shifts_x = [x.shift(k).over("vt_symbol") for k in range(window)]
+    shifts_y = [y.shift(k).over("vt_symbol") for k in range(window)]
+
+    # Count non-null pairs
+    count_expr = pl.sum_horizontal([
+        pl.when(sx.is_not_null() & sy.is_not_null()).then(1).otherwise(0)
+        for sx, sy in zip(shifts_x, shifts_y)
+    ])
+
+    # Means
+    mean_x = pl.sum_horizontal(shifts_x) / count_expr
+    mean_y = pl.sum_horizontal(shifts_y) / count_expr
+
+    # E[xy]
+    sum_xy = pl.sum_horizontal([sx * sy for sx, sy in zip(shifts_x, shifts_y)])
+    mean_xy = sum_xy / count_expr
+
+    # Standard deviations
+    sum_x_sq = pl.sum_horizontal([sx.pow(2) for sx in shifts_x])
+    sum_y_sq = pl.sum_horizontal([sy.pow(2) for sy in shifts_y])
+    var_x = (sum_x_sq / count_expr) - mean_x.pow(2)
+    var_y = (sum_y_sq / count_expr) - mean_y.pow(2)
+    std_x = var_x.sqrt()
+    std_y = var_y.sqrt()
+
+    # Correlation
+    corr_expr = (mean_xy - mean_x * mean_y) / (std_x * std_y)
+
+    df = df_merged.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        pl.when(corr_expr.is_infinite() | corr_expr.is_nan())
+        .then(None)
+        .otherwise(corr_expr)
+        .alias("data")
+    )
+    return DataProxy(df)
+
+
+def fast_ts_cov_v2(feature1: DataProxy, feature2: DataProxy, window: int) -> DataProxy:
+    """Covariance with shift expansion (over-free).
+
+    Uses: cov = corr * std_x * std_y
+    Delegates to fast_ts_corr_v2 for correlation
+    """
+    # Get correlation
+    corr_result = fast_ts_corr_v2(feature1, feature2, window)
+
+    # Get individual std deviations
+    std_x_result = fast_ts_std(feature1, window)
+    std_y_result = fast_ts_std(feature2, window)
+
+    # Merge and compute covariance
+    df_merged = corr_result.df.join(
+        std_x_result.df.select(["datetime", "vt_symbol", pl.col("data").alias("std_x")]),
+        on=["datetime", "vt_symbol"]
+    ).join(
+        std_y_result.df.select(["datetime", "vt_symbol", pl.col("data").alias("std_y")]),
+        on=["datetime", "vt_symbol"]
+    )
+
+    df = df_merged.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        (pl.col("data") * pl.col("std_x") * pl.col("std_y")).alias("data")
+    )
+
+    # Handle infinite/NaN values
+    df = df.select(
+        pl.col("datetime"),
+        pl.col("vt_symbol"),
+        pl.when(pl.col("data").is_infinite() | pl.col("data").is_nan())
+        .then(None)
+        .otherwise(pl.col("data"))
+        .alias("data")
+    )
+    return DataProxy(df)
+
+
 def register_fast_ops() -> list[str]:
     """Register fast polars operators into vnpy's EXPRESSION_FUNCTIONS.
 
@@ -299,9 +513,14 @@ def register_fast_ops() -> list[str]:
         print(f"Registered {len(overrides)} fast operators")
     """
     operators = {
+        "ts_mean": fast_ts_mean,
+        "ts_std": fast_ts_std,
+        "ts_sum": fast_ts_sum,
+        "ts_min": fast_ts_min,
+        "ts_max": fast_ts_max,
+        "ts_corr": fast_ts_corr_v2,
+        "ts_cov": fast_ts_cov_v2,
         "ts_rank": fast_ts_rank,
-        "ts_corr": fast_ts_corr,
-        "ts_cov": fast_ts_cov,
         "ts_decay_linear": fast_ts_decay_linear,
         "ts_slope": fast_ts_slope,
         "ts_rsquare": fast_ts_rsquare,
@@ -328,4 +547,11 @@ __all__ = [
     "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",
 ]

tests/factor/test_fast_ops.py

@@ -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"])