From 1a7fe375a935a4556c1cc5632a8d6b2035880e61 Mon Sep 17 00:00:00 2001 From: Yusuke Tanaka Date: Tue, 9 Jun 2026 16:15:07 +0900 Subject: [PATCH 1/2] fix: resolve DataFrame type through intermediate generic classes get_dataframe_type_from_task auto-detects the DataFrame backend (pandas/polars) from the TaskOnKart[T] type parameter to pick the right file processor. When the type was bound through an intermediate generic class (e.g. Base(TaskOnKart[T]) with Concrete(Base[pl.DataFrame])), the TypeVar was left unresolved and detection fell back to pandas, so polars caches were silently read/written with the pandas processor. Walk the MRO to collect TypeVar bindings and resolve the TaskOnKart type argument so the intermediate-generic case is detected correctly. Downstream projects that worked around this by overriding output() can then drop that workaround. Co-Authored-By: Claude Opus 4.8 (1M context) --- gokart/utils.py | 25 +++++++++++++++++++++++-- test/test_utils.py | 36 +++++++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/gokart/utils.py b/gokart/utils.py index ec251044..148d84b6 100644 --- a/gokart/utils.py +++ b/gokart/utils.py @@ -94,6 +94,16 @@ def load_dill_with_pandas_backward_compatibility(file: FileLike | BytesIO) -> An return pd.read_pickle(cast(Any, file)) +def _resolve_type_var(type_arg: Any, substitutions: dict[TypeVar, Any]) -> Any: + # Follow a chain of TypeVar substitutions (e.g. T_a -> T_b -> pl.DataFrame) to a concrete type. + # seen guards against cyclic substitutions. + seen: set[TypeVar] = set() + while isinstance(type_arg, TypeVar) and type_arg in substitutions and type_arg not in seen: + seen.add(type_arg) + type_arg = substitutions[type_arg] + return type_arg + + def get_dataframe_type_from_task(task: Any) -> Literal['pandas', 'polars', 'polars-lazy']: """ Extract DataFrame type from TaskOnKart[T] type parameter. @@ -121,14 +131,25 @@ def get_dataframe_type_from_task(task: Any) -> Literal['pandas', 'polars', 'pola # Walk the MRO to find TaskOnKart[...] even when defined on a parent class mro = task_class.mro() if hasattr(task_class, 'mro') else [task_class] + # Collect TypeVar bindings so TaskOnKart[T] bound through an intermediate generic class + # (e.g. class Base(TaskOnKart[T]) with class Concrete(Base[pl.DataFrame])) resolves instead of + # falling back to pandas. The MRO runs most-derived first, so bindings are recorded before the + # TaskOnKart base that needs them, letting us collect and resolve in a single pass. + type_var_substitutions: dict[TypeVar, Any] = {} for cls in mro: for base in getattr(cls, '__orig_bases__', ()): origin = get_origin(base) - if origin and hasattr(origin, '__name__') and origin.__name__ == 'TaskOnKart': + if origin is None: + continue + for parameter, arg in zip(getattr(origin, '__parameters__', ()), get_args(base), strict=False): + if isinstance(parameter, TypeVar) and parameter not in type_var_substitutions: + type_var_substitutions[parameter] = arg + + if hasattr(origin, '__name__') and origin.__name__ == 'TaskOnKart': args = get_args(base) if not args: continue - df_type = args[0] + df_type = _resolve_type_var(args[0], type_var_substitutions) module = getattr(df_type, '__module__', '') # Check module name to determine DataFrame type diff --git a/test/test_utils.py b/test/test_utils.py index 2d703214..46e49e47 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,5 +1,5 @@ import unittest -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Generic, TypeVar import pandas as pd import pytest @@ -190,3 +190,37 @@ class _DerivedTaskWithMixin(_BasePolarsTaskWithMixin, _Mixin): task = _DerivedTaskWithMixin() self.assertEqual(get_dataframe_type_from_task(task), 'polars') + + @unittest.skipUnless(HAS_POLARS, 'polars is not installed') + def test_intermediate_generic_class_resolves_polars_typevar(self): + """Detect polars when the type is bound through an intermediate generic class. + + The TypeVar of ``TaskOnKart[T]`` stays unbound on the intermediate class and is + only bound to ``pl.DataFrame`` on a further derived class. Without resolving the + TypeVar substitution while walking the MRO, the unbound ``TaskOnKart[T]`` is picked + up, its module is ``typing``, and detection falls back to pandas. + """ + _T = TypeVar('_T') + + class _GenericTask(TaskOnKart[_T], Generic[_T]): + pass + + class _ConcretePolarsTask(_GenericTask[pl.DataFrame]): + pass + + task = _ConcretePolarsTask() + self.assertEqual(get_dataframe_type_from_task(task), 'polars') + + @unittest.skipUnless(HAS_POLARS, 'polars is not installed') + def test_intermediate_generic_class_resolves_polars_lazyframe_typevar(self): + """Detect polars-lazy when LazyFrame is bound through an intermediate generic class.""" + _T = TypeVar('_T') + + class _GenericTask(TaskOnKart[_T], Generic[_T]): + pass + + class _ConcreteLazyTask(_GenericTask[pl.LazyFrame]): + pass + + task = _ConcreteLazyTask() + self.assertEqual(get_dataframe_type_from_task(task), 'polars-lazy') From 42d79d8949cfbcd73f8638329ca4f9f6bd5cf455 Mon Sep 17 00:00:00 2001 From: Yusuke Tanaka Date: Tue, 9 Jun 2026 19:12:51 +0900 Subject: [PATCH 2/2] test: avoid luigi task registry name collision in DataFrame type tests luigi's Register metaclass globally registers every Task subclass by task_family name at class-definition time. The two new tests both defined a `_GenericTask`, producing an ambiguous registry entry that raised TaskClassAmbigiousException when PandasTypeConfigMap later resolved task names, intermittently failing unrelated worker/tree tests under parallel/random test ordering. Give each test's intermediate generic base a unique name (matching the existing per-test unique-naming convention) and align the polars skip decorator with the rest of the file. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/test_utils.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index 46e49e47..224a48c1 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -191,7 +191,7 @@ class _DerivedTaskWithMixin(_BasePolarsTaskWithMixin, _Mixin): task = _DerivedTaskWithMixin() self.assertEqual(get_dataframe_type_from_task(task), 'polars') - @unittest.skipUnless(HAS_POLARS, 'polars is not installed') + @pytest.mark.skipif(not HAS_POLARS, reason='polars not installed') def test_intermediate_generic_class_resolves_polars_typevar(self): """Detect polars when the type is bound through an intermediate generic class. @@ -202,25 +202,25 @@ def test_intermediate_generic_class_resolves_polars_typevar(self): """ _T = TypeVar('_T') - class _GenericTask(TaskOnKart[_T], Generic[_T]): + class _GenericBasePolarsTask(TaskOnKart[_T], Generic[_T]): pass - class _ConcretePolarsTask(_GenericTask[pl.DataFrame]): + class _DerivedGenericPolarsTask(_GenericBasePolarsTask[pl.DataFrame]): pass - task = _ConcretePolarsTask() + task = _DerivedGenericPolarsTask() self.assertEqual(get_dataframe_type_from_task(task), 'polars') - @unittest.skipUnless(HAS_POLARS, 'polars is not installed') + @pytest.mark.skipif(not HAS_POLARS, reason='polars not installed') def test_intermediate_generic_class_resolves_polars_lazyframe_typevar(self): """Detect polars-lazy when LazyFrame is bound through an intermediate generic class.""" _T = TypeVar('_T') - class _GenericTask(TaskOnKart[_T], Generic[_T]): + class _GenericBaseLazyTask(TaskOnKart[_T], Generic[_T]): pass - class _ConcreteLazyTask(_GenericTask[pl.LazyFrame]): + class _DerivedGenericLazyTask(_GenericBaseLazyTask[pl.LazyFrame]): pass - task = _ConcreteLazyTask() + task = _DerivedGenericLazyTask() self.assertEqual(get_dataframe_type_from_task(task), 'polars-lazy')