Skip to content

yggdrasil.arrow.ops

ops

Set-style operations on Arrow tabular data.

Currently exposes :func:upsert_arrow_tabular, which combines two pa.Table / pa.RecordBatch operands keyed by one or more shared columns, and :func:upsert_arrow_batches, the streaming counterpart that accepts two pa.RecordBatch iterables and yields the merged batches without materializing both sides as Tables. The match semantics are governed by :class:Mode:

  • :attr:Mode.APPEND — keep left intact for matching keys; only append rows from right whose keys are absent from left.
  • anything else (e.g. :attr:Mode.UPSERT, :attr:Mode.MERGE, :attr:Mode.OVERWRITE) — drop left rows whose keys are present in right, then append all of right so its values win on conflicts.

The return shape mirrors left (pa.Table in → pa.Table out; pa.RecordBatch in → pa.RecordBatch out) unless row_size / byte_size are set: explicit rechunking always materializes a pa.Table since a pa.RecordBatch is by definition a single chunk.

fill_arrow_table

fill_arrow_table(
    table: Table,
    *,
    sort_by: "str | None" = None,
    partition_by: "Sequence[str] | None" = None,
    fill_strategy: "str | None" = "ffill",
    fill_columns: "Sequence[str] | None" = None
) -> pa.Table

Forward / backward fill nulls per partition.

Parameters

table Input table. The fill is applied in-place on a copy — the input is returned by identity when fill_strategy disables the pass. sort_by Column to sort by within each partition before filling. When given, the table is first sorted by (*partition_by, sort_by) so ffill / bfill runs on the time-ordered axis. When None, the table is assumed to be in the correct order already (the resample path already emits sorted output). partition_by Columns that bound the fill — nulls don't carry across partition boundaries. Each partition's fill is independent. None / empty runs a flat global fill. fill_strategy "ffill" (default) propagates the last non-null value forward into subsequent nulls; "bfill" propagates the next non-null value backward. None / "none" / "" is a no-op (returns the input by identity). fill_columns Restrict the fill to these columns. None runs the fill on every non-partition / non-sort column. Nested types (struct / list / map) are always skipped — pyarrow's fill_null_forward kernel doesn't accept them.

Returns

pa.Table The filled table. The input is returned unchanged when the strategy disables the pass, the table is empty, or no fillable column remains after filtering.

resample_arrow_table

resample_arrow_table(
    table: Table,
    *,
    time_column: str,
    sampling_seconds: int,
    partition_by: "Sequence[str] | None" = None,
    fill_strategy: "str | None" = "ffill"
) -> pa.Table

Align table to a fixed sampling grid on time_column.

Every timestamp is floored to the largest multiple of sampling_seconds that's <= the original — the column ends up on the sampling_seconds grid and rows that landed in the same bucket collapse via "first" aggregation (matches the :func:dedup_arrow_table semantics: pick the first occurrence per group). Pure pyarrow compute, no Python row walk.

partition_by carries the entity columns the resample is independent on — passing ["symbol"] on a multi-instrument price feed groups by (symbol, bucket) so each instrument's rows bucket on their own timeline instead of getting collapsed across instruments. Default None runs a flat resample (one global timeline). The caller / :meth:CastOptions.resample_on_read auto-derives this list from the target schema's :attr:Field.primary_key set, minus time_column itself.

The contract is aggregate (downsample) or identity. When the source already lives on a coarser grid than the target, every bucket has one row and the result equals the input modulo the optional timestamp snap. When the source is finer, the dense rows collapse into the coarser bucket. Expanding (upsample) a coarse source to a finer grid is not in scope — gap-filling is application-specific and best done explicitly.

fill_strategy runs on the resampled output before return — "ffill" (default) carries the last non-null value forward into subsequent nulls within the same partition, "bfill" propagates the next non-null backward, None / "none" skips the pass. Buckets where the first row had a null column inherit from neighbouring buckets on the same partition's timeline; bucket "0" of a partition that has no prior non-null in that column stays null (no cross-partition leak). The fill sorts the resampled output by (*partition_by, time_column) so the result is canonically ordered on return.

Short-circuits in three cases (returns input by identity):

  • time_column missing from the schema,
  • the column isn't a timestamp,
  • sampling_seconds <= 0 (caller's "no resample requested" knob).

resample_arrow_batches

resample_arrow_batches(
    batches: "Iterable[pa.RecordBatch]",
    *,
    time_column: str,
    sampling_seconds: int,
    partition_by: "Sequence[str] | None" = None,
    fill_strategy: "str | None" = "ffill"
) -> "Iterator[pa.RecordBatch]"

Iterator-shaped wrapper around :func:resample_arrow_table.

Materialises the stream into a single :class:pa.Table (a duplicate's bucket-mates can straddle any chunk boundary, just like dedup), runs the resample, and re-batches on pyarrow's natural chunk boundaries.

partition_by carries the entity columns the resample is independent on — see :func:resample_arrow_table for the semantics.

Empty / zero-budget short-circuits to the input iterator unchanged so the caller can route every read through this without paying when there's nothing to resample.

dedup_arrow_table

dedup_arrow_table(table: Table, keys: Sequence[str]) -> pa.Table

Drop duplicate rows on the keys columns, keep the first occurrence.

Implementation runs entirely in pyarrow's C++ kernels:

  1. Append a synthetic __ygg_idx__ column carrying the row index (one pa.array allocation, no row walk).
  2. group_by(keys, use_threads=False).aggregate([(__ygg_idx__, "first")]) collapses the table to one row per key tuple, picking the first occurrence's row index. "first" is an ordered aggregator (pyarrow requires use_threads=False); "min" would coincide on monotonic row indices but benched slower at every table size from 100 to 10 000 rows (multi-threading overhead exceeds its benefit on dedup-shaped work). "first" is also the semantic answer — pick the first row, not the smallest synthetic index.
  3. Sort those indices so the output preserves the input order (group_by makes no ordering promise on its output rows), then Table.take rebuilds the deduped table.

Empty input / empty key list short-circuits to the input unchanged so the caller can call this unconditionally on every read pass.

dedup_arrow_batches

dedup_arrow_batches(
    batches: "Iterable[pa.RecordBatch]", keys: Sequence[str]
) -> "Iterator[pa.RecordBatch]"

Iterator-shaped wrapper around :func:dedup_arrow_table.

An iterator dedup is fundamentally a stop-the-world op: a duplicate's first occurrence can straddle any chunk boundary, so we have to materialise every batch before deciding which rows survive. Pre-materialise into a :class:pa.Table, hand that to :func:dedup_arrow_table, and re-batch the result with pyarrow's natural chunk boundaries on the way out.

Empty key list short-circuits to the input iterator unchanged so the read pipeline can call this unconditionally — the common case (no unique columns in the target schema) stays zero-cost.

upsert_arrow_tabular

upsert_arrow_tabular(
    left: ArrowTabular,
    right: ArrowTabular,
    match_by: "Sequence[str | Field]",
    mode: ModeLike,
    *,
    row_size: int | None = None,
    byte_size: int | None = None,
    memory_pool: MemoryPool | None = None
) -> ArrowTabular

Upsert right into left, matching rows by match_by.

Parameters

left, right pa.Table or pa.RecordBatch. right is projected and cast onto left's schema before concatenation, so any extra columns on right are dropped and shared columns are coerced to left's dtypes. match_by One or more column references — names (str) or :class:yggdrasil.data.Field instances — present on both operands and identifying a row. Field entries contribute their :attr:Field.name; per-frame alias / position fallbacks live on the select_in_* side. Nested key types (struct / list / map / union) are supported via a Python-set fallback; flat keys take a vectorized left-anti join. mode :attr:Mode.APPEND keeps left for matching keys and only appends rows from right whose keys are not in left. Any other mode replaces matching rows in left with right's values and appends the rest of right. Accepts the full :class:ModeLike grammar ("upsert", "append", :class:Mode member, integer code, …). row_size, byte_size Optional output chunking caps. When either is set the result is streamed through :func:yggdrasil.arrow.cast.rechunk_arrow_batches and emitted as a pa.Table whose chunks honour the requested size — overriding the "same kind as left" rule, since a pa.RecordBatch cannot represent multiple chunks. memory_pool Forwarded to the rechunker for pa.concat_batches when coalescing buffered batches under byte_size.

Returns

pa.Table | pa.RecordBatch Same kind as left when neither chunking knob is set; otherwise a pa.Table carrying the rechunked batches.

upsert_arrow_batches

upsert_arrow_batches(
    left: Iterable[RecordBatch],
    right: Iterable[RecordBatch],
    match_by: "Sequence[str | Field] | None",
    mode: ModeLike,
    *,
    schema: Schema | None = None,
    row_size: int | None = None,
    byte_size: int | None = None,
    memory_pool: MemoryPool | None = None
) -> Iterator[pa.RecordBatch]

Streaming upsert over two pa.RecordBatch iterables.

Tuned to avoid concatenating either side into a single Table:

  • :attr:Mode.APPEND — only left's key tuples are buffered (a Python set of hashable rows). left batches stream through untouched while their keys accumulate; right is then streamed batch-by-batch and filtered against the seen keys.
  • upsert-like modes (anything other than APPEND) — right is drained first and held as a list of batches (kept separate, not concatenated) so its keys can drive the left filter; left is streamed batch-by-batch, then the buffered right is emitted so its values win on conflicts.
Parameters

left, right Iterables of pa.RecordBatch (e.g. a generator, a pa.RecordBatchReader, or pa.Table.to_batches()). right's batches are projected onto left's schema before emission, so extra columns are dropped and shared columns are coerced. match_by Column references — names (str) or :class:yggdrasil.data.Field instances — present on both operands and identifying a row. Field entries contribute their :attr:Field.name; alias / position fallbacks live on the select_in_* side and don't bleed into the dedup hash. Nested key types work transparently — keys are materialized to hashable Python values for set membership. None / empty degrades to a plain key-less concatenation (left first, then right) so callers can wire the same dispatch through whether or not they have keys to dedup on; alignment to left's schema is still applied. mode Same grammar as :func:upsert_arrow_tabular. Ignored when match_by is empty — without keys, every mode collapses to "concat left then right". schema Optional output schema override. When omitted, the schema is taken from the first left batch (or the first right batch when left is empty). row_size, byte_size Optional output chunking caps. When set, the result is piped through :func:rechunk_arrow_batches before yielding. memory_pool Forwarded to :func:rechunk_arrow_batches for buffered pa.concat_batches calls.

Yields

pa.RecordBatch The merged stream. Empty (zero-row) batches are dropped.