DENG-11439-Incremental / streaming profile loads

[gkatre]

Description

Removes the profiler's in-memory ceiling — the root cause of the DENG-11332 OOM — while keeping the partition write atomic.

The column_profiles_v1 profiler accumulated every profiled column row for the whole run in a single in-memory list, then did one final WRITE_TRUNCATE load. Peak memory scaled with the number of tables, and the weekly run was OOM evicted at ~255 tables (DENG-11334 capped per-row size as a stopgap, but didn't remove the ceiling).

This change flushes the buffer incrementally so peak memory is bounded regardless of table count, which unblocks warehouse-scale profiling (the keystone of the DENG-11286 redesign).

How

• Incremental flushing (bounded memory): rows are flushed to BigQuery once the in-memory buffer reaches a byte/row bound, then the buffer is cleared, peak memory no longer grows with table count.
• Atomic partition write via staging: each flush appends into a per-run staging table, and a single query job then replaces the destination profiled_at partition once the whole run succeeds.

Existing semantics are preserved:

• Idempotent re-runs: the atomic partition swap replaces the partition, never doubles it.
• "Leave partition unchanged": when nothing profiles (no rows ⇒ no swap).
• Fail-loud guard: when every table hard-fails (unchanged).

PR Type

• 1st Part of the Profiling Redesign project [Epic DENG-11286]

Related Tickets & Documents

• DENG-11439

Note

_LOAD_BATCH_BYTES/_LOAD_BATCH_ROWS are reasonable default. Tuning it against an explicit pod memory request/limit, if needed, could be handled as a separate follow-up work under the epic (not in this PR).

Checklist

• Reviewed

Reviewer, please follow this checklist

[github-actions]

This PR changes column_profiles_v1/query.py so profiled rows are flushed to BigQuery in 50k-row batches instead of one final WRITE_TRUNCATE load: save_profiles gains a write_disposition parameter, and main() gets a _flush() closure that truncates the date partition on the first flush and appends thereafter.

Two things stood out. First, replacing one load with N loads against the same partition decorator trades atomicity for bounded memory — a failure partway through now leaves a partial partition where previously the prior data stayed intact; staging into a temp table and finishing with a single copy job would get both. Second, the memory bound isn't fully achieved: the futures dict keeps every completed future's result (the per-table profile JSON) alive for the whole run, so peak memory still grows with table count. Details inline, along with a note on the batch-size threshold and the now-stale module docstring.

No tests accompany the change; _flush()'s first-flush-truncates state machine and the "no rows ⇒ leave partition unchanged" path are the kind of logic that's cheap to cover with a fake client, if the repo's conventions for these query.py scripts allow it.

[github-actions]

issue: this drops the atomic partition replacement. Before, the partition was written by exactly one WRITE_TRUNCATE load, so it held either the previous run's complete data or the new run's complete data. Now the first flush truncates and later flushes append, so any failure after flush #1 — a transient BQ/load error, an Airflow task timeout, or the very pod eviction this PR is addressing — leaves column_profiles_v1$YYYYMMDD containing an arbitrary subset of tables that is indistinguishable from a complete run to downstream consumers, until a retry happens to succeed.

Suggested fix: flush with WRITE_APPEND into a per-run staging table (e.g. {destination_table}__staging_{partition_date} with a short expires), then finish main() with a single copy job (WRITE_TRUNCATE) from staging into {destination_table}${partition_date} and delete the staging table. Copy jobs are free, the intermediate state is invisible to consumers, and the memory bound this PR buys is unaffected.

If the partial-partition window is considered acceptable, it should at least be called out in the module docstring, since the PR description currently states that existing semantics are preserved.

[gkatre]

switched to a per-run staging table with a single atomic swap

[github-actions]

issue: peak memory still scales with the number of tables, so the ceiling isn't actually removed. futures holds a strong reference to every Future for the entire lifetime of the with block, and a completed Future retains its result — here (item, profile_json). profile_json is the full per-table profile including the top-K values, i.e. the same payload the buffered rows are built from, so the run still accumulates one profile JSON per table on top of the (now bounded) row buffer. Flushing roughly halves peak memory rather than bounding it.

as_completed() snapshots its input into a set at call time, so it's safe to drain the dict while iterating:

dataset, table_name = futures.pop(future)

That drops the last reference to the Future (and its result) as soon as the row is built.

[gkatre]

Fixed by draining the futures with futures.pop(future) as each completes, so the per-table profile payloads are released instead of held for the whole run; peak memory no longer grows with table count.

[github-actions]

suggestion: 50,000 is a large bound because rows are variable-size, not uniform. A single row can carry _TOP_VALUES_K (50) values each capped at _MAX_VALUE_CHARS (256) chars plus an example value, so a worst-case row is ~13 KB of string data before Python object overhead — 50k of those is several hundred MB of buffer alone. On top of that, load_table_from_json serializes the whole batch to newline-delimited JSON in memory before uploading, so a flush roughly doubles the buffer's footprint at exactly the moment memory is highest.

Given the job was already OOM-evicted, consider either a much smaller row count (5–10k) or accumulating an approximate byte count alongside the rows and flushing on whichever threshold trips first. That keeps the bound tied to the thing that actually OOMs.

[gkatre]

Switched to a byte-based bound as the primary flush trigger, with the row count kept as secondary cap.

Now the bound tracks actual memory footprint rather than a fixed row count.

[github-actions]

nitpick: the module docstring (lines 9-10) is now stale — it still says "A single run accumulates all rows and overwrites its own profiled_at date partition in one load, so the job is idempotent." Worth updating to describe the incremental truncate-then-append flushing, since that sentence is the first thing a reader hits.

[gkatre]

Updated. The module docstring now describes the staging + atomic-swap

[lucia-vargas-a]

A comment about the TRUNCATE then APPEND approach, for your thoughts

[lucia-vargas-a]

Do you see a risk if the first TRUNCATE fails and the table is left with partial data?
How about a safer atomic approach, e.g. write to a temp table and move to production only if the whole job is successful?

[gkatre]

Thanks for the review @lucia-vargas-a !

You have a good point.

Although the first flush's WRITE_TRUNCATE is itself atomic (ie. if it fails, BigQuery leaves the partition unchanged and prior run intact), so that specific case is safe.

But the real risk is a later append (WRITE_APPEND) failing after the truncate succeeded, leaving a partial partition that may look like complete to a downstream process. The AI review also seems to have indicated something similar.

I will make the changes with the staging approach you suggested: append each flush into a temp table, then a single atomic copy into the partition once the whole run succeeds.

Thanks!

[gkatre]

@lucia-vargas-a added the fixes and tested.
Final review when you get a chance?
Thanks!

[scholtzan]

Integration report

• View full diff, if available (SQL + DAGs) (requires private-bigquery-etl access)
• Private CI run
• Bigeye monitoring dry-run: ✅ plan clean — view plan (requires private-bigquery-etl access)

[akkomar]

note (non-blocking): this is compatible with the classification pipeline plan from https://mozilla-hub.atlassian.net/browse/DENG-11453.

todo (non-blocking): the description defers pod request/limit tuning to "separate follow-up work under the epic", but I don't see a ticket for it, and the epic's scope pairs "explicit pod memory limits" with streaming loads.

thought (non-blocking): the piece I couldn't find in the epic is how a full-warehouse run is shaped. That's an epic question rather than a blocker here, but it needs to be addressed before DENG-11442 changes how work is discovered. I'll write up what classification needs into DENG-11445 when it lands.

[akkomar]

issue (non-blocking): the staging name derives only from destination and date, so two invocations sharing a --destination-table on the same day race over it.

Relevant context from the classification side (https://mozilla-hub.atlassian.net/browse/DENG-11453): we invoke this script several times on the same date, once per dataset, and isolate the runs purely by pointing --destination-* at per-run scratch tables. Staging inheriting the destination is what keeps those runs from colliding, so that property is load-bearing for us.

Our invocations are sequential today, so we're not blocked. The race becomes reachable if we enable the concurrency we've currently deferred, or if DENG-11442 shards discovery across parallel tasks writing to one destination. A run id in the staging name would close it now.

[lucia-vargas-a]

The comments pending make sense to me