Add RoutedObsSource composite observation data source

docs/modules/datasources_dataframe.rst

@@ -34,5 +34,6 @@ Data sources that provide tabular data as DataFrames.
       data.NNJAObsConv
       data.NomadsGDASObsConv
       data.RandomDataFrame
+      data.RoutedObsSource
       data.UFSObsConv
       data.UFSObsSat

earth2studio/data/__init__.py

@@ -61,6 +61,7 @@
     PlanetaryComputerSentinel3AOD,
 )
 from .rand import Random, Random_FX, RandomDataFrame
+from .routed import RoutedObsSource
 from .rx import CosineSolarZenith, LandSeaMask, SurfaceGeoPotential
 from .time_window import TimeWindow
 from .ufs import UFSObsConv, UFSObsSat

earth2studio/data/routed.py

@@ -0,0 +1,156 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Sequence
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+from earth2studio.data.base import DataFrameSource
+from earth2studio.utils.type import TimeArray, VariableArray
+
+
+class RoutedObsSource:
+    """Composite observation data source that routes variables to backends.
+
+    Presents multiple :class:`~earth2studio.data.base.DataFrameSource` backends
+    as a single source: each requested variable is dispatched to the backend
+    that owns it, results are fetched per backend and concatenated into one
+    DataFrame. Useful when a single logical observation stream is served by
+    different archives (e.g. microwave sounders from NNJA and infrared
+    sounders from the UFS Replay).
+
+    All backends must share the columns needed by the consumer. The composite
+    ``SCHEMA`` is the intersection of the backend schemas (fields matching in
+    name and type), and by default only those common columns are returned so
+    the concatenated frame is well-formed.
+
+    Parameters
+    ----------
+    routes : dict[str | Sequence[str], DataFrameSource]
+        Mapping from a variable name (or group of variable names) to the
+        backend data source that serves it. Each variable may appear in only
+        one route.
+
+    Examples
+    --------
+    >>> sat_source = RoutedObsSource(
+    ...     {
+    ...         ("atms", "mhs", "amsua", "amsub"): NNJAObsSat(),
+    ...         ("iasi", "crisfsr", "airs"): UFSObsSat(),
+    ...     }
+    ... )
+    >>> df = sat_source(time, ["atms", "iasi"])  # doctest: +SKIP
+    """
+
+    def __init__(
+        self,
+        routes: dict[str | Sequence[str], DataFrameSource],
+    ) -> None:
+        if not routes:
+            raise ValueError("At least one route must be provided")
+
+        self._routes: dict[str, DataFrameSource] = {}
+        self._sources: list[DataFrameSource] = []
+        for key, source in routes.items():
+            names = [key] if isinstance(key, str) else list(key)
+            if not names:
+                raise ValueError("Route variable groups cannot be empty")
+            for name in names:
+                if name in self._routes:
+                    raise ValueError(
+                        f"Variable {name!r} is mapped to more than one source"
+                    )
+                self._routes[name] = source
+            if all(source is not existing for existing in self._sources):
+                self._sources.append(source)
+
+        self.SCHEMA = self._common_schema([s.SCHEMA for s in self._sources])
+        if len(self.SCHEMA) == 0:
+            raise ValueError("Backend sources share no common schema fields")
+
+    @property
+    def variables(self) -> list[str]:
+        """Variables served by this source, in route order."""
+        return list(self._routes.keys())
+
+    @staticmethod
+    def _common_schema(schemas: list[pa.Schema]) -> pa.Schema:
+        """Intersection of schemas: fields matching in name and type,
+        ordered as in the first schema."""
+        common = [
+            field
+            for field in schemas[0]
+            if all(
+                schema.get_field_index(field.name) >= 0
+                and schema.field(field.name).type == field.type
+                for schema in schemas[1:]
+            )
+        ]
+        return pa.schema(common)
+
+    def __call__(
+        self,
+        time: datetime | list[datetime] | TimeArray,
+        variable: str | list[str] | VariableArray,
+        fields: str | list[str] | pa.Schema | None = None,
+    ) -> pd.DataFrame:
+        """Fetch observations, routing each variable to its backend.
+
+        Parameters
+        ----------
+        time : datetime | list[datetime] | TimeArray
+            Datetime, list of datetimes or array of np.datetime64 to return
+            data for.
+        variable : str | list[str] | VariableArray
+            String, list of strings or array of strings that refer to
+            variables to return.
+        fields : str | list[str] | pa.Schema | None, optional
+            Fields / columns to return; must be available on every backend
+            that serves a requested variable. If None, the common schema
+            columns are returned, by default None
+
+        Returns
+        -------
+        pd.DataFrame
+            Concatenated observations from all involved backends
+        """
+        if isinstance(variable, str):
+            variables = [variable]
+        else:
+            variables = [str(v) for v in np.asarray(variable).ravel()]
+
+        unknown = [v for v in variables if v not in self._routes]
+        if unknown:
+            raise ValueError(
+                f"Variable(s) {unknown} have no route; available: " f"{self.variables}"
+            )
+
+        if fields is None:
+            fields = self.SCHEMA.names
+
+        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)
+        df.attrs["source"] = "earth2studio.data.RoutedObsSource"
+        return df

test/data/test_routed.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from earth2studio.data import RoutedObsSource, fetch_dataframe
+
+
+class PhooObsSource:
+    """Minimal DataFrameSource returning one row per requested variable."""
+
+    def __init__(self, schema: pa.Schema, tag: str):
+        self.SCHEMA = schema
+        self.tag = tag
+        self.calls: list[tuple] = []
+
+    def __call__(self, time, variable, fields=None):
+        self.calls.append((time, list(variable), fields))
+        columns = fields if fields is not None else self.SCHEMA.names
+        rows = [
+            {
+                name: (v if name == "variable" else self.tag)
+                for name in columns
+                if name != "observation"
+            }
+            | {"observation": 1.0}
+            for v in variable
+        ]
+        return pd.DataFrame(rows)[list(columns)]
+
+
+SCHEMA_A = pa.schema(
+    [
+        pa.field("time", pa.timestamp("ns")),
+        pa.field("lat", pa.float32()),
+        pa.field("lon", pa.float32()),
+        pa.field("observation", pa.float32()),
+        pa.field("variable", pa.string()),
+        pa.field("only_in_a", pa.float32()),
+    ]
+)
+SCHEMA_B = pa.schema(
+    [
+        pa.field("time", pa.timestamp("ns")),
+        pa.field("lat", pa.float32()),
+        pa.field("lon", pa.float32()),
+        pa.field("observation", pa.float32()),
+        pa.field("variable", pa.string()),
+        pa.field("only_in_b", pa.string()),
+    ]
+)
+
+TIME = np.array([np.datetime64("2024-01-01T12:00:00")])
+
+
+def _build_routed():
+    src_a = PhooObsSource(SCHEMA_A, "a")
+    src_b = PhooObsSource(SCHEMA_B, "b")
+    routed = RoutedObsSource({("atms", "mhs"): src_a, ("iasi", "airs"): src_b})
+    return routed, src_a, src_b
+
+
+def test_common_schema():
+    routed, _, _ = _build_routed()
+    assert routed.SCHEMA.names == ["time", "lat", "lon", "observation", "variable"]
+    assert routed.variables == ["atms", "mhs", "iasi", "airs"]
+
+
+def test_routing_dispatch():
+    routed, src_a, src_b = _build_routed()
+    df = routed(TIME, ["atms", "iasi", "mhs"])
+
+    # Each backend called once with only its own variables
+    assert len(src_a.calls) == 1
+    assert src_a.calls[0][1] == ["atms", "mhs"]
+    assert len(src_b.calls) == 1
+    assert src_b.calls[0][1] == ["iasi"]
+
+    # Concatenated result covers all requested variables on common columns
+    assert sorted(df["variable"]) == ["atms", "iasi", "mhs"]
+    assert list(df.columns) == routed.SCHEMA.names
+
+
+def test_single_backend_only():
+    routed, src_a, src_b = _build_routed()
+    df = routed(TIME, "atms")
+    assert len(src_a.calls) == 1
+    assert len(src_b.calls) == 0
+    assert list(df["variable"]) == ["atms"]
+
+
+def test_fields_passthrough():
+    routed, src_a, _ = _build_routed()
+    df = routed(TIME, ["atms"], fields=["time", "observation", "variable"])
+    assert src_a.calls[0][2] == ["time", "observation", "variable"]
+    assert list(df.columns) == ["time", "observation", "variable"]
+
+
+def test_unknown_variable_raises():
+    routed, _, _ = _build_routed()
+    with pytest.raises(ValueError, match="no route"):
+        routed(TIME, ["bogus"])
+
+
+def test_duplicate_route_raises():
+    src = PhooObsSource(SCHEMA_A, "a")
+    with pytest.raises(ValueError, match="more than one source"):
+        RoutedObsSource({"atms": src, ("atms", "mhs"): src})
+
+
+def test_empty_routes_raises():
+    with pytest.raises(ValueError, match="At least one route"):
+        RoutedObsSource({})
+
+
+def test_disjoint_schemas_raise():
+    src_a = PhooObsSource(pa.schema([pa.field("x", pa.float32())]), "a")
+    src_b = PhooObsSource(pa.schema([pa.field("y", pa.float32())]), "b")
+    with pytest.raises(ValueError, match="no common schema"):
+        RoutedObsSource({"atms": src_a, "iasi": src_b})
+
+
+def test_fetch_dataframe_integration():
+    routed, _, _ = _build_routed()
+    df = fetch_dataframe(routed, TIME, np.array(["atms", "iasi"]))
+    assert np.all(df.attrs["request_time"] == TIME)
+    assert sorted(df["variable"]) == ["atms", "iasi"]