1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
use std::sync::Arc;
use polars::{
datatypes::{TimeUnit, TimeZone},
prelude::{DataType, Expr, LazyFrame, NULL, Schema, SortMultipleOptions, col, lit, when},
};
use crate::{
data::domain::{SessionWindow, WindowKind},
error::{ChapatyError, ChapatyResult, DataError},
transport::schema::CanonicalCol,
};
pub mod event;
pub mod ohlcv;
pub mod trades;
/// A trait enabling builder-pattern injection of batch indicators.
pub trait WithBatchIndicators: Sized {
type BatchIndicator: Clone;
#[must_use]
fn with_indicator(self, kind: Self::BatchIndicator) -> Self;
#[must_use]
fn with_indicators(self, kinds: &[Self::BatchIndicator]) -> Self {
kinds
.iter()
.fold(self, |acc, kind| acc.with_indicator(kind.clone()))
}
}
pub(crate) trait BatchCompute {
fn pre_compute(&self, lf: LazyFrame) -> ChapatyResult<LazyFrame>;
fn output_schema(&self) -> Arc<Schema>;
}
trait IndicatorExprExt {
/// Computes a continuous, running VWAP.
///
/// **Shape:** Yields a Vector (Series) of the same length as the input.
/// **Usage:** Use this strictly inside projection contexts like
/// `.with_columns()` or `.select()` when you need the running history
/// of the VWAP at every tick.
fn vwap_with_volume(self, volume: Expr) -> Expr;
/// Computes the final, static VWAP for a grouped aggregation.
///
/// **Shape:** Yields a Scalar (`f64`).
/// **Usage:** Use this strictly inside `.group_by().agg()` blocks.
/// Using the standard [`Self::vwap_with_volume`] method in an aggregation
/// context will incorrectly yield a `List<f64>` instead of a singular
/// value.
fn agg_vwap_with_volume(self, volume: Expr) -> Expr;
/// Classifies timestamps into session dates based on a `SessionWindow`.
/// Handles timezone conversions and overnight session wrapping.
fn session_date(self, session: SessionWindow) -> Expr;
/// Detects the precision of an integer epoch timestamp (Seconds, Millis,
/// Micros, or Nanos) based on its magnitude and normalizes it to a UTC
/// Datetime in Microseconds.
///
/// **Thresholds:**
/// * Seconds (10 digits): `< 10^11`
/// * Millis (13 digits): `< 10^14`
/// * Micros (16 digits): `< 10^17`
/// * Nanos (19 digits): `> 10^17`
fn ts_as_microseconds(self) -> Expr;
}
impl IndicatorExprExt for Expr {
fn vwap_with_volume(self, volume: Expr) -> Expr {
let (price_x_volume, vol) = prepare_vwap_components(self, volume);
let cumulative_v = vol.cum_sum(false);
when(cumulative_v.clone().gt(lit(0.0)))
.then(price_x_volume.cum_sum(false) / cumulative_v)
.otherwise(lit(NULL))
}
fn agg_vwap_with_volume(self, volume: Expr) -> Expr {
let (price_x_volume, vol) = prepare_vwap_components(self, volume);
let sum_price_x_vol = price_x_volume.sum();
let sum_vol = vol.sum();
when(sum_vol.clone().gt(lit(0.0)))
.then(sum_price_x_vol / sum_vol)
.otherwise(lit(NULL))
}
fn session_date(self, session: SessionWindow) -> Expr {
let local_ts = self.dt().convert_time_zone(session.pl_time_zone());
let local_time = local_ts.clone().dt().time();
let local_date = local_ts.dt().date();
let start_time = lit(session.start_nanos_since_midnight());
let end_time = lit(session.end_nanos_since_midnight());
match session.window_kind() {
WindowKind::Intraday => when(
local_time
.clone()
.gt(start_time)
.and(local_time.lt_eq(end_time)),
)
.then(local_date)
.otherwise(lit(NULL)),
WindowKind::Overnight => when(local_time.clone().gt(start_time))
.then(local_date.clone())
.when(local_time.lt_eq(end_time))
.then(local_date - lit(chrono::Duration::days(1)))
.otherwise(lit(NULL)),
}
}
fn ts_as_microseconds(self) -> Expr {
let ts = self.cast(DataType::Int64);
when(ts.clone().lt(lit(100_000_000_000_i64))) // < 10^11 -> Seconds
.then(ts.clone() * lit(1_000_000))
.when(ts.clone().lt(lit(100_000_000_000_000_i64))) // < 10^14 -> Millis
.then(ts.clone() * lit(1_000))
.when(ts.clone().gt(lit(100_000_000_000_000_000_i64))) // > 10^17 -> Nanos
.then(ts.clone() / lit(1000)) // Drop nano remainder by integer division
.otherwise(ts) // Already Micros
.cast(DataType::Datetime(
TimeUnit::Microseconds,
Some(TimeZone::UTC),
))
}
}
/// Returns a tuple of `(price_x_volume, valid_volume)` where any rows with
/// zero or negative volume are masked out.
fn prepare_vwap_components(price: Expr, volume: Expr) -> (Expr, Expr) {
let has_volume = volume.clone().gt(lit(0.0));
let price_x_volume = when(has_volume.clone())
.then(price * volume.clone())
.otherwise(lit(0.0));
let valid_vol = when(has_volume).then(volume).otherwise(lit(0.0));
(price_x_volume, valid_vol)
}
trait LazyFrameIndicatorExt {
/// Projects a single-value indicator into the canonical `[Timestamp,
/// Price]` shape.
fn into_price_timeseries(self, value: Expr) -> Self;
}
impl LazyFrameIndicatorExt for LazyFrame {
fn into_price_timeseries(self, value: Expr) -> Self {
self.sort(
[CanonicalCol::PointInTime],
SortMultipleOptions::default().with_maintain_order(false),
)
.select([
col(CanonicalCol::PointInTime),
value.alias(CanonicalCol::Price),
])
.filter(col(CanonicalCol::Price).is_not_null())
}
}
#[expect(
clippy::needless_pass_by_value,
reason = "error-mapping helper consumes the owned error to convert it into a ChapatyError"
)]
fn convert_err(e: polars::error::PolarsError) -> ChapatyError {
ChapatyError::Data(DataError::DataFrame(format!(
"Error while building batch indicator: {e}"
)))
}
#[cfg(test)]
mod tests {
#![expect(
clippy::unwrap_used,
clippy::expect_used,
reason = "tests assert against known-valid fixtures; unwrap and expect surface failures as panics that fail the test"
)]
use chrono::{TimeZone, Utc};
use polars::prelude::{IntoLazy, NamedFrom, Series};
use super::*;
use crate::transport::schema::CanonicalCol;
fn utc_micros(y: i32, mo: u32, d: u32, h: u32, mi: u32) -> i64 {
Utc.with_ymd_and_hms(y, mo, d, h, mi, 0)
.unwrap()
.timestamp_micros()
}
/// Runs `IndicatorExprExt::session_date` over a column of UTC timestamps
/// and returns the resulting session dates as ISO strings (`None` when
/// outside the window).
fn session_dates(window: SessionWindow, ts_utc: &[i64]) -> Vec<Option<String>> {
let series = Series::new(CanonicalCol::PointInTime.as_str().into(), ts_utc.to_vec())
.cast(&DataType::Datetime(
TimeUnit::Microseconds,
Some(polars::prelude::TimeZone::UTC),
))
.expect("cast to datetime");
let df = series
.into_frame()
.lazy()
.select([col(CanonicalCol::PointInTime)
.session_date(window)
.alias("session_date")])
.collect()
.expect("collect");
let as_str = df
.column("session_date")
.unwrap()
.cast(&DataType::String)
.unwrap();
let series = as_str.as_materialized_series();
let ca = series.str().unwrap();
(0..ca.len())
.map(|i| ca.get(i).map(str::to_string))
.collect()
}
fn date(value: &str) -> String {
value.to_string()
}
/// Assumption: We feed OHLCV-1m close timestamps as `PointInTime` Data
///
/// Intraday window (`start_mins < end_mins`): US core session 09:30–16:00
/// ET. `PointInTime` is the bar's close, so the window is `(start, end]`:
/// the bar closing at 09:30 opened at 09:29 (before the session) and is
/// excluded; the bar closing at 09:31 is the first valid bar. June keeps
/// ET at a fixed UTC−4 offset (no DST ambiguity).
#[test]
fn session_date_intraday_us_core_session() {
let window = SessionWindow::us_core_session();
let timestamps = [
utc_micros(2026, 6, 13, 13, 29), // 09:29 ET — bar [09:28,09:29), before open
utc_micros(2026, 6, 13, 13, 30), /* 09:30 ET — bar [09:29,09:30), still before open
* (exclusive) */
utc_micros(2026, 6, 13, 13, 31), // 09:31 ET — bar [09:30,09:31), first in-session bar
utc_micros(2026, 6, 13, 16, 0), // 12:00 ET — mid-session
utc_micros(2026, 6, 13, 19, 59), // 15:59 ET — still inside
utc_micros(2026, 6, 13, 20, 0), // 16:00 ET — bar [15:59,16:00), last bar (inclusive)
utc_micros(2026, 6, 13, 20, 1), // 16:01 ET — after close
utc_micros(2026, 6, 13, 22, 0), // 18:00 ET — after close
];
assert_eq!(
session_dates(window, ×tamps),
vec![
None,
None,
Some(date("2026-06-13")),
Some(date("2026-06-13")),
Some(date("2026-06-13")),
Some(date("2026-06-13")),
None,
None,
]
);
}
/// Assumption: We feed OHLCV-1m close timestamps as `PointInTime` Data
///
/// Overnight window (`start_mins >= end_mins`): US extended overnight
/// 18:00–09:30 ET. `PointInTime` is the bar's close, so the window is
/// `(start, end]`: the bar closing at 18:00 opened at 17:59 (before the
/// session) and is excluded; the bar closing at 18:01 is the first valid
/// bar. The post-midnight morning leg is pulled back to the previous day.
#[test]
fn session_date_overnight_pulls_morning_leg_to_previous_day() {
let window = SessionWindow::us_extended_overnight();
let timestamps = [
utc_micros(2026, 6, 13, 21, 59), // 17:59 ET Jun 13 — bar [17:58,17:59), before open
utc_micros(2026, 6, 13, 22, 0), /* 18:00 ET Jun 13 — bar [17:59,18:00), still
* before open (exclusive) */
utc_micros(2026, 6, 13, 22, 1), /* 18:01 ET Jun 13 — bar [18:00,18:01), first
* in-session bar */
utc_micros(2026, 6, 14, 3, 0), // 23:00 ET Jun 13 — late evening
utc_micros(2026, 6, 14, 6, 0), // 02:00 ET Jun 14 — morning leg → Jun 13
utc_micros(2026, 6, 14, 13, 29), // 09:29 ET Jun 14 — morning leg → Jun 13
utc_micros(2026, 6, 14, 13, 30), /* 09:30 ET Jun 14 — bar [09:29,09:30), last bar
* (inclusive) */
utc_micros(2026, 6, 14, 13, 31), // 09:31 ET Jun 14 — after end
utc_micros(2026, 6, 14, 16, 0), // 12:00 ET Jun 14 — daytime gap
];
assert_eq!(
session_dates(window, ×tamps),
vec![
None,
None,
Some(date("2026-06-13")),
Some(date("2026-06-13")),
Some(date("2026-06-13")),
Some(date("2026-06-13")),
Some(date("2026-06-13")),
None,
None,
]
);
}
}