kumoai 2.10.0.dev202510021830__py3-none-any.whl → 2.12.1__py3-none-any.whl

kumoai/experimental/rfm/local_table.py

@@ -23,11 +23,13 @@ class Column:
         stype: Stype,
         is_primary_key: bool = False,
         is_time_column: bool = False,
+        is_end_time_column: bool = False,
     ) -> None:
         self._name = name
         self._dtype = Dtype(dtype)
         self._is_primary_key = is_primary_key
         self._is_time_column = is_time_column
+        self._is_end_time_column = is_end_time_column
         self.stype = Stype(stype)
 
     @property
@@ -50,9 +52,12 @@ class Column:
             if self._is_primary_key and val != Stype.ID:
                 raise ValueError(f"Primary key '{self.name}' must have 'ID' "
                                  f"semantic type (got '{val}')")
-            if self.name == self._is_time_column and val != Stype.timestamp:
+            if self._is_time_column and val != Stype.timestamp:
                 raise ValueError(f"Time column '{self.name}' must have "
                                  f"'timestamp' semantic type (got '{val}')")
+            if self._is_end_time_column and val != Stype.timestamp:
+                raise ValueError(f"End time column '{self.name}' must have "
+                                 f"'timestamp' semantic type (got '{val}')")
 
         super().__setattr__(key, val)
 
@@ -93,6 +98,7 @@ class LocalTable:
             name="my_table",
             primary_key="id",
             time_column="time",
+            end_time_column=None,
         )
 
         # Verify metadata:
@@ -106,6 +112,8 @@ class LocalTable:
         name: The name of the table.
         primary_key: The name of the primary key of this table, if it exists.
         time_column: The name of the time column of this table, if it exists.
+        end_time_column: The name of the end time column of this table, if it
+            exists.
     """
     def __init__(
         self,
@@ -113,6 +121,7 @@ class LocalTable:
         name: str,
         primary_key: Optional[str] = None,
         time_column: Optional[str] = None,
+        end_time_column: Optional[str] = None,
     ) -> None:
 
         if df.empty:
@@ -130,6 +139,7 @@ class LocalTable:
         self._name = name
         self._primary_key: Optional[str] = None
         self._time_column: Optional[str] = None
+        self._end_time_column: Optional[str] = None
 
         self._columns: Dict[str, Column] = {}
         for column_name in df.columns:
@@ -141,6 +151,9 @@ class LocalTable:
         if time_column is not None:
             self.time_column = time_column
 
+        if end_time_column is not None:
+            self.end_time_column = end_time_column
+
     @property
     def name(self) -> str:
         r"""The name of the table."""
@@ -230,6 +243,8 @@ class LocalTable:
             self.primary_key = None
         if self._time_column == name:
             self.time_column = None
+        if self._end_time_column == name:
+            self.end_time_column = None
         del self._columns[name]
 
         return self
@@ -253,9 +268,8 @@ class LocalTable:
         :class:`ValueError` if the primary key has a non-ID semantic type or
         if the column name does not match a column in the data frame.
         """
-        if not self.has_primary_key():
+        if self._primary_key is None:
             return None
-        assert self._primary_key is not None
         return self[self._primary_key]
 
     @primary_key.setter
@@ -264,6 +278,10 @@ class LocalTable:
             raise ValueError(f"Cannot specify column '{name}' as a primary "
                              f"key since it is already defined to be a time "
                              f"column")
+        if name is not None and name == self._end_time_column:
+            raise ValueError(f"Cannot specify column '{name}' as a primary "
+                             f"key since it is already defined to be an end "
+                             f"time column")
 
         if self.primary_key is not None:
             self.primary_key._is_primary_key = False
@@ -295,9 +313,8 @@ class LocalTable:
         :class:`ValueError` if the time column has a non-timestamp semantic
         type or if the column name does not match a column in the data frame.
         """
-        if not self.has_time_column():
+        if self._time_column is None:
             return None
-        assert self._time_column is not None
         return self[self._time_column]
 
     @time_column.setter
@@ -306,6 +323,10 @@ class LocalTable:
             raise ValueError(f"Cannot specify column '{name}' as a time "
                              f"column since it is already defined to be a "
                              f"primary key")
+        if name is not None and name == self._end_time_column:
+            raise ValueError(f"Cannot specify column '{name}' as a time "
+                             f"column since it is already defined to be an "
+                             f"end time column")
 
         if self.time_column is not None:
             self.time_column._is_time_column = False
@@ -318,6 +339,52 @@ class LocalTable:
         self[name]._is_time_column = True
         self._time_column = name
 
+    # End Time column #########################################################
+
+    def has_end_time_column(self) -> bool:
+        r"""Returns ``True`` if this table has an end time column; ``False``
+        otherwise.
+        """
+        return self._end_time_column is not None
+
+    @property
+    def end_time_column(self) -> Optional[Column]:
+        r"""The end time column of this table.
+
+        The getter returns the end time column of this table, or ``None`` if no
+        such end time column is present.
+
+        The setter sets a column as an end time column on this table, and
+        raises a :class:`ValueError` if the end time column has a non-timestamp
+        semantic type or if the column name does not match a column in the data
+        frame.
+        """
+        if self._end_time_column is None:
+            return None
+        return self[self._end_time_column]
+
+    @end_time_column.setter
+    def end_time_column(self, name: Optional[str]) -> None:
+        if name is not None and name == self._primary_key:
+            raise ValueError(f"Cannot specify column '{name}' as an end time "
+                             f"column since it is already defined to be a "
+                             f"primary key")
+        if name is not None and name == self._time_column:
+            raise ValueError(f"Cannot specify column '{name}' as an end time "
+                             f"column since it is already defined to be a "
+                             f"time column")
+
+        if self.end_time_column is not None:
+            self.end_time_column._is_end_time_column = False
+
+        if name is None:
+            self._end_time_column = None
+            return
+
+        self[name].stype = Stype.timestamp
+        self[name]._is_end_time_column = True
+        self._end_time_column = name
+
     # Metadata ################################################################
 
     @property
@@ -326,16 +393,18 @@ class LocalTable:
         information about the columns in this table.
 
         The returned dataframe has columns ``name``, ``dtype``, ``stype``,
-        ``is_primary_key``, and ``is_time_column``, which provide an aggregate
-        view of the properties of the columns of this table.
+        ``is_primary_key``, ``is_time_column`` and ``is_end_time_column``,
+        which provide an aggregate view of the properties of the columns of
+        this table.
 
         Example:
+            >>> # doctest: +SKIP
             >>> import kumoai.experimental.rfm as rfm
             >>> table = rfm.LocalTable(df=..., name=...).infer_metadata()
             >>> table.metadata
-                name        dtype       stype    is_primary_key is_time_column
-            0   CustomerID  float64     ID       True            False
-        """
+                name        dtype    stype  is_primary_key  is_time_column  is_end_time_column
+            0   CustomerID  float64  ID     True            False           False
+        """  # noqa: E501
         cols = self.columns
 
         return pd.DataFrame({
@@ -355,6 +424,11 @@ class LocalTable:
                 dtype=bool,
                 data=[self._time_column == c.name for c in cols],
             ),
+            'is_end_time_column':
+            pd.Series(
+                dtype=bool,
+                data=[self._end_time_column == c.name for c in cols],
+            ),
         })
 
     def print_metadata(self) -> None:
@@ -417,6 +491,7 @@ class LocalTable:
             candidates = [
                 column.name for column in self.columns
                 if column.stype == Stype.timestamp
+                and column.name != self._end_time_column
             ]
             if time_column := utils.detect_time_column(self._data, candidates):
                 self.time_column = time_column
@@ -430,24 +505,26 @@ class LocalTable:
     # Helpers #################################################################
 
     def _to_api_table_definition(self) -> TableDefinition:
-        cols: List[ColumnDefinition] = []
-        for col in self.columns:
-            cols.append(ColumnDefinition(col.name, col.stype, col.dtype))
-        pkey = self._primary_key
-        time_col = self._time_column
-        source_table = UnavailableSourceTable(table=self.name)
-
         return TableDefinition(
-            cols=cols,
-            source_table=source_table,
-            pkey=pkey,
-            time_col=time_col,
+            cols=[
+                ColumnDefinition(col.name, col.stype, col.dtype)
+                for col in self.columns
+            ],
+            source_table=UnavailableSourceTable(table=self.name),
+            pkey=self._primary_key,
+            time_col=self._time_column,
+            end_time_col=self._end_time_column,
         )
 
     # Python builtins #########################################################
 
     def __hash__(self) -> int:
-        return hash(tuple(self.columns + [self.primary_key, self.time_column]))
+        special_columns = [
+            self.primary_key,
+            self.time_column,
+            self.end_time_column,
+        ]
+        return hash(tuple(self.columns + special_columns))
 
     def __contains__(self, name: str) -> bool:
         return self.has_column(name)
@@ -464,4 +541,5 @@ class LocalTable:
                 f'  num_columns={len(self.columns)},\n'
                 f'  primary_key={self._primary_key},\n'
                 f'  time_column={self._time_column},\n'
+                f'  end_time_column={self._end_time_column},\n'
                 f')')

kumoai/experimental/rfm/pquery/__init__.py

@@ -1,7 +1,7 @@
-from .backend import PQueryBackend
-from .pandas_backend import PQueryPandasBackend
+from .executor import PQueryExecutor
+from .pandas_executor import PQueryPandasExecutor
 
 __all__ = [
-    'PQueryBackend',
-    'PQueryPandasBackend',
+    'PQueryExecutor',
+    'PQueryPandasExecutor',
 ]

kumoai/experimental/rfm/pquery/{backend.py → executor.py}

@@ -1,23 +1,14 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Generic, Optional, Tuple, TypeVar, Union
+from typing import Dict, Generic, Tuple, TypeVar
 
-from kumoapi.rfm import PQueryDefinition
-from kumoapi.rfm.pquery import (
+from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
     Aggregation,
-    AggregationType,
-    BoolOp,
     Column,
     Condition,
     Filter,
-    Float,
-    FloatList,
-    Int,
-    IntList,
+    Join,
     LogicalOperation,
-    MemberOp,
-    RelOp,
-    Str,
-    StrList,
 )
 
 TableData = TypeVar('TableData')
@@ -25,58 +16,33 @@ ColumnData = TypeVar('ColumnData')
 IndexData = TypeVar('IndexData')
 
 
-class PQueryBackend(Generic[TableData, ColumnData, IndexData], ABC):
+class PQueryExecutor(Generic[TableData, ColumnData, IndexData], ABC):
     @abstractmethod
-    def eval_aggregation_type(
+    def execute_column(
         self,
-        op: AggregationType,
-        feat: Optional[ColumnData],
-        batch: IndexData,
-        batch_size: int,
+        column: Column,
+        feat_dict: Dict[str, TableData],
         filter_na: bool = True,
     ) -> Tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
-    def eval_rel_op(
-        self,
-        left: ColumnData,
-        op: RelOp,
-        right: Union[Int, Float, Str, None],
-    ) -> ColumnData:
-        pass
-
-    @abstractmethod
-    def eval_member_op(
-        self,
-        left: ColumnData,
-        op: MemberOp,
-        right: Union[IntList, FloatList, StrList],
-    ) -> ColumnData:
-        pass
-
-    @abstractmethod
-    def eval_bool_op(
-        self,
-        left: ColumnData,
-        op: BoolOp,
-        right: Optional[ColumnData],
-    ) -> ColumnData:
-        pass
-
-    @abstractmethod
-    def eval_column(
+    def execute_aggregation(
         self,
-        column: Column,
+        aggr: Aggregation,
         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]:
         pass
 
     @abstractmethod
-    def eval_aggregation(
+    def execute_condition(
         self,
-        aggr: Aggregation,
+        condition: Condition,
         feat_dict: Dict[str, TableData],
         time_dict: Dict[str, ColumnData],
         batch_dict: Dict[str, IndexData],
@@ -87,9 +53,9 @@ class PQueryBackend(Generic[TableData, ColumnData, IndexData], ABC):
         pass
 
     @abstractmethod
-    def eval_condition(
+    def execute_logical_operation(
         self,
-        condition: Condition,
+        logical_operation: LogicalOperation,
         feat_dict: Dict[str, TableData],
         time_dict: Dict[str, ColumnData],
         batch_dict: Dict[str, IndexData],
@@ -100,9 +66,9 @@ class PQueryBackend(Generic[TableData, ColumnData, IndexData], ABC):
         pass
 
     @abstractmethod
-    def eval_logical_operation(
+    def execute_join(
         self,
-        logical_operation: LogicalOperation,
+        join: Join,
         feat_dict: Dict[str, TableData],
         time_dict: Dict[str, ColumnData],
         batch_dict: Dict[str, IndexData],
@@ -113,20 +79,20 @@ class PQueryBackend(Generic[TableData, ColumnData, IndexData], ABC):
         pass
 
     @abstractmethod
-    def eval_filter(
+    def execute_filter(
         self,
         filter: Filter,
         feat_dict: Dict[str, TableData],
         time_dict: Dict[str, ColumnData],
         batch_dict: Dict[str, IndexData],
         anchor_time: ColumnData,
-    ) -> IndexData:
+    ) -> Tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
-    def eval_pquery(
+    def execute(
         self,
-        query: PQueryDefinition,
+        query: ValidatedPredictiveQuery,
         feat_dict: Dict[str, TableData],
         time_dict: Dict[str, ColumnData],
         batch_dict: Dict[str, IndexData],