feat(datasource): implement get_dataframe_by_tag for drift metric reference data

All drift metric endpoints call data_source.get_dataframe_by_tag() to
retrieve reference data filtered by tag (e.g. "TRAINING"), but the
method did not exist on DataSource, causing AttributeError and HTTP 500
for every drift computation.

Also converts metadata_names to list before calling .index() since
column_names() returns numpy arrays which lack .index(). Test mocks
use numpy arrays to match the real return type.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Sudip Sinha <Sudip.Sinha@RedHat.com>

Author: SudipSinha

src/service/data/datasources/data_source.py

@@ -171,6 +171,58 @@ async def get_organic_dataframe(
 
         return df
 
+    async def get_dataframe_by_tag(self, model_id: str, tag: str) -> pd.DataFrame:
+        """Get a dataframe filtered to rows matching a specific tag.
+
+        Args:
+            model_id: The model ID
+            tag: The tag value to filter by (e.g. "TRAINING")
+
+        Returns:
+            A pandas DataFrame containing only rows whose tags list includes the tag
+
+        """
+        try:
+            model_data = ModelData(model_id)
+            input_data, _, metadata = await model_data.data()
+            input_names, _, metadata_names = await model_data.column_names()
+
+            if metadata is None or input_data is None:
+                return pd.DataFrame()
+
+            tags_col = (
+                list(metadata_names).index("tags") if "tags" in metadata_names else -1
+            )
+            if tags_col < 0:
+                return pd.DataFrame()
+
+            mask = [
+                tag in (row[tags_col] if isinstance(row[tags_col], list) else [])
+                for row in metadata
+            ]
+            filtered_input = input_data[mask]
+
+            df_data: dict[str, object] = {}
+            for i, col_name in enumerate(input_names):
+                if (
+                    len(filtered_input.shape) == ARRAY_DIM_2D
+                    and i < filtered_input.shape[1]
+                ):
+                    df_data[col_name] = filtered_input[:, i]
+                elif len(filtered_input.shape) == 1 and i == 0:
+                    df_data[col_name] = filtered_input
+
+            return pd.DataFrame(df_data)
+
+        except Exception as e:  # Broad catch intentional: dataframe creation involves dynamic storage operations
+            logger.exception(
+                "Error creating dataframe by tag for model=%s, tag=%s",
+                model_id,
+                tag,
+            )
+            msg = f"Error creating dataframe by tag for model={model_id}: {e!s}"
+            raise DataframeCreateError(msg) from e
+
     # METADATA READS
 
     async def get_metadata(self, model_id: str) -> StorageMetadata:

tests/service/data/datasources/test_datasource.py

@@ -428,3 +428,99 @@ async def test_batch_size_calculation_with_limited_data(
         call_args = mock_model_data.data.call_args[1]
         assert call_args["start_row"] == 0  # Should start from beginning
         assert call_args["n_rows"] == EXPECTED_AVAILABLE_ROWS
+
+
+class TestGetDataframeByTag:
+    """Tests for DataSource.get_dataframe_by_tag."""
+
+    @pytest.fixture
+    def data_source(self) -> DataSource:
+        """Create a DataSource instance."""
+        return DataSource()
+
+    @patch("src.service.data.datasources.data_source.ModelData")
+    @pytest.mark.asyncio
+    async def test_returns_only_matching_rows(
+        self, mock_model_data_class: Mock, data_source: DataSource
+    ) -> None:
+        """Rows tagged TRAINING are returned; unlabeled rows are excluded."""
+        mock = Mock(spec=ModelData)
+        mock.column_names = AsyncMock(
+            return_value=(
+                np.array(["feature1", "feature2"]),
+                np.array(["output"]),
+                np.array(["id", "iso_time", "unix_timestamp", "tags"]),
+            ),
+        )
+        mock.data = AsyncMock(
+            return_value=(
+                np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
+                np.array([[0.0], [1.0], [0.0]]),
+                np.array(
+                    [
+                        ["id_0", "t0", 0.0, ["TRAINING"]],
+                        ["id_1", "t1", 1.0, ["unlabeled"]],
+                        ["id_2", "t2", 2.0, ["TRAINING"]],
+                    ],
+                    dtype=object,
+                ),
+            ),
+        )
+        mock_model_data_class.return_value = mock
+
+        df = await data_source.get_dataframe_by_tag("test-model", "TRAINING")
+
+        assert len(df) == 2  # noqa: PLR2004
+        assert list(df.columns) == ["feature1", "feature2"]
+        assert df["feature1"].tolist() == [1.0, 5.0]
+        assert df["feature2"].tolist() == [2.0, 6.0]
+
+    @patch("src.service.data.datasources.data_source.ModelData")
+    @pytest.mark.asyncio
+    async def test_returns_empty_for_nonexistent_tag(
+        self, mock_model_data_class: Mock, data_source: DataSource
+    ) -> None:
+        """Non-existent tag returns empty DataFrame."""
+        mock = Mock(spec=ModelData)
+        mock.column_names = AsyncMock(
+            return_value=(
+                np.array(["feature1"]),
+                np.array(["output"]),
+                np.array(["id", "iso_time", "unix_timestamp", "tags"]),
+            ),
+        )
+        mock.data = AsyncMock(
+            return_value=(
+                np.array([[1.0], [2.0]]),
+                np.array([[0.0], [1.0]]),
+                np.array(
+                    [
+                        ["id_0", "t0", 0.0, ["unlabeled"]],
+                        ["id_1", "t1", 1.0, ["unlabeled"]],
+                    ],
+                    dtype=object,
+                ),
+            ),
+        )
+        mock_model_data_class.return_value = mock
+
+        df = await data_source.get_dataframe_by_tag("test-model", "TRAINING")
+
+        assert len(df) == 0
+
+    @patch("src.service.data.datasources.data_source.ModelData")
+    @pytest.mark.asyncio
+    async def test_returns_empty_when_no_data(
+        self, mock_model_data_class: Mock, data_source: DataSource
+    ) -> None:
+        """Missing data returns empty DataFrame."""
+        mock = Mock(spec=ModelData)
+        mock.column_names = AsyncMock(
+            return_value=(np.array([]), np.array([]), np.array([])),
+        )
+        mock.data = AsyncMock(return_value=(None, None, None))
+        mock_model_data_class.return_value = mock
+
+        df = await data_source.get_dataframe_by_tag("test-model", "TRAINING")
+
+        assert len(df) == 0