kumoai 2.13.0.dev202512061731__cp313-cp313-macosx_11_0_arm64.whl → 2.13.0.dev202512081731__cp313-cp313-macosx_11_0_arm64.whl

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.13.0.dev202512061731'
+__version__ = '2.13.0.dev202512081731'

kumoai/experimental/rfm/backend/local/sampler.py

@@ -4,7 +4,7 @@ import numpy as np
 import pandas as pd
 
 from kumoai.experimental.rfm.backend.local import LocalGraphStore
-from kumoai.experimental.rfm.base import EdgeSpec, Sampler, SamplerOutput
+from kumoai.experimental.rfm.base import BackwardSamplerOutput, Sampler
 from kumoai.utils import ProgressLogger
 
 if TYPE_CHECKING:
@@ -36,30 +36,19 @@ class LocalSampler(Sampler):
             self._graph_store.time_dict,
         )
 
-    def sample(
+    def _sample_backward(
         self,
         entity_table_name: str,
         entity_pkey: pd.Series,
         anchor_time: pd.Series,
-        column_spec_dict: dict[str, list[str]],
-        edge_spec_dict: dict[tuple[str, str, str], list[EdgeSpec]],
-        drop_duplicates: bool = False,
-        return_edges: bool = False,
-    ) -> SamplerOutput:
+        columns_dict: dict[str, set[str]],
+        num_neighbors: list[int],
+    ) -> BackwardSamplerOutput:
 
-        if anchor_time.dtype != 'datetime64[ns]':
-            anchor_time = anchor_time.astype('datetime64[ns]')
-
-        num_hops = max([len(specs) for specs in edge_spec_dict.values()] + [0])
-        num_neighbors_dict: dict[str, list[int]] = {}
-
-        for edge_type, specs in edge_spec_dict.items():
-            edge_type_str = '__'.join(edge_type)
-            num_neighbors_dict[edge_type_str] = [0] * num_hops
-            for hop, spec in enumerate(specs):
-                # TODO Add support for time-based sampling.
-                assert spec.num_neighbors is not None
-                num_neighbors_dict[edge_type_str][hop] = spec.num_neighbors
+        num_neighbors_dict: dict[str, list[int]] = {
+            '__'.join(edge_type): num_neighbors
+            for edge_type in self.edge_types
+        }
 
         (
             row_dict,
@@ -80,10 +69,12 @@ class LocalSampler(Sampler):
         inverse_dict: dict[str, np.ndarray] = {}
         for table_name, node in node_dict.items():
             df = self._graph_store.df_dict[table_name]
-            columns = column_spec_dict[table_name]
+            columns = columns_dict[table_name]
             if self.end_time_column_dict.get(table_name, None) in columns:
                 df = df.iloc[node]
-            elif len(columns) > 0 and drop_duplicates:
+            elif len(columns) == 0:
+                df = df.iloc[node]
+            else:
                 # Only store unique rows in `df` above a certain threshold:
                 unique_node, inverse = np.unique(node, return_inverse=True)
                 if len(node) > 1.05 * len(unique_node):
@@ -91,10 +82,8 @@ class LocalSampler(Sampler):
                     inverse_dict[table_name] = inverse
                 else:
                     df = df.iloc[node]
-            else:
-                df = df.iloc[node]
             df = df.reset_index(drop=True)
-            df = df[columns]
+            df = df[list(columns)]
             df_dict[table_name] = df
 
         num_sampled_nodes_dict = {
@@ -103,24 +92,20 @@ class LocalSampler(Sampler):
             num_sampled_nodes_dict.items()
         }
 
-        if return_edges:
-            row_dict = {
-                edge_type: row_dict['__'.join(edge_type)]
-                for edge_type in edge_spec_dict.keys()
-            }
-            col_dict = {
-                edge_type: col_dict['__'.join(edge_type)]
-                for edge_type in edge_spec_dict.keys()
-            }
-            num_sampled_edges_dict = {
-                edge_type:
-                num_sampled_edges_dict['__'.join(edge_type)].tolist()
-                for edge_type in edge_spec_dict.keys()
-            }
-        else:
-            row_dict = col_dict = num_sampled_edges_dict = None
+        row_dict = {
+            edge_type: row_dict['__'.join(edge_type)]
+            for edge_type in self.edge_types
+        }
+        col_dict = {
+            edge_type: col_dict['__'.join(edge_type)]
+            for edge_type in self.edge_types
+        }
+        num_sampled_edges_dict = {
+            edge_type: num_sampled_edges_dict['__'.join(edge_type)].tolist()
+            for edge_type in self.edge_types
+        }
 
-        return SamplerOutput(
+        return BackwardSamplerOutput(
             df_dict=df_dict,
             inverse_dict=inverse_dict,
             batch_dict=batch_dict,

kumoai/experimental/rfm/base/__init__.py

@@ -1,14 +1,14 @@
 from .source import SourceColumn, SourceForeignKey
 from .column import Column
 from .table import Table
-from .sampler import EdgeSpec, SamplerOutput, Sampler
+from .sampler import BackwardSamplerOutput, ForwardSamplerOutput, Sampler
 
 __all__ = [
     'SourceColumn',
     'SourceForeignKey',
     'Column',
     'Table',
-    'EdgeSpec',
-    'SamplerOutput',
+    'BackwardSamplerOutput',
+    'ForwardSamplerOutput',
     'Sampler',
 ]

kumoai/experimental/rfm/base/sampler.py

@@ -1,11 +1,14 @@
 import copy
 import re
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import pandas as pd
+from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery.AST import Aggregation, ASTNode
 from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
 from kumoapi.typing import Stype
 
@@ -14,28 +17,21 @@ if TYPE_CHECKING:
 
 
 @dataclass
-class EdgeSpec:
-    num_neighbors: int | None = None
-    time_offsets: tuple[
-        pd.DateOffset | None,
-        pd.DateOffset,
-    ] | None = None
-
-    def __post_init__(self) -> None:
-        if (self.num_neighbors is None) == (self.time_offsets is None):
-            raise ValueError("Only one of 'num_neighbors' and 'time_offsets' "
-                             "must be provided")
-
-
-@dataclass
-class SamplerOutput:
+class BackwardSamplerOutput:
     df_dict: dict[str, pd.DataFrame]
     inverse_dict: dict[str, np.ndarray]
     batch_dict: dict[str, np.ndarray]
     num_sampled_nodes_dict: dict[str, list[int]]
-    row_dict: dict[tuple[str, str, str], np.ndarray] | None = None
-    col_dict: dict[tuple[str, str, str], np.ndarray] | None = None
-    num_sampled_edges_dict: dict[tuple[str, str, str], list[int]] | None = None
+    row_dict: dict[tuple[str, str, str], np.ndarray]
+    col_dict: dict[tuple[str, str, str], np.ndarray]
+    num_sampled_edges_dict: dict[tuple[str, str, str], list[int]]
+
+
+@dataclass
+class ForwardSamplerOutput:
+    entity_pkey: pd.Series
+    anchor_time: pd.Series
+    target: pd.Series
 
 
 class Sampler(ABC):
@@ -104,11 +100,6 @@ class Sampler(ABC):
         exclude_cols_dict: dict[str, list[str]] | None = None,
     ) -> Subgraph:
 
-        edge_spec_dict: dict[tuple[str, str, str], list[EdgeSpec]] = {
-            edge_type: [EdgeSpec(value) for value in num_neighbors]
-            for edge_type in self.edge_types
-        }
-
         # Exclude all columns that leak target information:
         table_stype_dict: dict[str, dict[str, Stype]] = self._table_stype_dict
         if exclude_cols_dict is not None:
@@ -118,26 +109,23 @@ class Sampler(ABC):
                     del table_stype_dict[table_name][column_name]
 
         # Collect all columns being used as features:
-        column_spec_dict: dict[str, list[str]] = {
-            table_name: list(stype_dict.keys())
+        columns_dict: dict[str, set[str]] = {
+            table_name: set(stype_dict.keys())
             for table_name, stype_dict in table_stype_dict.items()
         }
         # Make sure to store primary key information for entity tables:
         for table_name in entity_table_names:
-            column_spec_dict[table_name] = (
-                [self.primary_key_dict[table_name]] +
-                column_spec_dict[table_name])
+            columns_dict[table_name].add(self.primary_key_dict[table_name])
 
         if anchor_time.dtype != 'datetime64[ns]':
             anchor_time = anchor_time.astype('datetime64[ns]')
-        out = self.sample(
+
+        out = self._sample_backward(
             entity_table_name=entity_table_names[0],
             entity_pkey=entity_pkey,
             anchor_time=anchor_time,
-            column_spec_dict=column_spec_dict,
-            edge_spec_dict=edge_spec_dict,
-            drop_duplicates=True,
-            return_edges=True,
+            columns_dict=columns_dict,
+            num_neighbors=num_neighbors,
         )
 
         subgraph = Subgraph(
@@ -150,14 +138,14 @@ class Sampler(ABC):
             if len(batch) == 0:
                 continue
 
-            primary_key = None
+            primary_key: str | None = None
             if table_name in entity_table_names:
-                primary_key = self.primary_key_dict.get(table_name, None)
+                primary_key = self.primary_key_dict[table_name]
 
             df = out.df_dict[table_name].reset_index(drop=True)
-            if table_name in self.end_time_column_dict:
+            if end_time_column := self.end_time_column_dict.get(table_name):
                 # Set end time to NaT for all values greater than anchor time:
-                end_time_column = self.end_time_column_dict[table_name]
+                assert table_name not in out.inverse_dict
                 ser = df[end_time_column]
                 if ser.dtype != 'datetime64[ns]':
                     ser = ser.astype('datetime64[ns]')
@@ -179,9 +167,6 @@ class Sampler(ABC):
                 primary_key=primary_key,
             )
 
-        assert out.row_dict is not None
-        assert out.col_dict is not None
-        assert out.num_sampled_edges_dict is not None
         for edge_type in out.row_dict.keys():
             row: np.ndarray | None = out.row_dict[edge_type]
             col: np.ndarray | None = out.col_dict[edge_type]
@@ -227,19 +212,101 @@ class Sampler(ABC):
 
         return subgraph
 
+    def sample_forward(
+        self,
+        query: ValidatedPredictiveQuery,
+        num_examples: int,
+        anchor_time: pd.Timestamp | Literal['entity'],
+        random_seed: int | None = None,
+    ) -> ForwardSamplerOutput:
+
+        columns_dict: dict[str, set[str]] = defaultdict(set)
+        for fqn in query.all_query_columns + [query.entity_column]:
+            table_name, column_name = fqn.split('.')
+            columns_dict[table_name].add(column_name)
+
+        if time_column := self.time_column_dict[query.entity_table]:
+            columns_dict[table_name].add(time_column)
+        if end_time_column := self.end_time_column_dict[query.entity_table]:
+            columns_dict[table_name].add(end_time_column)
+
+        time_offset_dict: dict[
+            tuple[str, str, str],
+            tuple[pd.DateOffset | None, pd.DateOffset],
+        ] = {}
+
+        def _add_time_offset(node: ASTNode, num_forecasts: int = 1) -> None:
+            if isinstance(node, Aggregation):
+                table_name = node._get_target_column_name().split('.')[0]
+                columns_dict[table_name].add(self.time_column_dict[table_name])
+
+                edge_types = [
+                    edge_type for edge_type in self.edge_types
+                    if edge_type[0] == table_name
+                    and edge_type[2] == query.entity_table
+                ]
+                if len(edge_types) != 1:
+                    raise ValueError(f"Could not find a unique foreign key "
+                                     f"from table '{table_name}' to "
+                                     f"'{query.entity_table}'")
+                if edge_types[0] not in time_offset_dict:
+                    start = node.aggr_time_range.start_date_offset
+                    end = node.aggr_time_range.end_date_offset * num_forecasts
+                else:
+                    start, end = time_offset_dict[edge_types[0]]
+                    start = min_date_offset(
+                        start,
+                        node.aggr_time_range.start_date_offset,
+                    )
+                    end = max_date_offset(
+                        end,
+                        node.aggr_time_range.end_date_offset * num_forecasts,
+                    )
+                time_offset_dict[edge_types[0]] = (start, end)
+
+            for child in node.children:
+                _add_time_offset(child, num_forecasts)
+
+        _add_time_offset(query.target_ast, query.num_forecasts)
+        _add_time_offset(query.entity_ast)
+        if query.whatif_ast is not None:
+            _add_time_offset(query.whatif_ast)
+
+        return self._sample_forward(
+            query=query,
+            num_examples=num_examples,
+            anchor_time=anchor_time,
+            columns_dict=columns_dict,
+            time_offset_dict=time_offset_dict,
+            random_seed=random_seed,
+        )
+
     # Abstract Methods ########################################################
 
     @abstractmethod
-    def sample(
+    def _sample_backward(
         self,
         entity_table_name: str,
         entity_pkey: pd.Series,
         anchor_time: pd.Series,
-        column_spec_dict: dict[str, list[str]],
-        edge_spec_dict: dict[tuple[str, str, str], list[EdgeSpec]],
-        drop_duplicates: bool = False,
-        return_edges: bool = False,
-    ) -> SamplerOutput:
+        columns_dict: dict[str, set[str]],
+        num_neighbors: list[int],
+    ) -> BackwardSamplerOutput:
+        pass
+
+    @abstractmethod
+    def _sample_forward(
+        self,
+        query: ValidatedPredictiveQuery,
+        num_examples: int,
+        anchor_time: pd.Timestamp | Literal['entity'],
+        columns_dict: dict[str, set[str]],
+        time_offset_dict: dict[
+            tuple[str, str, str],
+            tuple[pd.DateOffset | None, pd.DateOffset],
+        ],
+        random_seed: int | None = None,
+    ) -> ForwardSamplerOutput:
         pass
 
 
@@ -285,3 +352,22 @@ def _normalize_text(
     ser = ser.map(normalize_fn)
 
     return ser
+
+
+def min_date_offset(*args: pd.DateOffset | None) -> pd.DateOffset | None:
+    if any(arg is None for arg in args):
+        return None
+
+    anchor = pd.Timestamp('2000-01-01')
+    timestamps = [anchor + arg for arg in args]
+    assert len(timestamps) > 0
+    argmin = min(range(len(timestamps)), key=lambda i: timestamps[i])
+    return args[argmin]
+
+
+def max_date_offset(*args: pd.DateOffset) -> pd.DateOffset:
+    anchor = pd.Timestamp('2000-01-01')
+    timestamps = [anchor + arg for arg in args]
+    assert len(timestamps) > 0
+    argmax = max(range(len(timestamps)), key=lambda i: timestamps[i])
+    return args[argmax]

{kumoai-2.13.0.dev202512061731.dist-info → kumoai-2.13.0.dev202512081731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.13.0.dev202512061731
+Version: 2.13.0.dev202512081731
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT

{kumoai-2.13.0.dev202512061731.dist-info → kumoai-2.13.0.dev202512081731.dist-info}/RECORD

@@ -1,7 +1,7 @@
 kumoai/kumolib.cpython-313-darwin.so,sha256=waBv-DiZ3WcasxiCQ-OM9EbSTgTtCfBTZIibXAK-JiQ,232816
 kumoai/_logging.py,sha256=U2_5ROdyk92P4xO4H2WJV8EC7dr6YxmmnM-b7QX9M7I,886
 kumoai/mixin.py,sha256=MP413xzuCqWhxAPUHmloLA3j4ZyF1tEtfi516b_hOXQ,812
-kumoai/_version.py,sha256=_62zq1gHlGnw4_Vj3YHa8OxfDzDGt7wK4h63pJl50_s,39
+kumoai/_version.py,sha256=W0EIBX5oPkQ0eXYnfNBgKMhonz56bp9ySi_IPtjQoCA,39
 kumoai/__init__.py,sha256=Nn9YH_x9kAeEFn8RWbP95slZow0qFnakPZZ1WADe1hY,10843
 kumoai/formatting.py,sha256=jA_rLDCGKZI8WWCha-vtuLenVKTZvli99Tqpurz1H84,953
 kumoai/futures.py,sha256=oJFIfdCM_3nWIqQteBKYMY4fPhoYlYWE_JA2o6tx-ng,3737
@@ -24,7 +24,7 @@ kumoai/experimental/rfm/backend/sqlite/table.py,sha256=kcYpWaZKFez2Tru6Sdz-Ywk8j
 kumoai/experimental/rfm/backend/local/__init__.py,sha256=2s9sSA-E-8pfkkzCH4XPuaSxSznEURMfMgwEIfYYPsg,1014
 kumoai/experimental/rfm/backend/local/table.py,sha256=Ahob9HidpU6z_M41rK5FATa3d7CL2UzZl8pGVyrzLNc,3565
 kumoai/experimental/rfm/backend/local/graph_store.py,sha256=RpfJldemOG-4RzGSIS9EcytHbvC4gYm-Ps3a-4qfptk,13297
-kumoai/experimental/rfm/backend/local/sampler.py,sha256=cLs0mJEoKWvfG-Wgm7tr_7DbwhliMLX8V0DyQqfGgtc,4613
+kumoai/experimental/rfm/backend/local/sampler.py,sha256=L1S2qxvkS_O8wy4K-czTxojPmklRrReTR8P3-e_8-hM,3823
 kumoai/experimental/rfm/backend/snow/__init__.py,sha256=B-tG-p8WA-mBuwvK1f0S2gdRPEGwApdxlnyeVSnY2xg,927
 kumoai/experimental/rfm/backend/snow/table.py,sha256=sHagXhW7RifzOiB4yjxV_9FtR0KUFVIw1mYwZe4bpMg,4255
 kumoai/experimental/rfm/pquery/__init__.py,sha256=X0O3EIq5SMfBEE-ii5Cq6iDhR3s3XMXB52Cx5htoePw,152
@@ -38,9 +38,9 @@ kumoai/experimental/rfm/infer/id.py,sha256=ZIO0DWIoiEoS_8MVc5lkqBfkTWWQ0yGCgjkwL
 kumoai/experimental/rfm/infer/dtype.py,sha256=ZZ6ztqJnTR1CaC2z5Uhf0o0rSdNThnss5tem5JNQkck,2607
 kumoai/experimental/rfm/infer/__init__.py,sha256=krdMFN8iKZlSFOl-M5MW1KuSviQV3H1E18jj2uB8g6Q,469
 kumoai/experimental/rfm/infer/timestamp.py,sha256=vM9--7eStzaGG13Y-oLYlpNJyhL6f9dp17HDXwtl_DM,1094
-kumoai/experimental/rfm/base/__init__.py,sha256=BD_oGL5sqYqGGPMR1EHQIZYYkOoM5dBFnYE5FNvON0U,295
+kumoai/experimental/rfm/base/__init__.py,sha256=3haYsIYypeL-U-9RuOOPnRdWaRlh-g_yE4ACJ2KLjOY,335
 kumoai/experimental/rfm/base/table.py,sha256=yaY7Auvq2KblXOid3-a_Pw6RgnPK5Y1zGAY2xi1D2gg,19843
-kumoai/experimental/rfm/base/sampler.py,sha256=ylJqCTJFBMCwXBsB6DDaYesQsOyqq7urq8pmpR2JdUI,10241
+kumoai/experimental/rfm/base/sampler.py,sha256=b45kllqSm-lpXbP9XbrGQPMx_hEIfesJILViAanh6rk,13456
 kumoai/experimental/rfm/base/source.py,sha256=8_waFQVsctryHkm9BwmFZ9-vw5cXAXfjk7KDmcl_kic,272
 kumoai/experimental/rfm/base/column.py,sha256=izCJmufJcd1RSi-ptFMfrue-JYag38MJxizka7ya0-A,2319
 kumoai/encoder/__init__.py,sha256=VPGs4miBC_WfwWeOXeHhFomOUocERFavhKf5fqITcds,182
@@ -106,8 +106,8 @@ kumoai/trainer/baseline_trainer.py,sha256=LlfViNOmswNv4c6zJJLsyv0pC2mM2WKMGYx06o
 kumoai/trainer/__init__.py,sha256=zUdFl-f-sBWmm2x8R-rdVzPBeU2FaMzUY5mkcgoTa1k,939
 kumoai/trainer/online_serving.py,sha256=9cddb5paeZaCgbUeceQdAOxysCtV5XP-KcsgFz_XR5w,9566
 kumoai/trainer/trainer.py,sha256=hBXO7gwpo3t59zKFTeIkK65B8QRmWCwO33sbDuEAPlY,20133
-kumoai-2.13.0.dev202512061731.dist-info/RECORD,,
-kumoai-2.13.0.dev202512061731.dist-info/WHEEL,sha256=oqGJCpG61FZJmvyZ3C_0aCv-2mdfcY9e3fXvyUNmWfM,136
-kumoai-2.13.0.dev202512061731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.13.0.dev202512061731.dist-info/METADATA,sha256=XT8M-fpeDwjsGaY1j8UtUESdReLI24cbccBCSF0k-Jo,2510
-kumoai-2.13.0.dev202512061731.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
+kumoai-2.13.0.dev202512081731.dist-info/RECORD,,
+kumoai-2.13.0.dev202512081731.dist-info/WHEEL,sha256=oqGJCpG61FZJmvyZ3C_0aCv-2mdfcY9e3fXvyUNmWfM,136
+kumoai-2.13.0.dev202512081731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.13.0.dev202512081731.dist-info/METADATA,sha256=ulcPeS_yowF-CWxGh5m_20ummlecVuiXBDNMgvXH-VU,2510
+kumoai-2.13.0.dev202512081731.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102