perf(data): vectorize NNJA IR decode and accumulate Arrow tables per batch

[negin513]

Earth2Studio Pull Request

Description

Vectorizes the NNJA IR sounder decode inner loop and switches decode workers to return Arrow tables instead of pickled row dicts. Addresses the performance and memory findings (§5.1–5.2) from the review analysis on #1046, which were explicitly deferred to a follow-on PR.

Vectorized inner loop

_decode_ir_subset previously ran roughly a dozen single-element NumPy operations per channel (two np.array constructions, np.power, a grid lookup with range check, and a full Planck inversion). The function now converts each footprint's channels in one NumPy pass: radiance scaling, wavenumber grid lookup, and radiance_to_bt are each called once per footprint on the full channel vector.

Arrow tables per batch

Decode workers previously returned pickled list[dict] rows that the parent accumulated in full before a single from_pylist → to_pandas call, holding the dict list, Arrow table, and pandas frame in memory simultaneously. Workers now convert rows to an Arrow table immediately after each message (freeing the dicts), return one pa.Table per batch across the process boundary, and the parent concatenates (without consolidating — pd.ArrowDtype columns hold ChunkedArrays natively, so no combine_chunks copy is paid) and converts to pandas once at the end.

Two additional wins from the Arrow return path (measured by @aayushg55 in review):

• Cross-process serialization nearly disappears. 400k rows as list[dict] pickle to 58.6 MB at 0.17 s to dump plus 0.24 s to load; the equivalent pa.Table is 34.4 MB and both directions round to 0.00 s — roughly 1 µs/row the old path paid twice.
• from_pylist moves off the serial critical path. On one 32-channel CrIS aggregate (46.7M rows), converting all rows in the parent is ~50 s of strictly serial wall time; the same work in message-sized calls across eight workers takes 6.56 s and scales essentially linearly.

Full pd.ArrowDtype output + dictionary-encoded strings

_table_to_dataframe previously only mapped unsigned-integer columns to pd.ArrowDtype, leaving floats and strings as standard pandas types. It now passes types_mapper=pd.ArrowDtype for all columns and dictionary-encodes the three low-cardinality string columns (satellite, variable, class) before conversion, cutting per-row string overhead on large frames (~24% smaller DataFrame memory footprint). With every column Arrow-backed, pa.Table.from_pandas on the returned frame is zero-copy, so consumers can do projection/filtering/thinning in Arrow despite being handed a DataFrame.

API-visible dtype change beyond IR: _rows_to_dataframe now delegates to _table_to_dataframe, so decode_microwave output dtypes change too for ATMS, AMSU-A, AMSU-B, and MHS — floats become float[pyarrow]/double[pyarrow], strings become dictionary<int8, string>[pyarrow], times become timestamp[ns][pyarrow]. This makes MW and IR share one dtype contract.

Streaming IR decode — _decode_ir_sounder_chunks

New private generator that yields one pa.Table per completed decode batch. decode_ir_sounder is a thin wrapper over it, so callers processing multiple files can pipeline decode and downstream work without waiting for all batches to finish. Validation and file I/O run eagerly at call time, and the iterator raises _NCEPIRSounderDecodeError on the first batch containing failures, so an early-exiting consumer cannot silently accept a partial decode (preserving the #1046 completeness contract).

Unchanged

Output schema, row order, and all filtering semantics (non-finite observations, missing-CHSF skip counting, non-finite BT, quality packing, CHNM channel ordering). Dtypes are intentionally changed as described above.

Performance

Benchmarked on this machine (single core, warm cache). IASI uses a 616-channel footprint; CrIS uses 431 channels (FSR).

Throughput — _decode_ir_subset

Sensor	Channels	Before	After	Speedup
IASI	616 ch/footprint	116,010 ch-rows/s	909,885 ch-rows/s	7.8×
CrIS	431 ch/footprint	136,135 ch-rows/s	782,826 ch-rows/s	5.8×
AIRS	—	—	—	~1.2×

AIRS gains far less: its radiances are already brightness temperatures (no Planck inversion to vectorize), and its three descriptors per channel make the serial descriptor walk dominate.

Memory — batch accumulation

Measured with pa.total_allocated_bytes and RSS (an earlier draft quoted tracemalloc figures, which trace only the Python allocator and cannot see PyArrow's memory pool, overstating the reduction):

	Before	After
Row representation	list[dict]	pa.Table per batch
Per-row footprint	533 B/row	86 B/row
Reduction		6.2×

DataFrame footprint — _table_to_dataframe (200k rows)

	Before	After
satellite dtype	str (object)¹	dictionary<int8, string>[pyarrow]
lat dtype	float32	float[pyarrow]
wavenumber dtype	float64	double[pyarrow]
DataFrame.memory_usage	26.5 MB	20.1 MB
Reduction		1.3×

¹ On pandas 3 the "before" state is already an Arrow-backed str rather than object.

Note: the dict-encoding cast adds ~6 ms overhead per _table_to_dataframe call on a 200k-row table. For the typical per-file decode path this is negligible relative to BUFR I/O and decode time.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

No new dependencies.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

The PR vectorizes NNJA infrared sounder conversion and replaces cross-process row dictionaries with schema-bound Arrow tables to reduce decode time and memory use.

• Converts each footprint’s channel observations, wavenumbers, and brightness temperatures as aligned NumPy vectors.
• Builds Arrow tables per decoded message and batch, then concatenates them before pandas conversion.
• Adds coverage for table concatenation, empty results, shared dtypes, and the updated worker return contract.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness, schema, ordering, or security regression identified.

The vectorized arrays remain aligned in encoded channel order, preserve the prior filtering conditions, and flow through consistently typed Arrow tables whose concatenation preserves row order and output dtypes.

Important Files Changed

Filename	Overview
earth2studio/data/utils_ncep.py	Vectorizes IR channel decoding and changes worker aggregation from row dictionaries to consistently schema-bound Arrow tables without an identified behavioral regression.
test/data/test_nnja.py	Updates the failure stub for Arrow returns and adds focused tests for concatenation, empty output, and dataframe dtype preservation.

_{Reviews (1): Last reviewed commit: "Vectorize IR subset conversion; accumula..." | Re-trigger Greptile}

[negin513]

/ok to test b96a883

[aayushg55]

Reviewed at 53f6a47 against origin/main at 33eca38, with the decode path exercised directly on both revisions.

Three small things before merge, in the line comments: the combine_chunks() at the end (utils_ncep.py:2189) allocates a second full copy and doubles the peak this PR reduces; the new generator made validation lazy and the failure contract skippable; and the new test bakes in the combine_chunks call rather than testing what production should do.

Three figures in the description need restating, all measurement artifacts rather than problems with the work:

• The "105x memory reduction, 2 B/row" comes from tracemalloc, which traces the Python allocator and so cannot see PyArrow's memory pool — it reports approximately nothing for Arrow buffers. Measured with pa.total_allocated_bytes and RSS, the same rows go from 533 B/row as list[dict] to 86 B/row as a pa.Table: a genuine 6.2x. The "before" figure of 167 B/row also looks like a peak divided by rows across batches that were not simultaneously alive.
• *Throughput— measured AIRS separately at 1.2x: its radiances are already brightness temperatures, so there is no Planck inversion to vectorize, and its three descriptors per channel make the serial descriptor walk dominate.
• The dtype row attributes float64 -> double[pyarrow] to lat, which is declared pa.float32() in the shared schema (lexicon/base.py:410-411) and taken from it directly by NCEP_MICROWAVE_OUTPUT_SCHEMA (utils_ncep.py:1357). Converting through each revision's helper gives lat: float32 -> float[pyarrow]. The described transition is real but belongs to wavenumber, the one genuinely float64 column. The string half is also pandas-version-dependent: on pandas 3 the "before" state is already an Arrow-backed str rather than object.

Two wins here are larger than the ones advertised, and neither is mentioned.

• Row serialization across the process boundary essentially disappears. For 400k rows the list[dict] pickles to 58.6 MB at 0.17 s to dump plus 0.24 s to load; the equivalent pa.Table is 34.4 MB and both directions round to 0.00 s — roughly 1 µs/row that the old path paid twice.
• from_pylist moves off the serial critical path. On one 32-channel CrIS aggregate, converting all 46.7M rows in the parent takes ~50 s of strictly serial wall time; the same work in message-sized calls across eight workers takes 6.56 s (barrier-synchronized so row-building cannot overlap conversion; summed work 47.07 s at 1.01 µs/row, so it scales essentially linearly).
  Also worth a line in the description: _rows_to_dataframe now delegates to _table_to_dataframe, so decode_microwave output changes dtype too — floats to double[pyarrow], strings to dictionary<int8, string>[pyarrow]. That is a good and consistent change, but it is API-visible for ATMS, AMSU-A, AMSU-B and MHS users, which the title does not suggest.

One consequence of the dtype switch deserves more emphasis than the footprint reduction it is credited with: with every column Arrow-backed, pa.Table.from_pandas on the returned frame is zero-copy — identical buffer addresses, all chunks preserved, nothing allocated — so a consumer can do projection, filtering and thinning in Arrow despite being handed a DataFrame. Numpy-backed columns cannot, because to_pandas consolidates each column before the consumer ever sees it.

[aayushg55]

Follow-up ideas (probably future PRs)

Issue A — Wide return for the IR (and MW) sounder decode path

The IR decode emits long format: one row per (footprint, channel). For a 48-hour, seven-sensor,
32-channel request that is 1.87B rows and ~125 GB held as the returned result. Decomposed at the
measured 67 B/row, ~70% of that is footprint geometry restated once per channel and another 12% is a
float64 wavenumber column carrying at most n_channels distinct values. Only ~6% is the radiances.
The same content wide — scalars once per footprint, observation and quality per footprint-channel,
wavenumber and channel index as metadata — is roughly 18 GB, a 7x reduction with nothing dropped. One
6-hour cycle across seven sensors is ~16 GB long against ~2 GB wide.

Wide also enables an ordering that long format forecloses. Quality screening and spatial thinning need
only a few cheap columns to decide which footprints survive, so on a wide table the expensive per-value
work runs on survivors only: on one 32-channel aggregate keeping 5%, gathering before the Planck
inversion costs 21.7 ms against 159.0 ms for inverting first and gathering after, 7.3x. In long format
the values are already rows, so the inversion is paid on everything by construction.

Note run-end encoding is not an alternative route here. It compresses well — a lat column over 431
channels goes 344.8 MB to 1.6 MB, 216x — but on pyarrow 23.0.1 only run_end_encode/run_end_decode
exist. There is no cast, comparison, take or reduction kernel, and pq.write_table rejects the type,
so an encoded column must be decoded before anything can consume it.

Issue B — Share the frame conversion and the footprint-to-long expansion across observation sources

Six sources perform the same operation — per-footprint scalars plus a (footprints, channels) value
matrix, emitted one row per pair — and each hand-rolls it: jpss_cris with np.repeat, metop_iasi and
metop_amsua with np.tile, jpss_atms and the pre-#1059 NCEP path through row dicts. The consequence
is that this PR's wins land on one sensor and the next improvement has to be made six times.

Two extractions, in order of value:

1. _table_to_dataframe should not be private to utils_ncep. It encodes a decision every source
   makes. Moved beside radiance_to_bt in data/utils.py, or into a small data/_frames.py,
   jpss_cris would gain this PR's dtype win by deleting one bare to_pandas() (jpss_cris.py:1138);
   ghcn.py:680 is another. After this PR, the NCEP path is the only Arrow-backed source, so the .values
   contract now depends on which source a user called.
2. A shared footprint-to-long helper returning a pa.Table on the schema, which is the natural home
   for the two things no source does today: emitting wide, and choosing row order deliberately.

The helper is also where emitted row order should be decided, because it currently differs by source
and is documented nowhere: jpss_cris and the NCEP path emit footprint-major
(footprint * n_channels + channel), the MetOp family channel-major (channel * n_obs + footprint).
Same rows, same schema, opposite order. A caller concatenating two sources gets an interleaving that
depends on which decoder produced each part.

Measured, the emitted order should follow the orientation the values already have, since whichever
layout disagrees with the source pays a transpose — the two differ by 1.3-1.4x overall, and the entire
margin is the value matrix. Where a source has no natural orientation, footprint-major is the better
default: it is the only order in which the restated scalars can be compressed, and per-footprint work
reads a contiguous run rather than every n_obs-th row.

This is about the order rows come out in, and is distinct from a source's internal matrix layout, which
issue C covers for IASI — that one changes no output at all.

Issue C — MetOpIASI: float32 arithmetic and matrix orientation

Two independent improvements to the same stage, which dominates the decoder. At one orbit-pass product
(168k footprints, 35.4M rows) at the 211-channel default, the stage is 367.8 ms: 17.7 ms fill, 246.9 ms
radiance_to_bt, 103.2 ms emit.

1. Precision. The spectra arrive as big-endian int16 (metop_iasi.py:668-673), are widened to
   float64 at line 676, carried at double precision through the scaling and the Planck inversion, and
   narrowed to float32 at emit — so the double precision is bought and discarded. In float32
   throughout, the same work takes 108.0 ms against 186.6 ms, 1.7x, for a maximum brightness
   temperature difference of 9.0e-05 K across 73-598 K. For scale, instrument noise is O(0.1 K), and the
   CODATA-vintage difference between constant sets is itself ~2.8e-4 K at 250 K and 700 cm^-1. Keep the
   wavenumber-derived coefficients in float64 and cast each once — they are per-channel scalars, so it
   is free — but note radiance_to_bt is dtype-agnostic, so the wavenumber array must be cast too or
   numpy promotes the expression back.

2. Internal matrix orientation. The radiance matrix is allocated (n_obs, n_ch)
   (metop_iasi.py:558) and then read by column to emit channel-major rows, so it transposes. The
   transpose is not removable, only relocatable: allocating (n_ch, n_obs) makes the emit a ravel but
   makes every scan's write strided instead. Relocating it wins anyway — 281.3 ms against 367.8 ms,
   86.6 ms, 24% of the stage — because transposing a 120-row block at a time stays in cache where
   reading one column of the whole array does not. Emitted row order is identical either way, so
   this changes nothing a caller can observe.

   Worth distinguishing from issue B, which sounds adjacent but is not the same thing: B is about the
   output order being inconsistent across sources, whereas this is purely about which side of IASI's
   own fill-then-emit pair absorbs a transpose that exists regardless. The two interact in one direction
   only — if B standardizes output on footprint-major, IASI's emit changes, and the best internal layout
   should then be re-measured rather than assumed to be this one.

Also dropna(subset=...) followed by a boolean mask (metop_iasi.py:745-747) are two full-frame copies
taken after materialization; filtering the numpy columns before constructing the frame avoids both. The
same pattern is in metop_amsua.py:339 and siblings. AMSU-A, MHS and AVHRR are not worth touching for
the transpose alone — under a millisecond each.

Issue D — scan_angle is NaN for all three NCEP IR sensors

utils_ncep.py:1908 leaves scan_angle as NaN for AIRS, IASI and CrIS, while the class docstring
describes it as the signed nominal look angle. Inherited gap rather than added here. We shoudl correctly fill scan_angle for all MW/IR sounders based on the instrument geometry.

Issue E — wavenumber is float64 schema-wide and restated per row

E2STUDIO_SCHEMA declares wavenumber as pa.float64() (lexicon/base.py:522-527) and every IR
decoder tiles it per footprint. It therefore costs 8 B/row to carry at most n_channels distinct
values — twice what the float32 observation it annotates costs, and ~12% of a long-format payload.

Two possible fixes, with different ceilings. Narrowing to float32 halves it and changes nothing else.
Dictionary encoding — pa.DictionaryArray.from_arrays with int8 indices, since the value set is tiny
and fixed per sensor — takes it to 1 B/row plus a 32-entry dictionary, 8x, and unlike run-end encoding
it is a well-supported type: Parquet uses dictionary encoding natively, pandas maps it to
pd.ArrowDtype(pa.dictionary(...)), and this PR already relies on that for the three string columns.

The caveat is that the saving is in holding the column, not in reading it: a consumer that pulls the
values into numpy gets them decoded back to full width, so this trades memory for a decode at the
boundary. Worth noting dictionary_encode() defaults to int32 indices, which would only be a 2x
saving — the int8 index width has to be chosen explicitly. Per-chunk dictionaries also need
unify_dictionaries() before concatenated chunks can be treated as one column.

Schema-level and affects all sources, so it wants its own change. Note the problem disappears entirely
under the wide return of issue A, where wavenumber is per-channel metadata rather than a column.

[negin513]

Scope note for re-review: 564b5de (trimmed back in e287696) adds one refactor on top of the review fixes — _table_to_dataframe is promoted to data/utils.py as public table_to_dataframe(table, dict_string_columns=()) (first step of follow-up issue B.1). The dictionary-encoded column set is a parameter so the shared helper carries no NCEP-specific policy; utils_ncep keeps a thin delegate with its satellite/variable/class set. No behavior or performance change; the conversion logic is unchanged. Adoption in jpss_atms/jpss_cris/ghcn is deliberately left for a follow-up since it changes those sources' output dtypes.

[negin513]

/ok to test fd339b6

[aayushg55]

LGTM!

[negin513]

/ok to test e78044e

[negin513]

Thanks @aayushg55

[negin513]

/ok to test e78044e