Add daily derived data pipeline
This commit is contained in:
Yuxuan Yan
@@ -0,0 +1,261 @@
|
||||
"""Tests for daily derived-data ingestion and plugins."""
|
||||
|
||||
import textwrap
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from cli import cli
|
||||
from pipeline.alpha.compute import join_feature_frames
|
||||
from pipeline.derived.compute import compute_derived, validate_derived_frame
|
||||
from pipeline.derived.registry import available_derived, get_derived, load_derived_module
|
||||
|
||||
|
||||
def _daily_bars() -> pd.DataFrame:
|
||||
return pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sz000001", "sh600000"],
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-03"]),
|
||||
"open": [10.0, 20.0, 11.0],
|
||||
"close": [10.5, 20.5, 11.5],
|
||||
"volume": [1000.0, 2000.0, 1200.0],
|
||||
})
|
||||
|
||||
|
||||
def _minute_bars() -> pd.DataFrame:
|
||||
return pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000", "sz000001"],
|
||||
"datetime": pd.to_datetime([
|
||||
"2024-01-02 09:35:00",
|
||||
"2024-01-02 09:40:00",
|
||||
"2024-01-02 09:35:00",
|
||||
]),
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-02"]),
|
||||
"time": ["09:35:00", "09:40:00", "09:35:00"],
|
||||
"open": [10.0, 10.5, 20.0],
|
||||
"high": [11.0, 12.0, 21.0],
|
||||
"low": [9.0, 10.0, 19.0],
|
||||
"close": [10.5, 11.0, 20.5],
|
||||
"volume": [100.0, 300.0, 200.0],
|
||||
"amount": [1000.0, 3300.0, 4100.0],
|
||||
})
|
||||
|
||||
|
||||
def test_validate_derived_frame_normalizes_and_sorts():
|
||||
result = validate_derived_frame(pd.DataFrame({
|
||||
"symbol_id": ["sz000001", "sh600000"],
|
||||
"date": ["2024-01-02 15:00:00", "2024-01-02 09:30:00"],
|
||||
"custom_value": [2.0, 1.0],
|
||||
}))
|
||||
|
||||
assert result["symbol_id"].tolist() == ["sh600000", "sz000001"]
|
||||
assert result["date"].tolist() == [
|
||||
pd.Timestamp("2024-01-02"),
|
||||
pd.Timestamp("2024-01-02"),
|
||||
]
|
||||
|
||||
|
||||
def test_validate_derived_frame_rejects_missing_keys():
|
||||
with pytest.raises(ValueError, match="missing required"):
|
||||
validate_derived_frame(pd.DataFrame({"symbol_id": ["sh600000"], "x": [1.0]}))
|
||||
|
||||
|
||||
def test_validate_derived_frame_rejects_duplicate_normalized_keys():
|
||||
with pytest.raises(ValueError, match="duplicate symbol_id,date"):
|
||||
validate_derived_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000"],
|
||||
"date": ["2024-01-02 09:30:00", "2024-01-02 15:00:00"],
|
||||
"x": [1.0, 2.0],
|
||||
}))
|
||||
|
||||
|
||||
def test_validate_derived_frame_rejects_duplicate_columns():
|
||||
bad = pd.DataFrame(
|
||||
[["sh600000", pd.Timestamp("2024-01-02"), 1.0, 2.0]],
|
||||
columns=["symbol_id", "date", "dup", "dup"],
|
||||
)
|
||||
with pytest.raises(ValueError, match="duplicate columns"):
|
||||
validate_derived_frame(bad)
|
||||
|
||||
|
||||
def test_validate_derived_frame_rejects_non_numeric_values():
|
||||
with pytest.raises(ValueError, match="numeric"):
|
||||
validate_derived_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"date": [pd.Timestamp("2024-01-02")],
|
||||
"bad": ["not numeric"],
|
||||
}))
|
||||
|
||||
|
||||
def test_derived_ingest_cli_accepts_csv_and_parquet(tmp_path):
|
||||
runner = CliRunner()
|
||||
source = pd.DataFrame({
|
||||
"symbol_id": ["sz000001", "sh600000"],
|
||||
"date": ["2024-01-02", "2024-01-02"],
|
||||
"custom_value": [2.0, 1.0],
|
||||
})
|
||||
csv_path = tmp_path / "custom.csv"
|
||||
parquet_path = tmp_path / "custom.pq"
|
||||
out_dir = tmp_path / "derived"
|
||||
source.to_csv(csv_path, index=False)
|
||||
source.to_parquet(parquet_path, index=False)
|
||||
|
||||
csv_result = runner.invoke(cli, [
|
||||
"derived", "ingest",
|
||||
"--input-path", str(csv_path),
|
||||
"--derived-name", "csv_custom",
|
||||
"--output-dir", str(out_dir),
|
||||
])
|
||||
assert csv_result.exit_code == 0, csv_result.output
|
||||
|
||||
parquet_result = runner.invoke(cli, [
|
||||
"derived", "ingest",
|
||||
"--input-path", str(parquet_path),
|
||||
"--derived-name", "parquet_custom",
|
||||
"--output-dir", str(out_dir),
|
||||
])
|
||||
assert parquet_result.exit_code == 0, parquet_result.output
|
||||
|
||||
written = pd.read_parquet(out_dir / "csv_custom.pq")
|
||||
assert written["symbol_id"].tolist() == ["sh600000", "sz000001"]
|
||||
assert (out_dir / "parquet_custom.pq").exists()
|
||||
|
||||
|
||||
def test_derived_validate_cli_rejects_duplicate_csv_columns(tmp_path):
|
||||
runner = CliRunner()
|
||||
csv_path = tmp_path / "bad.csv"
|
||||
csv_path.write_text("symbol_id,date,x,x\nsh600000,2024-01-02,1.0,2.0\n")
|
||||
|
||||
result = runner.invoke(cli, [
|
||||
"derived", "validate",
|
||||
"--input-path", str(csv_path),
|
||||
])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "duplicate columns" in result.output
|
||||
|
||||
|
||||
def test_external_derived_plugin_loads_filters_params_and_uses_inputs(tmp_path):
|
||||
module_path = tmp_path / "external_derived.py"
|
||||
module_path.write_text(textwrap.dedent('''
|
||||
import pandas as pd
|
||||
from pipeline.derived.base import BaseDerivedData
|
||||
from pipeline.derived.registry import register_derived
|
||||
|
||||
@register_derived
|
||||
class FlexibleDerived(BaseDerivedData):
|
||||
name = "flexible_derived_test"
|
||||
|
||||
def __init__(self, scale: float = 1.0):
|
||||
self.scale = scale
|
||||
|
||||
def compute(self, daily=None, minute=None) -> pd.DataFrame:
|
||||
result = None
|
||||
if daily is not None:
|
||||
result = daily[["symbol_id", "date", "close"]].copy()
|
||||
result["daily_scaled_close"] = result.pop("close") * self.scale
|
||||
if minute is not None:
|
||||
minute_out = (
|
||||
minute.groupby(["symbol_id", "date"], as_index=False)["volume"]
|
||||
.sum()
|
||||
.rename(columns={"volume": "minute_volume_sum"})
|
||||
)
|
||||
minute_out["minute_volume_sum"] *= self.scale
|
||||
result = minute_out if result is None else result.merge(
|
||||
minute_out, on=["symbol_id", "date"], how="left"
|
||||
)
|
||||
return result
|
||||
'''))
|
||||
|
||||
load_derived_module(str(module_path))
|
||||
assert "flexible_derived_test" in available_derived()
|
||||
|
||||
instance = get_derived("flexible_derived_test", scale=2.0, ignored=99)
|
||||
assert instance.scale == 2.0
|
||||
assert not hasattr(instance, "ignored")
|
||||
|
||||
daily_result = compute_derived(
|
||||
"flexible_derived_test",
|
||||
daily=_daily_bars(),
|
||||
scale=2.0,
|
||||
ignored=99,
|
||||
)
|
||||
assert "daily_scaled_close" in daily_result.columns
|
||||
assert np.isclose(daily_result["daily_scaled_close"].iloc[0], 21.0)
|
||||
|
||||
minute_result = compute_derived(
|
||||
"flexible_derived_test",
|
||||
minute=_minute_bars(),
|
||||
scale=2.0,
|
||||
)
|
||||
assert "minute_volume_sum" in minute_result.columns
|
||||
assert np.isclose(
|
||||
minute_result.loc[minute_result["symbol_id"] == "sh600000", "minute_volume_sum"].iloc[0],
|
||||
800.0,
|
||||
)
|
||||
|
||||
both_result = compute_derived(
|
||||
"flexible_derived_test",
|
||||
daily=_daily_bars(),
|
||||
minute=_minute_bars(),
|
||||
scale=1.0,
|
||||
)
|
||||
assert {"daily_scaled_close", "minute_volume_sum"}.issubset(both_result.columns)
|
||||
|
||||
|
||||
def test_derived_compute_cli_writes_builtin_minute_summary(tmp_path):
|
||||
runner = CliRunner()
|
||||
minute_path = tmp_path / "minute.pq"
|
||||
out_dir = tmp_path / "derived"
|
||||
_minute_bars().to_parquet(minute_path, index=False)
|
||||
|
||||
result = runner.invoke(cli, [
|
||||
"derived", "compute",
|
||||
"--minute-path", str(minute_path),
|
||||
"--derived-type", "minute_daily_summary",
|
||||
"--derived-name", "minute_summary",
|
||||
"--output-dir", str(out_dir),
|
||||
])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
written = pd.read_parquet(out_dir / "minute_summary.pq")
|
||||
assert "minute_vwap" in written.columns
|
||||
|
||||
|
||||
def test_alpha_feature_join_rejects_derived_column_collisions():
|
||||
data = _daily_bars()
|
||||
derived_a = data[["symbol_id", "date"]].copy()
|
||||
derived_a["custom_value"] = 1.0
|
||||
derived_b = data[["symbol_id", "date"]].copy()
|
||||
derived_b["custom_value"] = 2.0
|
||||
|
||||
with pytest.raises(ValueError, match="conflict"):
|
||||
join_feature_frames(data, [derived_a, derived_b])
|
||||
|
||||
close_collision = data[["symbol_id", "date"]].copy()
|
||||
close_collision["close"] = 1.0
|
||||
with pytest.raises(ValueError, match="conflict"):
|
||||
join_feature_frames(data, [close_collision])
|
||||
|
||||
|
||||
def test_legacy_feature_cli_delegates_to_derived_registry(tmp_path):
|
||||
runner = CliRunner()
|
||||
minute_path = tmp_path / "minute.pq"
|
||||
out_dir = tmp_path / "features"
|
||||
_minute_bars().to_parquet(minute_path, index=False)
|
||||
|
||||
list_result = runner.invoke(cli, ["feature", "list"])
|
||||
assert list_result.exit_code == 0, list_result.output
|
||||
assert "minute_daily_summary" in list_result.output
|
||||
|
||||
compute_result = runner.invoke(cli, [
|
||||
"feature", "compute",
|
||||
"--minute-path", str(minute_path),
|
||||
"--feature-type", "minute_daily_summary",
|
||||
"--feature-name", "minute_summary",
|
||||
"--output-dir", str(out_dir),
|
||||
])
|
||||
assert compute_result.exit_code == 0, compute_result.output
|
||||
assert (out_dir / "minute_summary.pq").exists()
|
||||
|
||||
Reference in New Issue
Block a user