fix: unify load_dataset X/y convention across dataset families (#530)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Shashankss1205

src/sktime_mcp/runtime/executor.py

@@ -193,13 +193,20 @@ def summarize_available_handles(self, limit: int = 5) -> dict[str, Any]:
             "n_available_handles": len(handle_ids),
         }
 
-    def _resolve_source(self, source: str) -> dict[str, Any]:
-        """Resolve a source id to a series, trying data_handle then demo dataset."""
+    def _resolve_source(self, source: str, prefer: str = "y") -> dict[str, Any]:
+        """Resolve a source id to a series, trying data_handle then demo dataset.
+
+        ``prefer`` selects which component of a demo dataset to return
+        ("y" or "X"); the other is the fallback when the preferred one is
+        absent. Data handles always resolve to their primary series.
+        """
         if source in self._data_handles:
             return {"success": True, "data": self._data_handles[source]["y"]}
         res = self.load_dataset(source)
         if res["success"]:
-            return {"success": True, "data": res["data"]}
+            first, second = ("X", "y") if prefer == "X" else ("y", "X")
+            data = res[first] if res[first] is not None else res[second]
+            return {"success": True, "data": data}
         return res
 
     def instantiate(
@@ -277,7 +284,12 @@ def mock_all_estimators(*args, **kwargs):
 
     # L-7: We can also add custom load_dataset functions here
     def load_dataset(self, name: str) -> dict[str, Any]:
-        """Load a demo dataset."""
+        """Load a demo dataset.
+
+        Returns canonical keys with one consistent meaning for every
+        dataset family: ``y`` is always the target/labels, ``X`` is always
+        the features/panel (or None).
+        """
         demo_datasets = _get_demo_datasets()
         if name not in demo_datasets:
             return {
@@ -294,9 +306,8 @@ def load_dataset(self, name: str) -> dict[str, Any]:
             data = loader()
 
             if isinstance(data, tuple):
-                # sktime classifier/clusterer datasets typically return (X, y)
-                # whereas forecaster datasets typically return (y) or (y, X)
-                # Let's check the shape/type to be safe, or just hardcode known ones
+                # sktime classifier/clusterer datasets return (X-panel, y-labels)
+                # whereas forecaster datasets return (y-target, X-exog)
                 if name in (
                     "arrow_head",
                     "italy_power_demand",
@@ -306,26 +317,21 @@ def load_dataset(self, name: str) -> dict[str, Any]:
                     "plaid",
                 ):
                     X, y = data[0], data[1] if len(data) > 1 else None
-                    # swap them back for our internal representation where 'data' is the primary object requested
-                    return {
-                        "success": True,
-                        "name": name,
-                        "data": X,
-                        "exog": y,
-                        "type": str(type(X).__name__),
-                    }
+                    primary = X
                 else:
                     y, X = data[0], data[1] if len(data) > 1 else None
+                    primary = y
             else:
                 y, X = data, None
+                primary = y
 
             return {
                 "success": True,
                 "name": name,
-                "shape": y.shape if hasattr(y, "shape") else len(y),
-                "type": str(type(y).__name__),
-                "data": y,
-                "exog": X,
+                "shape": primary.shape if hasattr(primary, "shape") else len(primary),
+                "type": str(type(primary).__name__),
+                "y": y,
+                "X": X,
             }
         except Exception as e:
             return {"success": False, "error": str(e)}
@@ -560,19 +566,19 @@ async def predict_async(
                 data_res = self.load_dataset(X_dataset)
                 if not data_res["success"]:
                     raise ValueError(data_res.get("error", "Failed to load dataset"))
-                X = data_res["data"]
-                y = data_res.get("exog")
+                y = data_res["y"]
+                X = data_res["X"]
             else:
                 if X_dataset:
                     data_res = self.load_dataset(X_dataset)
                     if not data_res["success"]:
                         raise ValueError(data_res.get("error", "Failed to load dataset"))
-                    X = data_res["data"]
+                    X = data_res["X"] if data_res["X"] is not None else data_res["y"]
                 if y_dataset:
                     data_res = self.load_dataset(y_dataset)
                     if not data_res["success"]:
                         raise ValueError(data_res.get("error", "Failed to load dataset"))
-                    y = data_res["data"]
+                    y = data_res["y"]
 
             fh = list(range(1, horizon + 1))
 
@@ -656,9 +662,14 @@ def call_method(
                         if "available" in data_res:
                             error_res["available"] = data_res["available"]
                         return error_res
-                    # Replace the kwarg with the actual data (e.g. y_dataset -> y)
+                    # Replace the kwarg with the actual data (e.g. y_dataset -> y);
+                    # the prefix selects the dataset component
                     actual_key = k.replace("_dataset", "")
-                    kwargs[actual_key] = data_res["data"]
+                    if actual_key == "X":
+                        value = data_res["X"] if data_res["X"] is not None else data_res["y"]
+                    else:
+                        value = data_res["y"]
+                    kwargs[actual_key] = value
                     del kwargs[k]
                 elif k.endswith("_data_handle") and isinstance(v, str):
                     if v in self._data_handles:
@@ -802,23 +813,20 @@ async def fit_async(
                 data_res = self.load_dataset(X_dataset)
                 if not data_res["success"]:
                     raise ValueError(data_res["error"])
-                if data_res.get("exog") is not None:
-                    X = data_res["data"]
-                    y = data_res["exog"]
-                else:
-                    y = data_res["data"]
+                y = data_res["y"]
+                X = data_res["X"]
             else:
                 if X_dataset:
                     data_res = self.load_dataset(X_dataset)
                     if not data_res["success"]:
                         raise ValueError(data_res["error"])
-                    X = data_res["data"]
+                    X = data_res["X"] if data_res["X"] is not None else data_res["y"]
 
                 if y_dataset:
                     data_res = self.load_dataset(y_dataset)
                     if not data_res["success"]:
                         raise ValueError(data_res["error"])
-                    y = data_res["data"]
+                    y = data_res["y"]
 
             # Step 2: Fit model
             self._job_manager.update_job(
@@ -901,7 +909,7 @@ async def evaluate_async(
 
             _X = None
             if X:
-                x_res = self._resolve_source(X)
+                x_res = self._resolve_source(X, prefer="X")
                 if not x_res["success"]:
                     raise ValueError(x_res["error"])
                 _X = x_res["data"]

src/sktime_mcp/tools/evaluate.py

@@ -82,7 +82,7 @@ def evaluate_tool(
 
     _X = None
     if X:
-        x_res = executor._resolve_source(X)
+        x_res = executor._resolve_source(X, prefer="X")
         if not x_res["success"]:
             return x_res
         _X = x_res["data"]

src/sktime_mcp/tools/fit_predict.py

@@ -69,24 +69,20 @@ def fit_tool(
         data_res = executor.load_dataset(X_dataset)
         if not data_res["success"]:
             return data_res
-        if data_res.get("exog") is not None:
-            X = data_res["data"]
-            y = data_res["exog"]
-        else:
-            y = data_res["data"]
-            X = None
+        y = data_res["y"]
+        X = data_res["X"]
     else:
         if X_dataset:
             data_res = executor.load_dataset(X_dataset)
             if not data_res["success"]:
                 return data_res
-            X = data_res["data"]
+            X = data_res["X"] if data_res["X"] is not None else data_res["y"]
 
         if y_dataset:
             data_res = executor.load_dataset(y_dataset)
             if not data_res["success"]:
                 return data_res
-            y = data_res["data"]
+            y = data_res["y"]
 
     if run_async:
         import asyncio
@@ -213,20 +209,20 @@ def predict_tool(
         data_res = executor.load_dataset(X_dataset)
         if not data_res["success"]:
             return data_res
-        X = data_res["data"]
-        y = data_res.get("exog")
+        y = data_res["y"]
+        X = data_res["X"]
     else:
         if X_dataset:
             data_res = executor.load_dataset(X_dataset)
             if not data_res["success"]:
                 return data_res
-            X = data_res["data"]
+            X = data_res["X"] if data_res["X"] is not None else data_res["y"]
 
         if y_dataset:
             data_res = executor.load_dataset(y_dataset)
             if not data_res["success"]:
                 return data_res
-            y = data_res["data"]
+            y = data_res["y"]
 
     fh = list(range(1, horizon + 1))
 
@@ -279,20 +275,20 @@ def update_tool(
         data_res = executor.load_dataset(X_dataset)
         if not data_res["success"]:
             return data_res
-        X = data_res["data"]
-        y = data_res.get("exog")
+        y = data_res["y"]
+        X = data_res["X"]
     else:
         if X_dataset:
             data_res = executor.load_dataset(X_dataset)
             if not data_res["success"]:
                 return data_res
-            X = data_res["data"]
+            X = data_res["X"] if data_res["X"] is not None else data_res["y"]
 
         if y_dataset:
             data_res = executor.load_dataset(y_dataset)
             if not data_res["success"]:
                 return data_res
-            y = data_res["data"]
+            y = data_res["y"]
 
     return executor.update(estimator_handle, y=y, X=X)
 

tests/test_core.py

@@ -303,7 +303,11 @@ def test_recent_handles_survive_eviction(self):
         assert f"data_{5:08x}" in ex._data_handles
 
     def test_format_releases_original_handle(self):
-        """format_data_handle releases the source handle after creating the formatted copy."""
+        """format_data_handle releases the source only with release_original=True.
+
+        The explicit transform_data path preserves the caller's input handle;
+        only internal auto-format-on-load passes release_original=True.
+        """
         import pandas as pd
 
         ex = self._make_executor(max_handles=50)
@@ -325,12 +329,13 @@ def test_format_releases_original_handle(self):
             auto_infer_freq=True,
             fill_missing=False,
             remove_duplicates=False,
+            release_original=True,
         )
 
         assert result["success"]
         new_id = result["data_handle"]
         assert new_id != original_id
-        # Original must be gone
+        # Original must be gone (release_original=True is the load-path behavior)
         assert original_id not in ex._data_handles
         # Formatted handle must exist
         assert new_id in ex._data_handles

tests/test_dataset_xy_convention.py

@@ -0,0 +1,88 @@
+"""load_dataset must use one X/y convention for every dataset family.
+
+Previously it returned {"data", "exog"} whose meaning flipped between
+classification datasets (data=X-panel, exog=y-labels) and forecasting
+datasets (data=y-target, exog=X-features). fit/predict/update assumed the
+classification convention, so fit(y_dataset="longley", X_dataset="longley")
+silently fitted with y and X swapped — the exogenous regressors became the
+target.
+"""
+
+import pandas as pd
+import pytest
+
+from sktime_mcp.runtime.executor import get_executor
+from sktime_mcp.tools.fit_predict import fit_tool
+from sktime_mcp.tools.instantiate import instantiate_tool
+
+
+def _release(handle):
+    import contextlib
+
+    from sktime_mcp.runtime.handles import get_handle_manager
+
+    with contextlib.suppress(KeyError):
+        get_handle_manager().release_handle(handle)
+
+
+class TestLoadDatasetCanonicalKeys:
+    def test_forecasting_series_dataset(self):
+        res = get_executor().load_dataset("airline")
+        assert res["success"]
+        assert isinstance(res["y"], pd.Series)
+        assert res["X"] is None
+
+    def test_forecasting_dataset_with_exog(self):
+        res = get_executor().load_dataset("longley")
+        assert res["success"]
+        # y is the univariate target (TOTEMP), X the exogenous regressors
+        assert res["y"].ndim == 1 or res["y"].shape[1] == 1
+        assert isinstance(res["X"], pd.DataFrame)
+        assert res["X"].shape[1] == 5
+
+    def test_classification_dataset(self):
+        res = get_executor().load_dataset("arrow_head")
+        assert res["success"]
+        # X is the panel DataFrame, y the class labels — one label per instance
+        assert isinstance(res["X"], pd.DataFrame)
+        assert res["y"].ndim == 1
+        assert len(res["X"]) == len(res["y"])
+
+
+class TestFitResolution:
+    def test_same_dataset_fit_uses_target_as_y(self):
+        """fit(y_dataset="longley", X_dataset="longley") must fit on TOTEMP."""
+        result = instantiate_tool(spec="NaiveForecaster()")
+        assert result["success"]
+        handle = result["handle"]
+        try:
+            fit_res = fit_tool(
+                estimator_handle=handle,
+                y_dataset="longley",
+                X_dataset="longley",
+            )
+            assert fit_res["success"], fit_res
+            instance = get_executor()._handle_manager.get_instance(handle)
+            fitted_y = instance._y
+            # The buggy path fitted on the 5-column exogenous frame
+            assert fitted_y.ndim == 1 or fitted_y.shape[1] == 1, (
+                f"y was swapped with X: fitted on shape {fitted_y.shape}"
+            )
+        finally:
+            _release(handle)
+
+    def test_same_dataset_fit_classifier_still_works(self):
+        """Classification datasets must keep panel→X, labels→y routing."""
+        result = instantiate_tool(spec="KNeighborsTimeSeriesClassifier()")
+        if not result["success"]:
+            pytest.skip("classifier not available")
+        handle = result["handle"]
+        try:
+            fit_res = fit_tool(
+                estimator_handle=handle,
+                y_dataset="arrow_head",
+                X_dataset="arrow_head",
+            )
+            assert fit_res["success"], fit_res
+        finally:
+            _release(handle)

tests/test_evaluate.py

@@ -68,7 +68,7 @@ def test_evaluate_with_data_handle():
     executor = get_executor()
     data_res = executor.load_dataset("airline")
     assert data_res["success"]
-    executor._data_handles["test_dh"] = {"y": data_res["data"]}
+    executor._data_handles["test_dh"] = {"y": data_res["y"]}
     handle = executor._handle_manager.create_handle("NaiveForecaster", NaiveForecaster(), {})
 
     try:

tests/test_predict_async.py

@@ -26,7 +26,7 @@ def _fitted_handle():
     executor = get_executor()
     handle = executor._handle_manager.create_handle("NaiveForecaster", NaiveForecaster(), {})
     data = executor.load_dataset("airline")
-    fit_res = executor.fit(handle, y=data["data"], fh=list(range(1, 4)))
+    fit_res = executor.fit(handle, y=data["y"], fh=list(range(1, 4)))
     assert fit_res["success"]
     return handle