Skip to content

Add RoutedObsSource composite observation data source#974

Open
negin513 wants to merge 2 commits into
NVIDIA:mainfrom
negin513:routed-obs-source
Open

Add RoutedObsSource composite observation data source#974
negin513 wants to merge 2 commits into
NVIDIA:mainfrom
negin513:routed-obs-source

Conversation

@negin513

Copy link
Copy Markdown
Member

Description

Adds RoutedObsSource, a composite DataFrameSource that routes each requested variable to the backend that owns it and concatenates the results on the backends' common schema (intersection of fields by name and type).

Motivation: a single logical observation stream increasingly spans multiple archives — e.g. microwave sounders from NNJA (#968) while infrared sounders remain on the UFS Replay. This lets consumers (such as DA model pipelines) fetch through one source without knowing about the split:

sat_source = RoutedObsSource({
    ("atms", "mhs", "amsua", "amsub"): NNJAObsSat(),
    ("iasi", "crisfsr", "airs"): UFSObsSat(),
})
df = fetch_dataframe(sat_source, time, SAT_VARS)
  • Conforms to the DataFrameSource protocol (works with fetch_dataframe unchanged)
  • Validates routes at construction: duplicate variables, empty routes, and backends with no shared schema fields all raise
  • Unit tested with mock backends (no optional dependencies required)

Checklist

  • Tests added (test/data/test_routed.py)
  • Docs entry (datasources_dataframe.rst)
  • black / ruff / mypy clean

🤖 Generated with Claude Code

negin513 added 2 commits July 16, 2026 20:35
Routes each requested variable to the backend DataFrameSource that owns
it and concatenates the results on the backends' common schema. Lets a
single logical observation stream span multiple archives (e.g. microwave
sounders from NNJA, infrared sounders from UFS Replay) without consumers
knowing about the split.
@greptile-apps

greptile-apps Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds RoutedObsSource, a composite DataFrameSource that dispatches each requested variable to the backend that owns it and concatenates the results over the intersection of the backend schemas. The implementation is well-structured and introduces construction-time validation for duplicate variables, empty routes, and disjoint schemas.

  • RoutedObsSource.__call__ iterates over registered backends, collects per-backend variables, calls each backend once, and concatenates the partial DataFrames. The common SCHEMA (field-name-and-type intersection) is used as the default column projection.
  • The test file covers routing dispatch, single-backend short-circuit, schema intersection, fields passthrough, and all error-path constructors using lightweight mock backends with no optional dependencies.

Confidence Score: 3/5

Safe to merge after addressing the empty-variable crash; the fields validation gap is lower risk but worth fixing before the source is used with arbitrary callers.

The pd.concat([]) crash on an empty variable list is a present defect on the call hot path — any caller that passes an empty variable array will get a confusing traceback instead of a clear error or empty result. The missing validation of caller-supplied fields against the common schema means backends can be called with columns they don't own, producing inconsistent failures.

earth2studio/data/routed.py — specifically the variable-list handling before pd.concat and the fields validation block.

Important Files Changed

Filename Overview
earth2studio/data/routed.py New RoutedObsSource composite; crashes on empty variable lists and silently forwards caller-supplied fields that may not be in the common schema.
test/data/test_routed.py Good coverage of routing, dispatch, schema intersection, and error cases; missing test for empty variable list input.
earth2studio/data/init.py RoutedObsSource correctly exported from the package.
docs/modules/datasources_dataframe.rst RoutedObsSource correctly added to the DataFrame data sources documentation.

Reviews (1): Last reviewed commit: "Add RoutedObsSource to dataframe datasou..." | Re-trigger Greptile

Comment on lines +147 to +154
parts = []
for source in self._sources:
source_variables = [v for v in variables if self._routes[v] is source]
if not source_variables:
continue
parts.append(source(time, source_variables, fields=fields))

df = pd.concat(parts, ignore_index=True)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Empty variable list crashes with unhelpful error

When variable is an empty list or empty array, variables resolves to [], the unknown check passes silently, parts stays empty, and pd.concat([]) raises ValueError: No objects to concatenate. A caller passing np.array([]) through fetch_dataframe would hit this immediately. Adding a guard before the loop — returning an empty DataFrame or raising a clear ValueError — prevents the confusing traceback.

Comment on lines +144 to +145
if fields is None:
fields = self.SCHEMA.names

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Caller-supplied fields not validated against the common schema

When the caller passes explicit fields, those column names are forwarded verbatim to every involved backend with no check against self.SCHEMA. A field present in only one backend will cause the other to raise or silently drop it, producing inconsistent errors that are hard to attribute. Validating against the intersection upfront produces a clear, actionable message.

Suggested change
if fields is None:
fields = self.SCHEMA.names
if fields is None:
fields = self.SCHEMA.names
else:
# Normalise to a list for uniform handling below
if isinstance(fields, pa.Schema):
fields = fields.names
elif isinstance(fields, str):
fields = [fields]
unknown_fields = [f for f in fields if f not in self.SCHEMA.names]
if unknown_fields:
raise ValueError(
f"Field(s) {unknown_fields} are not in the common schema; "
f"available: {self.SCHEMA.names}"
)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant