kumoai 2.8.0.dev202508221830__cp312-cp312-win_amd64.whl → 2.13.0.dev202512041141__cp312-cp312-win_amd64.whl

kumoai/experimental/rfm/infer/__init__.py

@@ -1,9 +1,15 @@
+from .dtype import infer_dtype
+from .pkey import infer_primary_key
+from .time_col import infer_time_column
 from .id import contains_id
 from .timestamp import contains_timestamp
 from .categorical import contains_categorical
 from .multicategorical import contains_multicategorical
 
 __all__ = [
+    'infer_dtype',
+    'infer_primary_key',
+    'infer_time_column',
     'contains_id',
     'contains_timestamp',
     'contains_categorical',

kumoai/experimental/rfm/infer/dtype.py

@@ -0,0 +1,79 @@
+from typing import Dict
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+from kumoapi.typing import Dtype
+
+PANDAS_TO_DTYPE: Dict[str, Dtype] = {
+    'bool': Dtype.bool,
+    'boolean': Dtype.bool,
+    'int8': Dtype.int,
+    'int16': Dtype.int,
+    'int32': Dtype.int,
+    'int64': Dtype.int,
+    'float16': Dtype.float,
+    'float32': Dtype.float,
+    'float64': Dtype.float,
+    'object': Dtype.string,
+    'string': Dtype.string,
+    'string[python]': Dtype.string,
+    'string[pyarrow]': Dtype.string,
+    'binary': Dtype.binary,
+}
+
+
+def infer_dtype(ser: pd.Series) -> Dtype:
+    """Extracts the :class:`Dtype` from a :class:`pandas.Series`.
+
+    Args:
+        ser: A :class:`pandas.Series` to analyze.
+
+    Returns:
+        The data type.
+    """
+    if pd.api.types.is_datetime64_any_dtype(ser.dtype):
+        return Dtype.date
+    if pd.api.types.is_timedelta64_dtype(ser.dtype):
+        return Dtype.timedelta
+    if isinstance(ser.dtype, pd.CategoricalDtype):
+        return Dtype.string
+
+    if (pd.api.types.is_object_dtype(ser.dtype)
+            and not isinstance(ser.dtype, pd.ArrowDtype)):
+        index = ser.iloc[:1000].first_valid_index()
+        if index is not None and pd.api.types.is_list_like(ser[index]):
+            pos = ser.index.get_loc(index)
+            assert isinstance(pos, int)
+            ser = ser.iloc[pos:pos + 1000].dropna()
+            arr = pa.array(ser.tolist())
+            ser = pd.Series(arr, dtype=pd.ArrowDtype(arr.type))
+
+    if isinstance(ser.dtype, pd.ArrowDtype):
+        if pa.types.is_list(ser.dtype.pyarrow_dtype):
+            elem_dtype = ser.dtype.pyarrow_dtype.value_type
+            if pa.types.is_integer(elem_dtype):
+                return Dtype.intlist
+            if pa.types.is_floating(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_decimal(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_string(elem_dtype):
+                return Dtype.stringlist
+            if pa.types.is_null(elem_dtype):
+                return Dtype.floatlist
+
+    if isinstance(ser.dtype, np.dtype):
+        dtype_str = str(ser.dtype).lower()
+    elif isinstance(ser.dtype, pd.api.extensions.ExtensionDtype):
+        dtype_str = ser.dtype.name.lower()
+        dtype_str = dtype_str.split('[')[0]  # Remove backend metadata
+    elif isinstance(ser.dtype, pa.DataType):
+        dtype_str = str(ser.dtype).lower()
+    else:
+        dtype_str = 'object'
+
+    if dtype_str not in PANDAS_TO_DTYPE:
+        raise ValueError(f"Unsupported data type '{ser.dtype}'")
+
+    return PANDAS_TO_DTYPE[dtype_str]

kumoai/experimental/rfm/infer/pkey.py

@@ -0,0 +1,126 @@
+import re
+import warnings
+from typing import Optional
+
+import pandas as pd
+
+
+def infer_primary_key(
+    table_name: str,
+    df: pd.DataFrame,
+    candidates: list[str],
+) -> Optional[str]:
+    r"""Auto-detect potential primary key column.
+
+    Args:
+        table_name: The table name.
+        df: The pandas DataFrame to analyze.
+        candidates: A list of potential candidates.
+
+    Returns:
+        The name of the detected primary key, or ``None`` if not found.
+    """
+    # A list of (potentially modified) table names that are eligible to match
+    # with a primary key, i.e.:
+    # - UserInfo -> User
+    # - snakecase <-> camelcase
+    # - camelcase <-> snakecase
+    # - plural <-> singular (users -> user, eligibilities -> eligibility)
+    # - verb -> noun (qualifying -> qualify)
+    _table_names = {table_name}
+    if table_name.lower().endswith('_info'):
+        _table_names.add(table_name[:-5])
+    elif table_name.lower().endswith('info'):
+        _table_names.add(table_name[:-4])
+
+    table_names = set()
+    for _table_name in _table_names:
+        table_names.add(_table_name.lower())
+        snakecase = re.sub(r'(.)([A-Z][a-z]+)', r'\1_\2', _table_name)
+        snakecase = re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', snakecase)
+        table_names.add(snakecase.lower())
+        camelcase = _table_name.replace('_', '')
+        table_names.add(camelcase.lower())
+        if _table_name.lower().endswith('s'):
+            table_names.add(_table_name.lower()[:-1])
+            table_names.add(snakecase.lower()[:-1])
+            table_names.add(camelcase.lower()[:-1])
+        else:
+            table_names.add(_table_name.lower() + 's')
+            table_names.add(snakecase.lower() + 's')
+            table_names.add(camelcase.lower() + 's')
+        if _table_name.lower().endswith('ies'):
+            table_names.add(_table_name.lower()[:-3] + 'y')
+            table_names.add(snakecase.lower()[:-3] + 'y')
+            table_names.add(camelcase.lower()[:-3] + 'y')
+        elif _table_name.lower().endswith('y'):
+            table_names.add(_table_name.lower()[:-1] + 'ies')
+            table_names.add(snakecase.lower()[:-1] + 'ies')
+            table_names.add(camelcase.lower()[:-1] + 'ies')
+        if _table_name.lower().endswith('ing'):
+            table_names.add(_table_name.lower()[:-3])
+            table_names.add(snakecase.lower()[:-3])
+            table_names.add(camelcase.lower()[:-3])
+
+    scores: list[tuple[str, int]] = []
+    for col_name in candidates:
+        col_name_lower = col_name.lower()
+
+        score = 0
+
+        if col_name_lower == 'id':
+            score += 4
+
+        for table_name_lower in table_names:
+
+            if col_name_lower == table_name_lower:
+                score += 4  # USER -> USER
+                break
+
+            for suffix in ['id', 'hash', 'key', 'code', 'uuid']:
+                if not col_name_lower.endswith(suffix):
+                    continue
+
+                if col_name_lower == f'{table_name_lower}_{suffix}':
+                    score += 5  # USER -> USER_ID
+                    break
+
+                if col_name_lower == f'{table_name_lower}{suffix}':
+                    score += 5  # User -> UserId
+                    break
+
+                if col_name_lower.endswith(f'{table_name_lower}_{suffix}'):
+                    score += 2
+
+                if col_name_lower.endswith(f'{table_name_lower}{suffix}'):
+                    score += 2
+
+            # `rel-bench` hard-coding :(
+            if table_name == 'studies' and col_name == 'nct_id':
+                score += 1
+
+        ser = df[col_name].iloc[:1_000_000]
+        score += 3 * (ser.nunique() / len(ser))
+
+        scores.append((col_name, score))
+
+    scores = [x for x in scores if x[-1] >= 4]
+    scores.sort(key=lambda x: x[-1], reverse=True)
+
+    if len(scores) == 0:
+        return None
+
+    if len(scores) == 1:
+        return scores[0][0]
+
+    # In case of multiple candidates, only return one if its score is unique:
+    if scores[0][1] != scores[1][1]:
+        return scores[0][0]
+
+    max_score = max(scores, key=lambda x: x[1])
+    candidates = [col_name for col_name, score in scores if score == max_score]
+    warnings.warn(f"Found multiple potential primary keys in table "
+                  f"'{table_name}': {candidates}. Please specify the primary "
+                  f"key for this table manually.")
+
+    return None

kumoai/experimental/rfm/infer/time_col.py

@@ -0,0 +1,62 @@
+import re
+import warnings
+from typing import Optional
+
+import pandas as pd
+
+
+def infer_time_column(
+    df: pd.DataFrame,
+    candidates: list[str],
+) -> Optional[str]:
+    r"""Auto-detect potential time column.
+
+    Args:
+        df: The pandas DataFrame to analyze.
+        candidates: A list of potential candidates.
+
+    Returns:
+        The name of the detected time column, or ``None`` if not found.
+    """
+    candidates = [  # Exclude all candidates with `*last*` in column names:
+        col_name for col_name in candidates
+        if not re.search(r'(^|_)last(_|$)', col_name, re.IGNORECASE)
+    ]
+
+    if len(candidates) == 0:
+        return None
+
+    if len(candidates) == 1:
+        return candidates[0]
+
+    # If there exists a dedicated `create*` column, use it as time column:
+    create_candidates = [
+        candidate for candidate in candidates
+        if candidate.lower().startswith('create')
+    ]
+    if len(create_candidates) == 1:
+        return create_candidates[0]
+    if len(create_candidates) > 1:
+        candidates = create_candidates
+
+    # Find the most optimal time column. Usually, it is the one pointing to
+    # the oldest timestamps:
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', message='Could not infer format')
+        min_timestamp_dict = {
+            key: pd.to_datetime(df[key].iloc[:10_000], 'coerce')
+            for key in candidates
+        }
+    min_timestamp_dict = {
+        key: value.min().tz_localize(None)
+        for key, value in min_timestamp_dict.items()
+    }
+    min_timestamp_dict = {
+        key: value
+        for key, value in min_timestamp_dict.items() if not pd.isna(value)
+    }
+
+    if len(min_timestamp_dict) == 0:
+        return None
+
+    return min(min_timestamp_dict, key=min_timestamp_dict.get)  # type: ignore

kumoai/experimental/rfm/infer/timestamp.py

@@ -2,6 +2,7 @@ import re
 import warnings
 
 import pandas as pd
+from dateutil.parser import UnknownTimezoneWarning
 from kumoapi.typing import Dtype, Stype
 
 
@@ -20,9 +21,7 @@ def contains_timestamp(ser: pd.Series, column_name: str, dtype: Dtype) -> bool:
         column_name,
         re.IGNORECASE,
     )
-
-    if match is not None:
-        return True
+    score = 0.3 if match is not None else 0.0
 
     ser = ser.iloc[:100]
     ser = ser.dropna()
@@ -34,5 +33,9 @@ def contains_timestamp(ser: pd.Series, column_name: str, dtype: Dtype) -> bool:
     ser = ser.astype(str)  # Avoid parsing numbers as unix timestamps.
 
     with warnings.catch_warnings():
+        warnings.simplefilter('ignore', UnknownTimezoneWarning)
         warnings.filterwarnings('ignore', message='Could not infer format')
-        return pd.to_datetime(ser, errors='coerce').notna().all()
+        mask = pd.to_datetime(ser, errors='coerce').notna()
+        score += int(mask.sum()) / len(mask)
+
+    return score >= 1.0

kumoai/experimental/rfm/local_graph_sampler.py

@@ -1,14 +1,54 @@
+import re
 from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import pandas as pd
-from kumoapi.model_plan import RunMode
 from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
 from kumoapi.typing import Stype
 
 import kumoai.kumolib as kumolib
 from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
-from kumoai.experimental.rfm.utils import normalize_text
+
+PUNCTUATION = re.compile(r"[\'\"\.,\(\)\!\?\;\:]")
+MULTISPACE = re.compile(r"\s+")
+
+
+def normalize_text(
+    ser: pd.Series,
+    max_words: Optional[int] = 50,
+) -> pd.Series:
+    r"""Normalizes text into a list of lower-case words.
+
+    Args:
+        ser: The :class:`pandas.Series` to normalize.
+        max_words: The maximum number of words to return.
+            This will auto-shrink any large text column to avoid blowing up
+            context size.
+    """
+    if len(ser) == 0 or pd.api.types.is_list_like(ser.iloc[0]):
+        return ser
+
+    def normalize_fn(line: str) -> list[str]:
+        line = PUNCTUATION.sub(" ", line)
+        line = re.sub(r"<br\s*/?>", " ", line)  # Handle <br /> or <br>
+        line = MULTISPACE.sub(" ", line)
+        words = line.split()
+        if max_words is not None:
+            words = words[:max_words]
+        return words
+
+    ser = ser.fillna('').astype(str)
+
+    if max_words is not None:
+        # We estimate the number of words as 5 characters + 1 space in an
+        # English text on average. We need this pre-filter here, as word
+        # splitting on a giant text can be very expensive:
+        ser = ser.str[:6 * max_words]
+
+    ser = ser.str.lower()
+    ser = ser.map(normalize_fn)
+
+    return ser
 
 
 class LocalGraphSampler:
@@ -33,7 +73,6 @@ class LocalGraphSampler:
         entity_table_names: Tuple[str, ...],
         node: np.ndarray,
         time: np.ndarray,
-        run_mode: RunMode,
         num_neighbors: List[int],
         exclude_cols_dict: Dict[str, List[str]],
     ) -> Subgraph:
@@ -92,15 +131,23 @@ class LocalGraphSampler:
                 )
                 continue
 
-            # Only store unique rows in `df` above a certain threshold:
-            unique_node, inverse_node = np.unique(node, return_inverse=True)
-            if len(node) > 1.05 * len(unique_node):
-                df = df.iloc[unique_node]
-                row = inverse_node
+            row: Optional[np.ndarray] = None
+            if table_name in self._graph_store.end_time_column_dict:
+                # Set end time to NaT for all values greater than anchor time:
+                df = df.iloc[node].reset_index(drop=True)
+                col_name = self._graph_store.end_time_column_dict[table_name]
+                ser = df[col_name]
+                value = ser.astype('datetime64[ns]').astype(int).to_numpy()
+                mask = value > time[batch]
+                df.loc[mask, col_name] = pd.NaT
             else:
-                df = df.iloc[node]
-                row = None
-            df = df.reset_index(drop=True)
+                # 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):
+                    df = df.iloc[unique_node].reset_index(drop=True)
+                    row = inverse
+                else:
+                    df = df.iloc[node].reset_index(drop=True)
 
             # Filter data frame to minimal set of columns:
             df = df[columns]

kumoai/experimental/rfm/local_graph_store.py

@@ -1,14 +1,13 @@
 import warnings
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
 from kumoapi.rfm.context import Subgraph
 from kumoapi.typing import Stype
 
-from kumoai.experimental.rfm import LocalGraph
-from kumoai.experimental.rfm.utils import normalize_text
-from kumoai.utils import ProgressLogger
+from kumoai.experimental.rfm import Graph, LocalTable
+from kumoai.utils import InteractiveProgressLogger, ProgressLogger
 
 try:
     import torch
@@ -20,13 +19,18 @@ except ImportError:
 class LocalGraphStore:
     def __init__(
         self,
-        graph: LocalGraph,
-        preprocess: bool = False,
-        verbose: bool = True,
+        graph: Graph,
+        verbose: Union[bool, ProgressLogger] = True,
     ) -> None:
 
-        with ProgressLogger("Materializing graph", verbose=verbose) as logger:
-            self.df_dict, self.mask_dict = self.sanitize(graph, preprocess)
+        if not isinstance(verbose, ProgressLogger):
+            verbose = InteractiveProgressLogger(
+                "Materializing graph",
+                verbose=verbose,
+            )
+
+        with verbose as logger:
+            self.df_dict, self.mask_dict = self.sanitize(graph)
             self.stype_dict = self.get_stype_dict(graph)
             logger.log("Sanitized input data")
 
@@ -39,6 +43,7 @@ class LocalGraphStore:
 
             (
                 self.time_column_dict,
+                self.end_time_column_dict,
                 self.time_dict,
                 self.min_time,
                 self.max_time,
@@ -98,8 +103,7 @@ class LocalGraphStore:
 
     def sanitize(
         self,
-        graph: LocalGraph,
-        preprocess: bool = False,
+        graph: Graph,
     ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, np.ndarray]]:
         r"""Sanitizes raw data according to table schema definition:
 
@@ -108,17 +112,12 @@ class LocalGraphStore:
         * drops timezone information from timestamps
         * drops duplicate primary keys
         * removes rows with missing primary keys or time values
-
-        If ``preprocess`` is set to ``True``, it will additionally pre-process
-        data for faster model processing. In particular, it:
-        * tokenizes any text column that is not a foreign key
         """
-        df_dict: Dict[str, pd.DataFrame] = {
-            table_name: table._data.copy(deep=False).reset_index(drop=True)
-            for table_name, table in graph.tables.items()
-        }
-
-        foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
+        df_dict: Dict[str, pd.DataFrame] = {}
+        for table_name, table in graph.tables.items():
+            assert isinstance(table, LocalTable)
+            df = table._data
+            df_dict[table_name] = df.copy(deep=False).reset_index(drop=True)
 
         mask_dict: Dict[str, np.ndarray] = {}
         for table in graph.tables.values():
@@ -137,12 +136,6 @@ class LocalGraphStore:
                         ser = ser.dt.tz_localize(None)
                         df_dict[table.name][col.name] = ser
 
-                # Normalize text in advance (but exclude foreign keys):
-                if (preprocess and col.stype == Stype.text
-                        and (table.name, col.name) not in foreign_keys):
-                    ser = df_dict[table.name][col.name]
-                    df_dict[table.name][col.name] = normalize_text(ser)
-
             mask: Optional[np.ndarray] = None
             if table._time_column is not None:
                 ser = df_dict[table.name][table._time_column]
@@ -158,7 +151,7 @@ class LocalGraphStore:
 
         return df_dict, mask_dict
 
-    def get_stype_dict(self, graph: LocalGraph) -> Dict[str, Dict[str, Stype]]:
+    def get_stype_dict(self, graph: Graph) -> Dict[str, Dict[str, Stype]]:
         stype_dict: Dict[str, Dict[str, Stype]] = {}
         foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
         for table in graph.tables.values():
@@ -173,7 +166,7 @@ class LocalGraphStore:
 
     def get_pkey_data(
         self,
-        graph: LocalGraph,
+        graph: Graph,
     ) -> Tuple[
             Dict[str, str],
             Dict[str, pd.DataFrame],
@@ -195,11 +188,15 @@ class LocalGraphStore:
                 pkey_map = pkey_map[self.mask_dict[table.name]]
 
             if len(pkey_map) == 0:
-                raise ValueError(
-                    f"Found no valid rows in table '{table.name}' since there "
-                    f"exists not a single row with a non-N/A primary key."
-                    f"Consider fixing your underlying data or removing this "
-                    f"table from the graph.")
+                error_msg = f"Found no valid rows in table '{table.name}'. "
+                if table.has_time_column():
+                    error_msg += ("Please make sure that there exists valid "
+                                  "non-N/A primary key and time column pairs "
+                                  "in this table.")
+                else:
+                    error_msg += ("Please make sure that there exists valid "
+                                  "non-N/A primary keys in this table.")
+                raise ValueError(error_msg)
 
             pkey_map_dict[table.name] = pkey_map
 
@@ -207,18 +204,23 @@ class LocalGraphStore:
 
     def get_time_data(
         self,
-        graph: LocalGraph,
+        graph: Graph,
     ) -> Tuple[
+            Dict[str, str],
             Dict[str, str],
             Dict[str, np.ndarray],
             pd.Timestamp,
             pd.Timestamp,
     ]:
         time_column_dict: Dict[str, str] = {}
+        end_time_column_dict: Dict[str, str] = {}
         time_dict: Dict[str, np.ndarray] = {}
         min_time = pd.Timestamp.max
         max_time = pd.Timestamp.min
         for table in graph.tables.values():
+            if table._end_time_column is not None:
+                end_time_column_dict[table.name] = table._end_time_column
+
             if table._time_column is None:
                 continue
 
@@ -233,11 +235,17 @@ class LocalGraphStore:
                 min_time = min(min_time, time.min())
                 max_time = max(max_time, time.max())
 
-        return time_column_dict, time_dict, min_time, max_time
+        return (
+            time_column_dict,
+            end_time_column_dict,
+            time_dict,
+            min_time,
+            max_time,
+        )
 
     def get_csc(
         self,
-        graph: LocalGraph,
+        graph: Graph,
     ) -> Tuple[
             Dict[Tuple[str, str, str], np.ndarray],
             Dict[Tuple[str, str, str], np.ndarray],