kumoai 2.13.0.dev202512040252__cp310-cp310-win_amd64.whl → 2.15.0.dev202601141731__cp310-cp310-win_amd64.whl

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

@@ -27,9 +27,11 @@ def connect(**kwargs: Any) -> Connection:
 
 
 from .table import SnowTable  # noqa: E402
+from .sampler import SnowSampler  # noqa: E402
 
 __all__ = [
     'connect',
     'Connection',
     'SnowTable',
+    'SnowSampler',
 ]

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

@@ -0,0 +1,407 @@
+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 Connection, SnowTable
+from kumoai.experimental.rfm.base import SQLSampler, Table
+from kumoai.experimental.rfm.pquery import PQueryPandasExecutor
+from kumoai.utils import ProgressLogger, quote_ident
+
+if TYPE_CHECKING:
+    from kumoai.experimental.rfm import Graph
+
+
+@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, verbose=verbose)
+
+        for table in graph.tables.values():
+            assert isinstance(table, SnowTable)
+            self._connection = table._connection
+
+        self._num_rows_dict: dict[str, int] = {
+            table.name: cast(int, table._num_rows)
+            for table in graph.tables.values()
+        }
+
+    @property
+    def num_rows_dict(self) -> dict[str, int]:
+        return self._num_rows_dict
+
+    def _get_min_max_time_dict(
+        self,
+        table_names: list[str],
+    ) -> dict[str, tuple[pd.Timestamp, pd.Timestamp]]:
+        selects: list[str] = []
+        for table_name in table_names:
+            column = self.time_column_dict[table_name]
+            column_ref = self.table_column_ref_dict[table_name][column]
+            ident = quote_ident(table_name, char="'")
+            select = (f"SELECT\n"
+                      f"  {ident} 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:
+            cursor.execute(sql)
+            for table_name, _min, _max in cursor.fetchall():
+                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_entity_table(
+        self,
+        table_name: str,
+        columns: set[str],
+        num_rows: int,
+        random_seed: int | None = None,
+    ) -> pd.DataFrame:
+        # 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] = []
+
+        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)}"
+
+        with self._connection.cursor() as cursor:
+            # NOTE This may return duplicate primary keys. This is okay.
+            cursor.execute(sql)
+            table = cursor.fetch_arrow_all()
+
+        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,
+        query: ValidatedPredictiveQuery,
+        entity_df: pd.DataFrame,
+        train_index: np.ndarray,
+        train_time: pd.Series,
+        num_train_examples: int,
+        test_index: np.ndarray,
+        test_time: pd.Series,
+        num_test_examples: int,
+        columns_dict: dict[str, set[str]],
+        time_offset_dict: dict[
+            tuple[str, str, str],
+            tuple[pd.DateOffset | None, pd.DateOffset],
+        ],
+    ) -> tuple[pd.Series, np.ndarray, pd.Series, np.ndarray]:
+
+        # 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