kumoai 2.14.0.dev202512191731__cp311-cp311-macosx_11_0_arm64.whl → 2.15.0.dev202601141731__cp311-cp311-macosx_11_0_arm64.whl

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

@@ -1,16 +1,21 @@
 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
-from kumoai.experimental.rfm.base import SQLSampler
+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 quote_ident
+from kumoai.utils import ProgressLogger, quote_ident
+
+if TYPE_CHECKING:
+    from kumoai.experimental.rfm import Graph
 
 
 @contextmanager
@@ -22,30 +27,51 @@ def paramstyle(connection: Connection, style: str = 'qmark') -> Iterator[None]:
 
 
 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:
-            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]
+            ident = quote_ident(table_name, char="'")
             select = (f"SELECT\n"
-                      f"  ? 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"  {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 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),
-            )
+        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
 
@@ -59,17 +85,27 @@ class SnowSampler(SQLSampler):
         # 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)}"
@@ -79,7 +115,11 @@ class SnowSampler(SQLSampler):
             cursor.execute(sql)
             table = cursor.fetch_arrow_all()
 
-        return self._sanitize(table_name, table)
+        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,
@@ -114,11 +154,11 @@ class SnowSampler(SQLSampler):
             query.entity_table: np.arange(len(entity_df)),
         }
         for edge_type, (min_offset, max_offset) in time_offset_dict.items():
-            table_name, fkey, _ = edge_type
+            table_name, foreign_key, _ = edge_type
             feat_dict[table_name], batch_dict[table_name] = self._by_time(
                 table_name=table_name,
-                fkey=fkey,
-                pkey=entity_df[self.primary_key_dict[query.entity_table]],
+                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,
@@ -149,104 +189,219 @@ class SnowSampler(SQLSampler):
     def _by_pkey(
         self,
         table_name: str,
-        pkey: pd.Series,
+        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
+        ]
 
-        pkey_name = self.primary_key_dict[table_name]
-        source_table = self.source_table_dict[table_name]
-
-        payload = json.dumps(list(pkey))
+        payload = json.dumps(list(index))
 
         sql = ("WITH TMP as (\n"
                "  SELECT\n"
-               "    f.index as BATCH,\n")
-        if source_table[pkey_name].dtype.is_int():
-            sql += "    f.value::NUMBER as ID\n"
-        elif source_table[pkey_name].dtype.is_float():
-            sql += "    f.value::FLOAT as ID\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 ID\n"
+            sql += "    f.value::VARCHAR as __KUMO_ID__\n"
         sql += (f"  FROM TABLE(FLATTEN(INPUT => PARSE_JSON(?))) f\n"
                 f")\n"
-                f"SELECT TMP.BATCH as __BATCH__, "
-                f"{', '.join('ENT.' + quote_ident(col) for col in columns)}\n"
+                f"SELECT "
+                f"TMP.__KUMO_BATCH__ as __KUMO_BATCH__, "
+                f"{', '.join(projections)}\n"
                 f"FROM TMP\n"
-                f"JOIN {self.fqn_dict[table_name]} ENT\n"
-                f"  ON ENT.{quote_ident(pkey_name)} = TMP.ID")
+                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('__TMP__', pa.array(range(len(table))))
-        gb = tmp.group_by('__BATCH__').aggregate([('__TMP__', 'min')])
-        table = table.take(gb['__TMP___min'])
+        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)
 
-        batch = table['__BATCH__'].cast(pa.int64()).to_numpy()
-        table = table.remove_column(table.schema.get_field_index('__BATCH__'))
+        return Table._sanitize(
+            df=table.to_pandas(),
+            dtype_dict=self.table_dtype_dict[table_name],
+            stype_dict=self.table_stype_dict[table_name],
+        ), batch
 
-        return table.to_pandas(), batch  # TODO Use `self._sanitize`.
+    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,
-        fkey: str,
-        pkey: pd.Series,
+        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(pkey, end_time, start_time)))
+            payload = json.dumps(list(zip(index, end_time, start_time)))
         else:
-            payload = json.dumps(list(zip(pkey, end_time)))
-
-        # Based on benchmarking, JSON payload is the fastest way to query by
-        # custom indices (compared to large `IN` clauses or temporary tables):
-        source_table = self.source_table_dict[table_name]
-        time_column = self.time_column_dict[table_name]
+            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 BATCH,\n")
-        if source_table[fkey].dtype.is_int():
-            sql += "    f.value[0]::NUMBER as ID,\n"
-        elif source_table[fkey].dtype.is_float():
-            sql += "    f.value[0]::FLOAT as ID,\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 ID,\n"
-        sql += "    f.value[1]::TIMESTAMP_NTZ as END_TIME"
+            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 START_TIME"
+            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 TMP.BATCH as __BATCH__, "
-                f"{', '.join('FACT.' + quote_ident(col) for col in columns)}\n"
+                f"SELECT "
+                f"TMP.__KUMO_BATCH__ as __KUMO_BATCH__, "
+                f"{', '.join(projections)}\n"
                 f"FROM TMP\n"
-                f"JOIN {self.fqn_dict[table_name]} FACT\n"
-                f"  ON FACT.{quote_ident(fkey)} = TMP.ID\n"
-                f" AND FACT.{quote_ident(time_column)} <= TMP.END_TIME")
-        if min_offset is not None:
-            sql += f"\n AND FACT.{quote_ident(time_column)} > TMP.START_TIME"
+                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['__BATCH__'].cast(pa.int64()).to_numpy()
-        table = table.remove_column(table.schema.get_field_index('__BATCH__'))
-
-        return self._sanitize(table_name, table), batch
+        batch = table['__KUMO_BATCH__'].cast(pa.int64()).to_numpy()
+        batch_index = table.schema.get_field_index('__KUMO_BATCH__')
+        table = table.remove_column(batch_index)
 
-    def _sanitize(self, table_name: str, table: pa.table) -> pd.DataFrame:
-        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],
+        ), batch