kumoai 2.13.0.dev202512031731__cp312-cp312-macosx_11_0_arm64.whl → 2.14.0.dev202512181731__cp312-cp312-macosx_11_0_arm64.whl

kumoai/experimental/rfm/infer/dtype.py

@@ -1,36 +1,23 @@
-from typing import Any, Dict
-
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 from kumoapi.typing import Dtype
 
-PANDAS_TO_DTYPE: Dict[Any, Dtype] = {
-    np.dtype('bool'): Dtype.bool,
-    pd.BooleanDtype(): Dtype.bool,
-    pa.bool_(): Dtype.bool,
-    np.dtype('byte'): Dtype.int,
-    pd.UInt8Dtype(): Dtype.int,
-    np.dtype('int16'): Dtype.int,
-    pd.Int16Dtype(): Dtype.int,
-    np.dtype('int32'): Dtype.int,
-    pd.Int32Dtype(): Dtype.int,
-    np.dtype('int64'): Dtype.int,
-    pd.Int64Dtype(): Dtype.int,
-    np.dtype('float32'): Dtype.float,
-    pd.Float32Dtype(): Dtype.float,
-    np.dtype('float64'): Dtype.float,
-    pd.Float64Dtype(): Dtype.float,
-    np.dtype('object'): Dtype.string,
-    pd.StringDtype(storage='python'): Dtype.string,
-    pd.StringDtype(storage='pyarrow'): Dtype.string,
-    pa.string(): Dtype.string,
-    pa.binary(): Dtype.binary,
-    np.dtype('datetime64[ns]'): Dtype.date,
-    np.dtype('timedelta64[ns]'): Dtype.timedelta,
-    pa.list_(pa.float32()): Dtype.floatlist,
-    pa.list_(pa.int64()): Dtype.intlist,
-    pa.list_(pa.string()): Dtype.stringlist,
+PANDAS_TO_DTYPE: dict[str, Dtype] = {
+    'bool': Dtype.bool,
+    'boolean': Dtype.bool,
+    'int8': Dtype.int,
+    'int16': Dtype.int,
+    'int32': Dtype.int,
+    'int64': Dtype.int,
+    'float16': Dtype.float,
+    'float32': Dtype.float,
+    'float64': Dtype.float,
+    'object': Dtype.string,
+    'string': Dtype.string,
+    'string[python]': Dtype.string,
+    'string[pyarrow]': Dtype.string,
+    'binary': Dtype.binary,
 }
 
 
@@ -45,46 +32,46 @@ def infer_dtype(ser: pd.Series) -> Dtype:
     """
     if pd.api.types.is_datetime64_any_dtype(ser.dtype):
         return Dtype.date
-
+    if pd.api.types.is_timedelta64_dtype(ser.dtype):
+        return Dtype.timedelta
     if isinstance(ser.dtype, pd.CategoricalDtype):
         return Dtype.string
 
-    if pd.api.types.is_object_dtype(ser.dtype):
+    if (pd.api.types.is_object_dtype(ser.dtype)
+            and not isinstance(ser.dtype, pd.ArrowDtype)):
         index = ser.iloc[:1000].first_valid_index()
         if index is not None and pd.api.types.is_list_like(ser[index]):
             pos = ser.index.get_loc(index)
             assert isinstance(pos, int)
             ser = ser.iloc[pos:pos + 1000].dropna()
+            arr = pa.array(ser.tolist())
+            ser = pd.Series(arr, dtype=pd.ArrowDtype(arr.type))
 
-            if not ser.map(pd.api.types.is_list_like).all():
-                raise ValueError("Data contains a mix of list-like and "
-                                 "non-list-like values")
-
-            # Remove all empty Python lists without known data type:
-            ser = ser[ser.map(lambda x: not isinstance(x, list) or len(x) > 0)]
-
-            # Infer unique data types in this series:
-            dtypes = ser.apply(lambda x: PANDAS_TO_DTYPE.get(
-                np.array(x).dtype, Dtype.string)).unique().tolist()
-
-            invalid_dtypes = set(dtypes) - {
-                Dtype.string,
-                Dtype.int,
-                Dtype.float,
-            }
-            if len(invalid_dtypes) > 0:
-                raise ValueError(f"Data contains unsupported list data types: "
-                                 f"{list(invalid_dtypes)}")
-
-            if Dtype.string in dtypes:
-                return Dtype.stringlist
-
-            if dtypes == [Dtype.int]:
+    if isinstance(ser.dtype, pd.ArrowDtype):
+        if pa.types.is_list(ser.dtype.pyarrow_dtype):
+            elem_dtype = ser.dtype.pyarrow_dtype.value_type
+            if pa.types.is_integer(elem_dtype):
                 return Dtype.intlist
-
-            return Dtype.floatlist
-
-    if ser.dtype not in PANDAS_TO_DTYPE:
+            if pa.types.is_floating(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_decimal(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_string(elem_dtype):
+                return Dtype.stringlist
+            if pa.types.is_null(elem_dtype):
+                return Dtype.floatlist
+
+    if isinstance(ser.dtype, np.dtype):
+        dtype_str = str(ser.dtype).lower()
+    elif isinstance(ser.dtype, pd.api.extensions.ExtensionDtype):
+        dtype_str = ser.dtype.name.lower()
+        dtype_str = dtype_str.split('[')[0]  # Remove backend metadata
+    elif isinstance(ser.dtype, pa.DataType):
+        dtype_str = str(ser.dtype).lower()
+    else:
+        dtype_str = 'object'
+
+    if dtype_str not in PANDAS_TO_DTYPE:
         raise ValueError(f"Unsupported data type '{ser.dtype}'")
 
-    return PANDAS_TO_DTYPE[ser.dtype]
+    return PANDAS_TO_DTYPE[dtype_str]

kumoai/experimental/rfm/infer/pkey.py

@@ -1,6 +1,5 @@
 import re
 import warnings
-from typing import Optional
 
 import pandas as pd
 
@@ -9,7 +8,7 @@ def infer_primary_key(
     table_name: str,
     df: pd.DataFrame,
     candidates: list[str],
-) -> Optional[str]:
+) -> str | None:
     r"""Auto-detect potential primary key column.
 
     Args:
@@ -20,6 +19,9 @@ def infer_primary_key(
     Returns:
         The name of the detected primary key, or ``None`` if not found.
     """
+    if len(candidates) == 0:
+        return None
+
     # A list of (potentially modified) table names that are eligible to match
     # with a primary key, i.e.:
     # - UserInfo -> User

kumoai/experimental/rfm/infer/time_col.py

@@ -1,6 +1,5 @@
 import re
 import warnings
-from typing import Optional
 
 import pandas as pd
 
@@ -8,7 +7,7 @@ import pandas as pd
 def infer_time_column(
     df: pd.DataFrame,
     candidates: list[str],
-) -> Optional[str]:
+) -> str | None:
     r"""Auto-detect potential time column.
 
     Args:

kumoai/experimental/rfm/pquery/executor.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Generic, Tuple, TypeVar
+from typing import Generic, TypeVar
 
 from kumoapi.pquery import ValidatedPredictiveQuery
 from kumoapi.pquery.AST import (
@@ -21,82 +21,82 @@ class PQueryExecutor(Generic[TableData, ColumnData, IndexData], ABC):
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, TableData],
+        feat_dict: dict[str, TableData],
         filter_na: bool = True,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass

kumoai/experimental/rfm/pquery/pandas_executor.py

@@ -1,5 +1,3 @@
-from typing import Dict, List, Tuple
-
 import numpy as np
 import pandas as pd
 from kumoapi.pquery import ValidatedPredictiveQuery
@@ -22,9 +20,9 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, pd.DataFrame],
+        feat_dict: dict[str, pd.DataFrame],
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         table_name, column_name = column.fqn.split(".")
         if column_name == '*':
             out = pd.Series(np.ones(len(feat_dict[table_name]), dtype='int64'))
@@ -60,7 +58,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
         batch: np.ndarray,
         batch_size: int,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
 
         mask = feat.notna()
         feat, batch = feat[mask], batch[mask]
@@ -104,13 +102,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         target_table = aggr._get_target_column_name().split('.')[0]
         target_batch = batch_dict[target_table]
         target_time = time_dict[target_table]
@@ -131,10 +129,10 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
                 filter_na=True,
             )
 
-        outs: List[pd.Series] = []
-        masks: List[np.ndarray] = []
+        outs: list[pd.Series] = []
+        masks: list[np.ndarray] = []
         for _ in range(num_forecasts):
-            anchor_target_time = anchor_time[target_batch]
+            anchor_target_time = anchor_time.iloc[target_batch]
             anchor_target_time = anchor_target_time.reset_index(drop=True)
 
             time_filter_mask = (target_time <= anchor_target_time +
@@ -226,13 +224,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -306,13 +304,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -370,13 +368,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(join.rhs_target, Aggregation):
             return self.execute_aggregation(
                 aggr=join.rhs_target,
@@ -393,12 +391,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         out, mask = self.execute_column(
             column=filter.target,
             feat_dict=feat_dict,
@@ -431,12 +429,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(query.entity_ast, Column):
             out, mask = self.execute_column(
                 column=query.entity_ast,