kumoai 2.14.0.dev202512141732__py3-none-any.whl → 2.15.0.dev202601131732__py3-none-any.whl

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

@@ -1,15 +1,15 @@
-import warnings
-from typing import List, Optional, cast
+from typing import Sequence, cast
 
 import pandas as pd
+from kumoapi.model_plan import MissingType
 
 from kumoai.experimental.rfm.base import (
+    ColumnSpec,
     DataBackend,
     SourceColumn,
     SourceForeignKey,
     Table,
 )
-from kumoai.experimental.rfm.infer import infer_dtype
 
 
 class LocalTable(Table):
@@ -57,9 +57,9 @@ class LocalTable(Table):
         self,
         df: pd.DataFrame,
         name: str,
-        primary_key: Optional[str] = None,
-        time_column: Optional[str] = None,
-        end_time_column: Optional[str] = None,
+        primary_key: MissingType | str | None = MissingType.VALUE,
+        time_column: str | None = None,
+        end_time_column: str | None = None,
     ) -> None:
 
         if df.empty:
@@ -75,7 +75,6 @@ class LocalTable(Table):
 
         super().__init__(
             name=name,
-            columns=list(df.columns),
             primary_key=primary_key,
             time_column=time_column,
             end_time_column=end_time_column,
@@ -85,35 +84,30 @@ class LocalTable(Table):
     def backend(self) -> DataBackend:
         return cast(DataBackend, DataBackend.LOCAL)
 
-    def _get_source_columns(self) -> List[SourceColumn]:
-        source_columns: List[SourceColumn] = []
-        for column in self._data.columns:
-            ser = self._data[column]
-            try:
-                dtype = infer_dtype(ser)
-            except Exception:
-                warnings.warn(f"Data type inference for column '{column}' in "
-                              f"table '{self.name}' failed. Consider changing "
-                              f"the data type of the column to use it within "
-                              f"this table.")
-                continue
-
-            source_column = SourceColumn(
-                name=column,
-                dtype=dtype,
+    def _get_source_columns(self) -> list[SourceColumn]:
+        return [
+            SourceColumn(
+                name=column_name,
+                dtype=None,
                 is_primary_key=False,
                 is_unique_key=False,
                 is_nullable=True,
-            )
-            source_columns.append(source_column)
+            ) for column_name in self._data.columns
+        ]
 
-        return source_columns
-
-    def _get_source_foreign_keys(self) -> List[SourceForeignKey]:
+    def _get_source_foreign_keys(self) -> list[SourceForeignKey]:
         return []
 
-    def _get_sample_df(self) -> pd.DataFrame:
+    def _get_source_sample_df(self) -> pd.DataFrame:
         return self._data
 
-    def _get_num_rows(self) -> Optional[int]:
+    def _get_expr_sample_df(
+        self,
+        columns: Sequence[ColumnSpec],
+    ) -> pd.DataFrame:
+        raise RuntimeError(f"Column expressions are not supported in "
+                           f"'{self.__class__.__name__}'. Please apply your "
+                           f"expressions on the `pd.DataFrame` directly.")
+
+    def _get_num_rows(self) -> int | None:
         return len(self._data)

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

@@ -1,35 +1,51 @@
-from typing import TYPE_CHECKING, Literal
+import json
+import math
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 from kumoapi.pquery import ValidatedPredictiveQuery
 
-from kumoai.experimental.rfm.backend.snow import SnowTable
-from kumoai.experimental.rfm.base import Sampler, SamplerOutput
-from kumoai.utils import ProgressLogger, quote_ident
+from kumoai.experimental.rfm.backend.snow import Connection, SnowTable
+from kumoai.experimental.rfm.base import SQLSampler, Table
+from kumoai.experimental.rfm.pquery import PQueryPandasExecutor
+from kumoai.utils import ProgressLogger
 
 if TYPE_CHECKING:
     from kumoai.experimental.rfm import Graph
 
 
-class SnowSampler(Sampler):
+@contextmanager
+def paramstyle(connection: Connection, style: str = 'qmark') -> Iterator[None]:
+    _style = connection._paramstyle
+    connection._paramstyle = style
+    yield
+    connection._paramstyle = _style
+
+
+class SnowSampler(SQLSampler):
     def __init__(
         self,
         graph: 'Graph',
         verbose: bool | ProgressLogger = True,
     ) -> None:
-        super().__init__(graph=graph)
+        super().__init__(graph=graph, verbose=verbose)
 
-        self._fqn_dict: dict[str, str] = {}
         for table in graph.tables.values():
             assert isinstance(table, SnowTable)
             self._connection = table._connection
-            self._fqn_dict[table.name] = table.fqn
+
+        self._num_rows_dict: dict[str, int] = {
+            table.name: cast(int, table._num_rows)
+            for table in graph.tables.values()
+        }
 
     @property
-    def fqn_dict(self) -> dict[str, str]:
-        r"""The fully-qualified quoted names for all tables in the graph."""
-        return self._fqn_dict
+    def num_rows_dict(self) -> dict[str, int]:
+        return self._num_rows_dict
 
     def _get_min_max_time_dict(
         self,
@@ -37,37 +53,28 @@ class SnowSampler(Sampler):
     ) -> dict[str, tuple[pd.Timestamp, pd.Timestamp]]:
         selects: list[str] = []
         for table_name in table_names:
-            time_column = self.time_column_dict[table_name]
+            column = self.time_column_dict[table_name]
+            column_ref = self.table_column_ref_dict[table_name][column]
             select = (f"SELECT\n"
-                      f"  %s as table_name,\n"
-                      f"  MIN({quote_ident(time_column)}) as min_date,\n"
-                      f"  MAX({quote_ident(time_column)}) as max_date\n"
-                      f"FROM {self.fqn_dict[table_name]}")
+                      f"  ? as table_name,\n"
+                      f"  MIN({column_ref}) as min_date,\n"
+                      f"  MAX({column_ref}) as max_date\n"
+                      f"FROM {self.source_name_dict[table_name]}")
             selects.append(select)
         sql = "\nUNION ALL\n".join(selects)
 
         out_dict: dict[str, tuple[pd.Timestamp, pd.Timestamp]] = {}
-        with self._connection.cursor() as cursor:
+        with paramstyle(self._connection), self._connection.cursor() as cursor:
             cursor.execute(sql, table_names)
             rows = cursor.fetchall()
-            for table_name, _min, _max in rows:
-                out_dict[table_name] = (
-                    pd.Timestamp.max if _min is None else pd.Timestamp(_min),
-                    pd.Timestamp.min if _max is None else pd.Timestamp(_max),
-                )
+        for table_name, _min, _max in rows:
+            out_dict[table_name] = (
+                pd.Timestamp.max if _min is None else pd.Timestamp(_min),
+                pd.Timestamp.min if _max is None else pd.Timestamp(_max),
+            )
 
         return out_dict
 
-    def _sample_subgraph(
-        self,
-        entity_table_name: str,
-        entity_pkey: pd.Series,
-        anchor_time: pd.Series | Literal['entity'],
-        columns_dict: dict[str, set[str]],
-        num_neighbors: list[int],
-    ) -> SamplerOutput:
-        raise NotImplementedError
-
     def _sample_entity_table(
         self,
         table_name: str,
@@ -78,17 +85,27 @@ class SnowSampler(Sampler):
         # NOTE Snowflake does support `SEED` only as part of `SYSTEM` sampling.
         num_rows = min(num_rows, 1_000_000)  # Snowflake's upper limit.
 
+        source_table = self.source_table_dict[table_name]
         filters: list[str] = []
-        primary_key = self.primary_key_dict[table_name]
-        if self.source_table_dict[table_name][primary_key].is_nullable:
-            filters.append(f" {quote_ident(primary_key)} IS NOT NULL")
-        time_column = self.time_column_dict.get(table_name)
-        if (time_column is not None and
-                self.source_table_dict[table_name][time_column].is_nullable):
-            filters.append(f" {quote_ident(time_column)} IS NOT NULL")
 
-        sql = (f"SELECT {', '.join(quote_ident(col) for col in columns)}\n"
-               f"FROM {self.fqn_dict[table_name]}\n"
+        key = self.primary_key_dict[table_name]
+        if key not in source_table or source_table[key].is_nullable:
+            key_ref = self.table_column_ref_dict[table_name][key]
+            filters.append(f" {key_ref} IS NOT NULL")
+
+        column = self.time_column_dict.get(table_name)
+        if column is None:
+            pass
+        elif column not in source_table or source_table[column].is_nullable:
+            column_ref = self.table_column_ref_dict[table_name][column]
+            filters.append(f" {column_ref} IS NOT NULL")
+
+        projections = [
+            self.table_column_proj_dict[table_name][column]
+            for column in columns
+        ]
+        sql = (f"SELECT {', '.join(projections)}\n"
+               f"FROM {self.source_name_dict[table_name]}\n"
                f"SAMPLE ROW ({num_rows} ROWS)")
         if len(filters) > 0:
             sql += f"\nWHERE{' AND'.join(filters)}"
@@ -98,7 +115,11 @@ class SnowSampler(Sampler):
             cursor.execute(sql)
             table = cursor.fetch_arrow_all()
 
-        return table.to_pandas(types_mapper=pd.ArrowDtype)
+        return Table._sanitize(
+            df=table.to_pandas(types_mapper=pd.ArrowDtype),
+            dtype_dict=self.table_dtype_dict[table_name],
+            stype_dict=self.table_stype_dict[table_name],
+        )
 
     def _sample_target(
         self,
@@ -116,4 +137,271 @@ class SnowSampler(Sampler):
             tuple[pd.DateOffset | None, pd.DateOffset],
         ],
     ) -> tuple[pd.Series, np.ndarray, pd.Series, np.ndarray]:
-        raise NotImplementedError
+
+        # NOTE For Snowflake, we execute everything at once to pay minimal
+        # query initialization costs.
+        index = np.concatenate([train_index, test_index])
+        time = pd.concat([train_time, test_time], axis=0, ignore_index=True)
+
+        entity_df = entity_df.iloc[index].reset_index(drop=True)
+
+        feat_dict: dict[str, pd.DataFrame] = {query.entity_table: entity_df}
+        time_dict: dict[str, pd.Series] = {}
+        time_column = self.time_column_dict.get(query.entity_table)
+        if time_column in columns_dict[query.entity_table]:
+            time_dict[query.entity_table] = entity_df[time_column]
+        batch_dict: dict[str, np.ndarray] = {
+            query.entity_table: np.arange(len(entity_df)),
+        }
+        for edge_type, (min_offset, max_offset) in time_offset_dict.items():
+            table_name, foreign_key, _ = edge_type
+            feat_dict[table_name], batch_dict[table_name] = self._by_time(
+                table_name=table_name,
+                foreign_key=foreign_key,
+                index=entity_df[self.primary_key_dict[query.entity_table]],
+                anchor_time=time,
+                min_offset=min_offset,
+                max_offset=max_offset,
+                columns=columns_dict[table_name],
+            )
+            time_column = self.time_column_dict.get(table_name)
+            if time_column in columns_dict[table_name]:
+                time_dict[table_name] = feat_dict[table_name][time_column]
+
+        y, mask = PQueryPandasExecutor().execute(
+            query=query,
+            feat_dict=feat_dict,
+            time_dict=time_dict,
+            batch_dict=batch_dict,
+            anchor_time=time,
+            num_forecasts=query.num_forecasts,
+        )
+
+        train_mask = mask[:len(train_index)]
+        test_mask = mask[len(train_index):]
+
+        boundary = int(train_mask.sum())
+        train_y = y.iloc[:boundary]
+        test_y = y.iloc[boundary:].reset_index(drop=True)
+
+        return train_y, train_mask, test_y, test_mask
+
+    def _by_pkey(
+        self,
+        table_name: str,
+        index: pd.Series,
+        columns: set[str],
+    ) -> tuple[pd.DataFrame, np.ndarray]:
+        key = self.primary_key_dict[table_name]
+        key_ref = self.table_column_ref_dict[table_name][key]
+        projections = [
+            self.table_column_proj_dict[table_name][column]
+            for column in columns
+        ]
+
+        payload = json.dumps(list(index))
+
+        sql = ("WITH TMP as (\n"
+               "  SELECT\n"
+               "    f.index as __KUMO_BATCH__,\n")
+        if self.table_dtype_dict[table_name][key].is_int():
+            sql += "    f.value::NUMBER as __KUMO_ID__\n"
+        elif self.table_dtype_dict[table_name][key].is_float():
+            sql += "    f.value::FLOAT as __KUMO_ID__\n"
+        else:
+            sql += "    f.value::VARCHAR as __KUMO_ID__\n"
+        sql += (f"  FROM TABLE(FLATTEN(INPUT => PARSE_JSON(?))) f\n"
+                f")\n"
+                f"SELECT "
+                f"TMP.__KUMO_BATCH__ as __KUMO_BATCH__, "
+                f"{', '.join(projections)}\n"
+                f"FROM TMP\n"
+                f"JOIN {self.source_name_dict[table_name]}\n"
+                f"  ON {key_ref} = TMP.__KUMO_ID__")
+
+        with paramstyle(self._connection), self._connection.cursor() as cursor:
+            cursor.execute(sql, (payload, ))
+            table = cursor.fetch_arrow_all()
+
+        # Remove any duplicated primary keys in post-processing:
+        tmp = table.append_column('__KUMO_ID__', pa.array(range(len(table))))
+        gb = tmp.group_by('__KUMO_BATCH__').aggregate([('__KUMO_ID__', 'min')])
+        table = table.take(gb['__KUMO_ID___min'])
+
+        batch = table['__KUMO_BATCH__'].cast(pa.int64()).to_numpy()
+        batch_index = table.schema.get_field_index('__KUMO_BATCH__')
+        table = table.remove_column(batch_index)
+
+        return Table._sanitize(
+            df=table.to_pandas(),
+            dtype_dict=self.table_dtype_dict[table_name],
+            stype_dict=self.table_stype_dict[table_name],
+        ), batch
+
+    def _by_fkey(
+        self,
+        table_name: str,
+        foreign_key: str,
+        index: pd.Series,
+        num_neighbors: int,
+        anchor_time: pd.Series | None,
+        columns: set[str],
+    ) -> tuple[pd.DataFrame, np.ndarray]:
+        time_column = self.time_column_dict.get(table_name)
+
+        end_time: pd.Series | None = None
+        start_time: pd.Series | None = None
+        if time_column is not None and anchor_time is not None:
+            # In order to avoid a full table scan, we limit foreign key
+            # sampling to a certain time range, approximated by the number of
+            # rows, timestamp ranges and `num_neighbors` value.
+            # Downstream, this helps Snowflake to apply partition pruning:
+            dst_table_name = [
+                dst_table
+                for key, dst_table in self.foreign_key_dict[table_name]
+                if key == foreign_key
+            ][0]
+            num_facts = self.num_rows_dict[table_name]
+            num_entities = self.num_rows_dict[dst_table_name]
+            min_time = self.get_min_time([table_name])
+            max_time = self.get_max_time([table_name])
+            freq = num_facts / num_entities
+            freq = freq / max((max_time - min_time).total_seconds(), 1)
+            offset = pd.Timedelta(seconds=math.ceil(5 * num_neighbors / freq))
+
+            end_time = anchor_time.dt.strftime("%Y-%m-%d %H:%M:%S")
+            start_time = anchor_time - offset
+            start_time = start_time.dt.strftime("%Y-%m-%d %H:%M:%S")
+            payload = json.dumps(list(zip(index, end_time, start_time)))
+        else:
+            payload = json.dumps(list(zip(index)))
+
+        key_ref = self.table_column_ref_dict[table_name][foreign_key]
+        projections = [
+            self.table_column_proj_dict[table_name][column]
+            for column in columns
+        ]
+
+        sql = ("WITH TMP as (\n"
+               "  SELECT\n"
+               "    f.index as __KUMO_BATCH__,\n")
+        if self.table_dtype_dict[table_name][foreign_key].is_int():
+            sql += "    f.value[0]::NUMBER as __KUMO_ID__"
+        elif self.table_dtype_dict[table_name][foreign_key].is_float():
+            sql += "    f.value[0]::FLOAT as __KUMO_ID__"
+        else:
+            sql += "    f.value[0]::VARCHAR as __KUMO_ID__"
+        if end_time is not None and start_time is not None:
+            sql += (",\n"
+                    "    f.value[1]::TIMESTAMP_NTZ as __KUMO_END_TIME__,\n"
+                    "    f.value[2]::TIMESTAMP_NTZ as __KUMO_START_TIME__")
+        sql += (f"\n"
+                f"  FROM TABLE(FLATTEN(INPUT => PARSE_JSON(?))) f\n"
+                f")\n"
+                f"SELECT "
+                f"TMP.__KUMO_BATCH__ as __KUMO_BATCH__, "
+                f"{', '.join(projections)}\n"
+                f"FROM TMP\n"
+                f"JOIN {self.source_name_dict[table_name]}\n"
+                f"  ON {key_ref} = TMP.__KUMO_ID__\n")
+        if end_time is not None and start_time is not None:
+            assert time_column is not None
+            time_ref = self.table_column_ref_dict[table_name][time_column]
+            sql += (f" AND {time_ref} <= TMP.__KUMO_END_TIME__\n"
+                    f" AND {time_ref} > TMP.__KUMO_START_TIME__\n"
+                    f"WHERE {time_ref} <= '{end_time.max()}'\n"
+                    f"  AND {time_ref} > '{start_time.min()}'\n")
+        sql += ("QUALIFY ROW_NUMBER() OVER (\n"
+                "  PARTITION BY TMP.__KUMO_BATCH__\n")
+        if time_column is not None:
+            sql += f"  ORDER BY {time_ref} DESC\n"
+        else:
+            sql += f"  ORDER BY {key_ref}\n"
+        sql += f") <= {num_neighbors}"
+
+        with paramstyle(self._connection), self._connection.cursor() as cursor:
+            cursor.execute(sql, (payload, ))
+            table = cursor.fetch_arrow_all()
+
+        batch = table['__KUMO_BATCH__'].cast(pa.int64()).to_numpy()
+        batch_index = table.schema.get_field_index('__KUMO_BATCH__')
+        table = table.remove_column(batch_index)
+
+        return Table._sanitize(
+            df=table.to_pandas(),
+            dtype_dict=self.table_dtype_dict[table_name],
+            stype_dict=self.table_stype_dict[table_name],
+        ), batch
+
+    # Helper Methods ##########################################################
+
+    def _by_time(
+        self,
+        table_name: str,
+        foreign_key: str,
+        index: pd.Series,
+        anchor_time: pd.Series,
+        min_offset: pd.DateOffset | None,
+        max_offset: pd.DateOffset,
+        columns: set[str],
+    ) -> tuple[pd.DataFrame, np.ndarray]:
+        time_column = self.time_column_dict[table_name]
+
+        end_time = anchor_time + max_offset
+        end_time = end_time.dt.strftime("%Y-%m-%d %H:%M:%S")
+        start_time: pd.Series | None = None
+        if min_offset is not None:
+            start_time = anchor_time + min_offset
+            start_time = start_time.dt.strftime("%Y-%m-%d %H:%M:%S")
+            payload = json.dumps(list(zip(index, end_time, start_time)))
+        else:
+            payload = json.dumps(list(zip(index, end_time)))
+
+        key_ref = self.table_column_ref_dict[table_name][foreign_key]
+        time_ref = self.table_column_ref_dict[table_name][time_column]
+        projections = [
+            self.table_column_proj_dict[table_name][column]
+            for column in columns
+        ]
+        sql = ("WITH TMP as (\n"
+               "  SELECT\n"
+               "    f.index as __KUMO_BATCH__,\n")
+        if self.table_dtype_dict[table_name][foreign_key].is_int():
+            sql += "    f.value[0]::NUMBER as __KUMO_ID__,\n"
+        elif self.table_dtype_dict[table_name][foreign_key].is_float():
+            sql += "    f.value[0]::FLOAT as __KUMO_ID__,\n"
+        else:
+            sql += "    f.value[0]::VARCHAR as __KUMO_ID__,\n"
+        sql += "    f.value[1]::TIMESTAMP_NTZ as __KUMO_END_TIME__"
+        if min_offset is not None:
+            sql += ",\n    f.value[2]::TIMESTAMP_NTZ as __KUMO_START_TIME__"
+        sql += (f"\n"
+                f"  FROM TABLE(FLATTEN(INPUT => PARSE_JSON(?))) f\n"
+                f")\n"
+                f"SELECT "
+                f"TMP.__KUMO_BATCH__ as __KUMO_BATCH__, "
+                f"{', '.join(projections)}\n"
+                f"FROM TMP\n"
+                f"JOIN {self.source_name_dict[table_name]}\n"
+                f"  ON {key_ref} = TMP.__KUMO_ID__\n"
+                f" AND {time_ref} <= TMP.__KUMO_END_TIME__\n")
+        if start_time is not None:
+            sql += f"AND {time_ref} > TMP.__KUMO_START_TIME__\n"
+        # Add global time bounds to enable partition pruning:
+        sql += f"WHERE {time_ref} <= '{end_time.max()}'"
+        if start_time is not None:
+            sql += f"\nAND {time_ref} > '{start_time.min()}'"
+
+        with paramstyle(self._connection), self._connection.cursor() as cursor:
+            cursor.execute(sql, (payload, ))
+            table = cursor.fetch_arrow_all()
+
+        batch = table['__KUMO_BATCH__'].cast(pa.int64()).to_numpy()
+        batch_index = table.schema.get_field_index('__KUMO_BATCH__')
+        table = table.remove_column(batch_index)
+
+        return Table._sanitize(
+            df=table.to_pandas(types_mapper=pd.ArrowDtype),
+            dtype_dict=self.table_dtype_dict[table_name],
+            stype_dict=self.table_stype_dict[table_name],
+        ), batch