Route bars through Rust core and PyO3 bindings

Sean-Koval · Sean-Koval · commit 9c2c0b686ba2 · 2026-02-14T02:28:19.000-05:00
diff --git a/crates/openquant/src/data_structures.rs b/crates/openquant/src/data_structures.rs
@@ -23,6 +23,7 @@ pub enum StandardBarType {
 /// Bar output with OHLCV-like fields.
 #[derive(Debug, Clone, PartialEq)]
 pub struct StandardBar {
+    pub start_timestamp: NaiveDateTime,
     pub timestamp: NaiveDateTime,
     pub open: f64,
     pub high: f64,
@@ -211,6 +212,7 @@ fn build_bar(trades: &[Trade]) -> StandardBar {
 
     let open = trades.first().expect("non-empty slice").price;
     let close = trades.last().expect("non-empty slice").price;
+    let start_timestamp = trades.first().expect("non-empty slice").timestamp;
     let timestamp = trades.last().expect("non-empty slice").timestamp;
     let (high, low) = trades.iter().fold((f64::NEG_INFINITY, f64::INFINITY), |(h, l), trade| {
         (h.max(trade.price), l.min(trade.price))
@@ -222,6 +224,7 @@ fn build_bar(trades: &[Trade]) -> StandardBar {
     });
 
     StandardBar {
+        start_timestamp,
         timestamp,
         open,
         high,
diff --git a/crates/pyopenquant/src/lib.rs b/crates/pyopenquant/src/lib.rs
@@ -1,4 +1,5 @@
 use nalgebra::DMatrix;
+use openquant::data_structures::{standard_bars, time_bars, StandardBarType, Trade};
 use openquant::filters::Threshold;
 use openquant::pipeline::{
     run_mid_frequency_pipeline, ResearchPipelineConfig, ResearchPipelineInput,
@@ -50,6 +51,61 @@ fn format_naive_datetimes(values: Vec<chrono::NaiveDateTime>) -> Vec<String> {
     values.into_iter().map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string()).collect()
 }
 
+fn parse_one_naive_datetime(value: &str) -> PyResult<chrono::NaiveDateTime> {
+    chrono::NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S")
+        .or_else(|_| {
+            chrono::NaiveDate::parse_from_str(value, "%Y-%m-%d")
+                .map(|d| d.and_hms_opt(0, 0, 0).expect("valid fixed midnight"))
+        })
+        .map_err(|e| {
+            PyValueError::new_err(format!(
+                "invalid datetime '{value}' (expected '%Y-%m-%d %H:%M:%S' or '%Y-%m-%d'): {e}"
+            ))
+        })
+}
+
+fn build_trades(
+    timestamps: Vec<String>,
+    prices: Vec<f64>,
+    volumes: Vec<f64>,
+) -> PyResult<Vec<Trade>> {
+    if timestamps.len() != prices.len() || prices.len() != volumes.len() {
+        return Err(PyValueError::new_err(format!(
+            "timestamps/prices/volumes length mismatch: {} / {} / {}",
+            timestamps.len(),
+            prices.len(),
+            volumes.len()
+        )));
+    }
+    let mut trades = Vec::with_capacity(prices.len());
+    for i in 0..prices.len() {
+        trades.push(Trade {
+            timestamp: parse_one_naive_datetime(&timestamps[i])?,
+            price: prices[i],
+            volume: volumes[i],
+        });
+    }
+    Ok(trades)
+}
+
+fn bars_to_rows(bars: Vec<openquant::data_structures::StandardBar>) -> Vec<(String, String, f64, f64, f64, f64, f64, f64, usize)> {
+    bars.into_iter()
+        .map(|b| {
+            (
+                b.start_timestamp.format("%Y-%m-%d %H:%M:%S").to_string(),
+                b.timestamp.format("%Y-%m-%d %H:%M:%S").to_string(),
+                b.open,
+                b.high,
+                b.low,
+                b.close,
+                b.volume,
+                b.dollar_value,
+                b.tick_count,
+            )
+        })
+        .collect()
+}
+
 #[pyfunction(name = "calculate_value_at_risk")]
 fn risk_calculate_value_at_risk(returns: Vec<f64>, confidence_level: f64) -> PyResult<f64> {
     RiskMetrics::default().calculate_value_at_risk(&returns, confidence_level).map_err(to_py_err)
@@ -154,6 +210,66 @@ fn sampling_seq_bootstrap(
     openquant::sampling::seq_bootstrap(&ind_mat, sample_length, warmup_samples)
 }
 
+#[pyfunction(name = "build_time_bars")]
+fn bars_build_time_bars(
+    timestamps: Vec<String>,
+    prices: Vec<f64>,
+    volumes: Vec<f64>,
+    interval_seconds: i64,
+) -> PyResult<Vec<(String, String, f64, f64, f64, f64, f64, f64, usize)>> {
+    if interval_seconds <= 0 {
+        return Err(PyValueError::new_err("interval_seconds must be > 0"));
+    }
+    let trades = build_trades(timestamps, prices, volumes)?;
+    let bars = time_bars(&trades, chrono::Duration::seconds(interval_seconds));
+    Ok(bars_to_rows(bars))
+}
+
+#[pyfunction(name = "build_tick_bars")]
+fn bars_build_tick_bars(
+    timestamps: Vec<String>,
+    prices: Vec<f64>,
+    volumes: Vec<f64>,
+    ticks_per_bar: usize,
+) -> PyResult<Vec<(String, String, f64, f64, f64, f64, f64, f64, usize)>> {
+    if ticks_per_bar == 0 {
+        return Err(PyValueError::new_err("ticks_per_bar must be > 0"));
+    }
+    let trades = build_trades(timestamps, prices, volumes)?;
+    let bars = standard_bars(&trades, ticks_per_bar as f64, StandardBarType::Tick);
+    Ok(bars_to_rows(bars))
+}
+
+#[pyfunction(name = "build_volume_bars")]
+fn bars_build_volume_bars(
+    timestamps: Vec<String>,
+    prices: Vec<f64>,
+    volumes: Vec<f64>,
+    volume_per_bar: f64,
+) -> PyResult<Vec<(String, String, f64, f64, f64, f64, f64, f64, usize)>> {
+    if !volume_per_bar.is_finite() || volume_per_bar <= 0.0 {
+        return Err(PyValueError::new_err("volume_per_bar must be > 0"));
+    }
+    let trades = build_trades(timestamps, prices, volumes)?;
+    let bars = standard_bars(&trades, volume_per_bar, StandardBarType::Volume);
+    Ok(bars_to_rows(bars))
+}
+
+#[pyfunction(name = "build_dollar_bars")]
+fn bars_build_dollar_bars(
+    timestamps: Vec<String>,
+    prices: Vec<f64>,
+    volumes: Vec<f64>,
+    dollar_value_per_bar: f64,
+) -> PyResult<Vec<(String, String, f64, f64, f64, f64, f64, f64, usize)>> {
+    if !dollar_value_per_bar.is_finite() || dollar_value_per_bar <= 0.0 {
+        return Err(PyValueError::new_err("dollar_value_per_bar must be > 0"));
+    }
+    let trades = build_trades(timestamps, prices, volumes)?;
+    let bars = standard_bars(&trades, dollar_value_per_bar, StandardBarType::Dollar);
+    Ok(bars_to_rows(bars))
+}
+
 #[pyfunction(name = "get_signal")]
 fn bet_sizing_get_signal(prob: Vec<f64>, num_classes: usize, pred: Option<Vec<f64>>) -> Vec<f64> {
     openquant::bet_sizing::get_signal(&prob, num_classes, pred.as_deref())
@@ -325,6 +441,14 @@ fn _core(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_submodule(&sampling)?;
     m.add("sampling", sampling)?;
 
+    let bars = PyModule::new_bound(py, "bars")?;
+    bars.add_function(wrap_pyfunction!(bars_build_time_bars, &bars)?)?;
+    bars.add_function(wrap_pyfunction!(bars_build_tick_bars, &bars)?)?;
+    bars.add_function(wrap_pyfunction!(bars_build_volume_bars, &bars)?)?;
+    bars.add_function(wrap_pyfunction!(bars_build_dollar_bars, &bars)?)?;
+    m.add_submodule(&bars)?;
+    m.add("bars", bars)?;
+
     let bet_sizing = PyModule::new_bound(py, "bet_sizing")?;
     bet_sizing.add_function(wrap_pyfunction!(bet_sizing_get_signal, &bet_sizing)?)?;
     bet_sizing.add_function(wrap_pyfunction!(bet_sizing_discrete_signal, &bet_sizing)?)?;
diff --git a/docs/python_bindings.md b/docs/python_bindings.md
@@ -52,7 +52,7 @@ Input conventions:
 - `timestamps`: list of strings formatted as `%Y-%m-%d %H:%M:%S`
 - timestamp variants require `len(close) == len(timestamps)`
 
-### `openquant.bars` (AFML Ch.2 event-driven bars)
+### `openquant.bars` (AFML Ch.2 event-driven bars; Rust core via PyO3)
 - `build_time_bars(df, interval="1d")`
 - `build_tick_bars(df, ticks_per_bar=50)`
 - `build_volume_bars(df, volume_per_bar=100_000.0)`
diff --git a/python/openquant/bars.py b/python/openquant/bars.py
@@ -1,35 +1,62 @@
 from __future__ import annotations
 
 import math
+from typing import Callable
 
 import polars as pl
 
+from . import _core
 from . import data
 
 
-def _prepare(df: pl.DataFrame) -> pl.DataFrame:
-    return data.clean_ohlcv(df)
-
-
-def _aggregate(df: pl.DataFrame) -> pl.DataFrame:
-    out = (
-        df.group_by(["symbol", "bar_id"])
-        .agg(
-            pl.col("ts").min().alias("start_ts"),
-            pl.col("ts").max().alias("end_ts"),
-            pl.col("open").first().alias("open"),
-            pl.col("high").max().alias("high"),
-            pl.col("low").min().alias("low"),
-            pl.col("close").last().alias("close"),
-            pl.col("adj_close").last().alias("adj_close"),
-            pl.col("volume").sum().alias("volume"),
-            pl.len().alias("n_obs"),
+def _interval_to_seconds(interval: str) -> int:
+    s = interval.strip().lower()
+    if s.endswith("d"):
+        return int(s[:-1]) * 24 * 3600
+    if s.endswith("h"):
+        return int(s[:-1]) * 3600
+    if s.endswith("m"):
+        return int(s[:-1]) * 60
+    if s.endswith("s"):
+        return int(s[:-1])
+    raise ValueError(f"unsupported interval format: {interval}")
+
+
+def _rows_to_frame(symbol: str, rows: list[tuple[str, str, float, float, float, float, float, float, int]]) -> pl.DataFrame:
+    if not rows:
+        return pl.DataFrame(
+            {
+                "ts": [],
+                "symbol": [],
+                "open": [],
+                "high": [],
+                "low": [],
+                "close": [],
+                "volume": [],
+                "adj_close": [],
+                "start_ts": [],
+                "n_obs": [],
+                "dollar_value": [],
+            }
         )
-        .sort(["symbol", "end_ts", "bar_id"])
-        .drop("bar_id")
-        .rename({"end_ts": "ts"})
-    )
-    return out.select(
+    return pl.DataFrame(
+        {
+            "start_ts": [r[0] for r in rows],
+            "ts": [r[1] for r in rows],
+            "open": [r[2] for r in rows],
+            "high": [r[3] for r in rows],
+            "low": [r[4] for r in rows],
+            "close": [r[5] for r in rows],
+            "volume": [r[6] for r in rows],
+            "dollar_value": [r[7] for r in rows],
+            "n_obs": [r[8] for r in rows],
+        }
+    ).with_columns(
+        pl.lit(symbol).alias("symbol"),
+        pl.col("start_ts").str.strptime(pl.Datetime, strict=False),
+        pl.col("ts").str.strptime(pl.Datetime, strict=False),
+        pl.col("close").alias("adj_close"),
+    ).select(
         [
             "ts",
             "symbol",
@@ -41,44 +68,46 @@ def _aggregate(df: pl.DataFrame) -> pl.DataFrame:
             "adj_close",
             "start_ts",
             "n_obs",
+            "dollar_value",
         ]
     )
 
 
+def _build_by_symbol(
+    df: pl.DataFrame,
+    rust_builder: Callable[[list[str], list[float], list[float], float | int], list[tuple[str, str, float, float, float, float, float, float, int]]],
+    param: float | int,
+) -> pl.DataFrame:
+    clean = data.clean_ohlcv(df).sort(["symbol", "ts"])
+    out_frames: list[pl.DataFrame] = []
+    for symbol in clean["symbol"].unique().to_list():
+        sdf = clean.filter(pl.col("symbol") == symbol).sort("ts")
+        rows = rust_builder(
+            [str(x) for x in sdf["ts"].to_list()],
+            [float(x) for x in sdf["close"].to_list()],
+            [float(x) for x in sdf["volume"].to_list()],
+            param,
+        )
+        out_frames.append(_rows_to_frame(symbol, rows))
+    if not out_frames:
+        return _rows_to_frame("", [])
+    return pl.concat(out_frames, how="vertical").sort(["symbol", "ts"])
+
+
 def build_time_bars(df: pl.DataFrame, *, interval: str = "1d") -> pl.DataFrame:
-    clean = _prepare(df)
-    grouped = clean.with_columns(pl.col("ts").dt.truncate(interval).alias("bar_id"))
-    return _aggregate(grouped)
+    return _build_by_symbol(df, _core.bars.build_time_bars, _interval_to_seconds(interval))
 
 
 def build_tick_bars(df: pl.DataFrame, *, ticks_per_bar: int = 50) -> pl.DataFrame:
     if ticks_per_bar <= 0:
         raise ValueError("ticks_per_bar must be > 0")
-    clean = _prepare(df)
-    grouped = clean.with_columns(
-        (pl.int_range(0, pl.len()).over("symbol") // ticks_per_bar)
-        .cast(pl.Int64)
-        .alias("bar_id")
-    )
-    return _aggregate(grouped)
+    return _build_by_symbol(df, _core.bars.build_tick_bars, ticks_per_bar)
 
 
 def build_volume_bars(df: pl.DataFrame, *, volume_per_bar: float = 100_000.0) -> pl.DataFrame:
     if volume_per_bar <= 0:
         raise ValueError("volume_per_bar must be > 0")
-    clean = _prepare(df)
-    eps = volume_per_bar * 1e-9
-    grouped = (
-        clean.with_columns(pl.col("volume").cum_sum().over("symbol").alias("cum_volume"))
-        .with_columns(
-            (((pl.col("cum_volume") - eps).clip(lower_bound=0.0)) / volume_per_bar)
-            .floor()
-            .cast(pl.Int64)
-            .alias("bar_id")
-        )
-        .drop("cum_volume")
-    )
-    return _aggregate(grouped)
+    return _build_by_symbol(df, _core.bars.build_volume_bars, volume_per_bar)
 
 
 def build_dollar_bars(
@@ -88,20 +117,7 @@ def build_dollar_bars(
 ) -> pl.DataFrame:
     if dollar_value_per_bar <= 0:
         raise ValueError("dollar_value_per_bar must be > 0")
-    clean = _prepare(df)
-    eps = dollar_value_per_bar * 1e-9
-    grouped = (
-        clean.with_columns((pl.col("close") * pl.col("volume")).alias("dollar_value"))
-        .with_columns(pl.col("dollar_value").cum_sum().over("symbol").alias("cum_dollar"))
-        .with_columns(
-            (((pl.col("cum_dollar") - eps).clip(lower_bound=0.0)) / dollar_value_per_bar)
-            .floor()
-            .cast(pl.Int64)
-            .alias("bar_id")
-        )
-        .drop(["dollar_value", "cum_dollar"])
-    )
-    return _aggregate(grouped)
+    return _build_by_symbol(df, _core.bars.build_dollar_bars, dollar_value_per_bar)
 
 
 def _lag1_autocorr(values: list[float]) -> float:
@@ -120,7 +136,7 @@ def _lag1_autocorr(values: list[float]) -> float:
 
 
 def bar_diagnostics(df: pl.DataFrame) -> dict[str, float]:
-    clean = _prepare(df).sort(["symbol", "ts"])
+    clean = data.clean_ohlcv(df).sort(["symbol", "ts"])
     returns = (
         clean.with_columns(
             (
