Skip to content

LazyFrame and GroupBy

A LazyFrame wraps a query plan — a description, not a result. Every method returns a new LazyFrame; nothing reads data until collect(), collect_to_native(), or explain(). At collect, the plan runs through the optimizer and executes outside the interpreter with the GIL released, so other Python threads keep running while a query does.

Transformations

filter(predicate)          # keep rows where the predicate is true
select(exprs)              # replace the column set
with_columns(exprs)        # add or replace columns
group_by(keys) -> GroupBy  # then .agg(exprs) -> LazyFrame
group_by_bar(time_col, every) -> GroupBy  # bucket time_col, then group_by
sort(columns, descending=None)
limit(n)
slice(offset, length)

select, with_columns, group_by, and agg accept a single expression or string, or any iterable of them — strings are column names there. filter accepts a value, so a bare string is a string literal; write col("name") for a column. with_columns evaluates every expression against the original input columns (peers never see each other's outputs) and replaces by output name. sort is stable, takes column names (with descending a single bool or a list matching the columns), and ranks null smallest in every kind. An empty group_by key list aggregates the whole frame.

import keyten as kt

df = kt.DataFrame([
    kt.Series.string("sym", ["a", "b", "a", "b"]),
    kt.Series.int("qty", [1, 2, 3, 4]),
])
out = (
    df.lazy()
    .filter(kt.col("qty") > 1)
    .with_columns((kt.col("qty") * 10).alias("q10"))
    .group_by("sym")
    .agg([kt.col("q10").sum(), kt.col("qty").count().alias("n")])
    .sort("sym")
    .collect()
)
assert out.column("q10").to_list() == [30, 60]
assert out.column("n").to_list() == [1, 2]

Group output order is unspecified. Rows are complete and correct; their order is not part of the contract — the in-memory, parallel, and spilled paths are each deterministic but need not agree. Follow a group_by with an explicit sort when order matters. Null is a groupable value, and Float keys group by bit pattern.

group_by_bar(time_col, every) is truncate(every) on time_col fused with group_by on the result — the bar-building shortcut for market data. every is a duration string ("1m", "5s", "1h", …); see truncate for the unit grammar and per-kind rules. agg() accepts arithmetic combining aggregate results, not just bare aggregates — a ratio like VWAP writes directly as one expression:

import keyten as kt

MIN_NS = 60_000_000_000
ticks = kt.DataFrame([
    kt.Series.time("ts", [t * MIN_NS // 2 for t in range(3)]),
    kt.Series.float("px", [10.0, 10.2, 10.4]),
    kt.Series.float("qty", [5.0, 7.0, 9.0]),
])
vwap = (
    ticks.lazy()
    .group_by_bar("ts", "1m")
    .agg([((kt.col("px") * kt.col("qty")).sum() / kt.col("qty").sum()).alias("vwap")])
    .sort("ts")
    .collect()
)
assert abs(vwap.column("vwap").to_list()[0] - (10.0 * 5.0 + 10.2 * 7.0) / 12.0) < 1e-9

The split form — sum the pieces in agg, then divide in a following with_columns — still works and is sometimes clearer when the intermediate sums are useful on their own:

vwap_split = (
    ticks.lazy()
    .group_by_bar("ts", "1m")
    .agg([
        (kt.col("px") * kt.col("qty")).sum().alias("px_qty"),
        kt.col("qty").sum().alias("vol"),
    ])
    .with_columns((kt.col("px_qty") / kt.col("vol")).alias("vwap"))
    .sort("ts")
    .collect()
)
assert abs(vwap_split.column("vwap").to_list()[0] - vwap.column("vwap").to_list()[0]) < 1e-12

Every col reached inside a combining expression must be underneath an aggregate — (col("px") + col("qty").sum()) names px and raises SchemaError, since px is a row-level column with nothing to aggregate it. Nesting one aggregate inside another (col("qty").sum().sum()) is rejected the same way, with a message calling out the nesting, at plan-build time rather than deep inside execution. A combining expression with no bare-aggregate root (nothing left after removing the arithmetic) still needs an explicit .alias(...), the same rule corr() already follows.

Top-k and dropping nulls

drop_nulls(subset)   # keep rows where every listed column is non-null
head(n=5)             # first n rows (limit with a default)
top_k(k, by)          # k rows with the largest values under by (all descending)
bottom_k(k, by)       # k rows with the smallest values under by (all ascending)

drop_nulls requires at least one column named in subset — there is no whole-frame default the way there is on DataFrame.drop_nulls. top_k and bottom_k take by as a single column name or a list of names; think of them as a fused sort(...).limit(k) that never materializes the full order.

import keyten as kt

df = kt.DataFrame([
    kt.Series.string("sym", ["a", "b", "c", "d", "e"]),
    kt.Series.int("qty", [5, 1, 4, 2, 3]),
])
top = df.lazy().top_k(2, "qty").collect()
assert top.column("qty").to_list() == [5, 4]
assert df.lazy().drop_nulls(["qty"]).head(2).collect().shape == (2, 2)

Joins

inner_join(other, on)   # on = [(left_name, right_name), ...]
left_join(other, on)
semi_join(other, on)    # left rows WITH a match; left columns only
anti_join(other, on)    # left rows WITHOUT a match; left columns only
join_asof(other, on, by=(), *, strategy="backward", tolerance=None)   # on = (left, right); by = [(l, r), ...]

Equi-join semantics, in the order they surprise people: null keys never match (unlike group_by, where null is a key); output is all left columns then the right columns minus its keys; and a colliding right name gets a _right suffix. Join row order is unspecified across resident, parallel, spilled, and distributed strategies. Follow the join with sort when order matters.

join_asof matches each left row to a right row within the same by group, chosen by strategy — always left-outer shaped; a left row with no qualifying right row survives with nulls. Both sides must already be sorted ascending on the on column with no nulls; the engine verifies this (cheaply, from block metadata when it can) and raises rather than sorting your data behind your back.

import keyten as kt

trades = kt.DataFrame([
    kt.Series.string("sym", ["a", "a", "b"]),
    kt.Series.int("ts", [10, 20, 15]),
])
quotes = kt.DataFrame([
    kt.Series.string("sym", ["a", "a", "b"]),
    kt.Series.int("ts", [5, 15, 100]),
    kt.Series.float("bid", [1.0, 2.0, 9.0]),
])
priced = trades.lazy().join_asof(
    quotes.lazy(), ("ts", "ts"), by=[("sym", "sym")]
).sort(["sym", "ts"]).collect()
assert priced.column("bid").to_list() == [1.0, 2.0, None]

strategy picks which right row qualifies: "backward" (the default) takes the latest right row at or before the left row; "forward" takes the earliest right row at or after it; "nearest" takes whichever of the backward/forward candidates is closer in on-value, ties going backward. "nearest" needs a value with magnitude, so it rejects a Str on-key at build time.

quotes2 = kt.DataFrame([
    kt.Series.string("sym", ["a", "a"]),
    kt.Series.int("ts", [100, 200]),
    kt.Series.float("bid", [1.0, 2.0]),
])
trade2 = kt.DataFrame([
    kt.Series.string("sym", ["a"]),
    kt.Series.int("ts", [130]),
])
forward = trade2.lazy().join_asof(quotes2.lazy(), ("ts", "ts"), strategy="forward").collect()
nearest = trade2.lazy().join_asof(quotes2.lazy(), ("ts", "ts"), strategy="nearest").collect()
assert forward.column("bid").to_list() == [2.0]   # earliest quote at or after ts=130
assert nearest.column("bid").to_list() == [1.0]   # ts=100 is 30 away, ts=200 is 70 away

tolerance caps how far the matched on value may be from the left row's before the match is dropped (nulled, not row-dropped — the row still survives, matching the general asof shape). It takes the same representation as on's kind: a duration string ("5s", "1m", …, the truncate grammar) for a Timestamp on-key, or a plain int for an Int on-key; tolerance <= 0 and a mismatched representation both raise ValueError before anything runs.

qts = kt.DataFrame([
    kt.Series.int("ts", [0, 5_000_000_000]),
    kt.Series.float("bid", [1.0, 2.0]),
]).lazy().with_columns(kt.col("ts").cast("timestamp")).collect()
trd = kt.DataFrame([
    kt.Series.int("ts", [2_000_000_000, 11_000_000_000]),
]).lazy().with_columns(kt.col("ts").cast("timestamp")).collect()
stale = trd.lazy().join_asof(qts.lazy(), ("ts", "ts"), tolerance="5s").collect()
assert stale.column("bid").to_list() == [1.0, None]  # second trade is 6s past its backward quote

Combining frames

concat(frames)   # row-wise; identical schemas required

Frames stack in input order; names, kinds, and column order must match exactly.

Running a query

explain() -> str                     # the optimized plan, without running it
serialize() -> bytes                 # the plan's versioned wire form
LazyFrame.deserialize(bytes) -> LazyFrame
collect() -> DataFrame
collect_on(addr) -> DataFrame           # execute the whole plan on one worker
collect_distributed(workers, ctx=None)  # scatter/ship with exact fallback
collect_to_native(path)              # execute straight to a native table

explain shows what will actually run — after predicate pushdown, projection pruning, and the other rewrites — and is the first tool to reach for when a query is slower than expected. collect_to_native persists the result without materializing it as a Python-visible frame; a plan with no aggregation, sort, or join streams block by block in constant memory. serialize/deserialize round-trip a plan through bytes, for sending it somewhere else to run; a plan over an in-memory frame cannot serialize (scan sources carry their path and schema, in-memory data carries neither), and the wire form is versioned rather than a stable file format.

collect_on and collect_distributed require headless workers built from the same engine revision. See remote and distributed execution for supported scan sources, aggregation decomposition, path mapping, fallback behavior, and network security.

import keyten as kt

df = kt.DataFrame([kt.Series.int("v", [3, 1, 2])])
plan = df.lazy().filter(kt.col("v") > 1).select("v").explain()
assert "SCAN" in plan and "predicate" in plan
import keyten as kt

kt.DataFrame([kt.Series.int("v", [1, 2, 3])]).write_native("plan_demo.k10dir")
lf = kt.scan_native("plan_demo.k10dir").filter(kt.col("v") > 1)
wire = lf.serialize()
back = kt.LazyFrame.deserialize(wire).collect()
assert back.to_dict()["v"] == [2, 3]

Watching a query: progress and cancellation

QueryCtx is a running query's observable side: progress counters the engine maintains at block boundaries, and a cancel flag it checks there. Pass one to collect (or collect_distributed) and read it from another thread; cancel() makes the query abort at its next block boundary with CancelledError. Because both are touched only at block boundaries, neither costs the query anything while it runs.

import keyten as kt

df = kt.DataFrame([kt.Series.int("v", list(range(10_000)))])
ctx = kt.QueryCtx()
out = df.lazy().group_by(kt.col("v")).agg([kt.lit(1).count().alias("n")]).collect(ctx)
stage, stages, done, total, rows = ctx.progress()
assert done == total and rows >= 10_000
assert ctx.is_cancelled() is False

In Jupyter, keyten.jupyter.watch(lf) wraps this into a live panel: the optimized plan, a progress bar streaming at ten hertz, and row counts — and the notebook's own interrupt button cancels the query. With watch(lf, workers=[...]) the same panel tracks a distributed collect, the fleet's blocks summed into one bar.

Errors are sticky

A mistake — an unknown column, a kind mismatch — is recorded in the plan the moment it is introduced; later calls pass it along untouched, and collect() or explain() raises the original error. Nothing panics, and nothing partially applies.

import keyten as kt

df = kt.DataFrame([kt.Series.int("v", [1])])
q = df.lazy().filter(kt.col("missing") > 1).select("v")  # no error yet
try:
    q.collect()
except kt.SchemaError:
    pass
else:
    raise AssertionError("expected SchemaError")