kumoai 2.14.0.dev202512301731__cp312-cp312-macosx_11_0_arm64.whl → 2.14.0.dev202601041732__cp312-cp312-macosx_11_0_arm64.whl

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.14.0.dev202512301731'
+__version__ = '2.14.0.dev202601041732'

kumoai/connector/utils.py

@@ -1,10 +1,10 @@
 import asyncio
 import csv
-import gc
 import io
 import math
 import os
 import re
+import sys
 import tempfile
 import threading
 import time
@@ -920,7 +920,10 @@ def _read_remote_file_with_progress(
     if capture_first_line and not seen_nl:
         header_line = bytes(header_acc)
 
-    mv = buf.getbuffer()  # zero-copy view of BytesIO internal buffer
+    if sys.version_info >= (3, 13):
+        mv = memoryview(buf.getvalue())
+    else:
+        mv = buf.getbuffer()  # zero-copy view of BytesIO internal buffer
     return buf, mv, header_line
 
 
@@ -999,7 +1002,10 @@ def _iter_mv_chunks(mv: memoryview,
     n = mv.nbytes
     while pos < n:
         nxt = min(n, pos + part_size)
-        yield mv[pos:nxt]  # zero-copy slice
+        if sys.version_info >= (3, 13):
+            yield mv[pos:nxt].tobytes()
+        else:
+            yield mv[pos:nxt]  # zero-copy slice
         pos = nxt
 
 
@@ -1473,13 +1479,17 @@ def _remote_upload_file(name: str, fs: Filesystem, url: str, info: dict,
     if renamed_cols_msg:
         logger.info(renamed_cols_msg)
 
+    try:
+        if isinstance(data_mv, memoryview):
+            data_mv.release()
+    except Exception:
+        pass
+
     try:
         if buf:
             buf.close()
     except Exception:
         pass
-    del buf, data_mv, header_line
-    gc.collect()
 
     logger.info("Upload complete. Validated table %s.", name)
 
@@ -1719,13 +1729,17 @@ def _remote_upload_directory(
                     else:
                         break
 
+                try:
+                    if isinstance(data_mv, memoryview):
+                        data_mv.release()
+                except Exception:
+                    pass
+
                 try:
                     if buf:
                         buf.close()
                 except Exception:
                     pass
-                del buf, data_mv, header_line
-                gc.collect()
 
                 _safe_bar_update(file_bar, 1)
                 _merge_status_update(fpath)

kumoai/experimental/rfm/__init__.py

@@ -20,6 +20,7 @@ from .sagemaker import (
 from .base import Table
 from .backend.local import LocalTable
 from .graph import Graph
+from .task_table import TaskTable
 from .rfm import ExplainConfig, Explanation, KumoRFM
 
 logger = logging.getLogger('kumoai_rfm')
@@ -202,6 +203,7 @@ __all__ = [
     'Table',
     'LocalTable',
     'Graph',
+    'TaskTable',
     'KumoRFM',
     'ExplainConfig',
     'Explanation',

kumoai/experimental/rfm/base/table.py

@@ -271,8 +271,8 @@ class Table(ABC):
         no such primary key is present.
 
         The setter sets a column as a primary key on this table, and raises a
-        :class:`ValueError` if the primary key has a non-ID semantic type or
-        if the column name does not match a column in the data frame.
+        :class:`ValueError` if the primary key has a non-ID compatible data
+        type or if the column name does not match a column in the data frame.
         """
         if self._primary_key is None:
             return None
@@ -316,8 +316,9 @@ class Table(ABC):
         such time column is present.
 
         The setter sets a column as a time column on this table, and raises a
-        :class:`ValueError` if the time column has a non-timestamp semantic
-        type or if the column name does not match a column in the data frame.
+        :class:`ValueError` if the time column has a non-timestamp compatible
+        data type or if the column name does not match a column in the data
+        frame.
         """
         if self._time_column is None:
             return None
@@ -362,8 +363,8 @@ class Table(ABC):
 
         The setter sets a column as an end time column on this table, and
         raises a :class:`ValueError` if the end time column has a non-timestamp
-        semantic type or if the column name does not match a column in the data
-        frame.
+        compatible data type or if the column name does not match a column in
+        the data frame.
         """
         if self._end_time_column is None:
             return None

kumoai/experimental/rfm/rfm.py

@@ -726,30 +726,34 @@ class KumoRFM:
                              "`predict()` or `evaluate()` methods to perform "
                              "predictions or evaluations.")
 
-        try:
-            request = RFMParseQueryRequest(
-                query=query,
-                graph_definition=self._graph_def,
-            )
-
-            resp = self._api_client.parse_query(request)
-
-            if len(resp.validation_response.warnings) > 0:
-                msg = '\n'.join([
-                    f'{i+1}. {warning.title}: {warning.message}' for i, warning
-                    in enumerate(resp.validation_response.warnings)
-                ])
-                warnings.warn(f"Encountered the following warnings during "
-                              f"parsing:\n{msg}")
+        request = RFMParseQueryRequest(
+            query=query,
+            graph_definition=self._graph_def,
+        )
 
-            return resp.query
-        except HTTPException as e:
+        for attempt in range(self.num_retries + 1):
             try:
-                msg = json.loads(e.detail)['detail']
-            except Exception:
-                msg = e.detail
-            raise ValueError(f"Failed to parse query '{query}'. "
-                             f"{msg}") from None
+                resp = self._api_client.parse_query(request)
+                break
+            except HTTPException as e:
+                if attempt == self.num_retries:
+                    try:
+                        msg = json.loads(e.detail)['detail']
+                    except Exception:
+                        msg = e.detail
+                    raise ValueError(f"Failed to parse query '{query}'. {msg}")
+
+                time.sleep(2**attempt)  # 1s, 2s, 4s, 8s, ...
+
+        if len(resp.validation_response.warnings) > 0:
+            msg = '\n'.join([
+                f'{i+1}. {warning.title}: {warning.message}'
+                for i, warning in enumerate(resp.validation_response.warnings)
+            ])
+            warnings.warn(f"Encountered the following warnings during "
+                          f"parsing:\n{msg}")
+
+        return resp.query
 
     @staticmethod
     def _get_task_type(

kumoai/experimental/rfm/task_table.py

@@ -0,0 +1,247 @@
+import copy
+from collections.abc import Sequence
+
+import pandas as pd
+from kumoapi.task import TaskType
+from kumoapi.typing import Dtype, Stype
+from typing_extensions import Self
+
+from kumoai.experimental.rfm.base import Column
+from kumoai.experimental.rfm.infer import contains_timestamp, infer_dtype
+
+
+class TaskTable:
+    r"""A :class:`TaskTable` fully specifies the task, *i.e.* its context and
+    prediction examples with entity IDs, targets and timestamps.
+
+    Args:
+        task_type: The task type.
+        context_df: The data frame holding context examples.
+        pred_df: The data frame holding prediction examples.
+        entity_table_name: The entity table to predict for. For link prediction
+            tasks, needs to hold both entity and target table names.
+        entity_column: The name of the entity column.
+        target_column: The name of the target column.
+        time_column: The name of the time column, if it exists.
+    """
+    def __init__(
+        self,
+        task_type: TaskType,
+        context_df: pd.DataFrame,
+        pred_df: pd.DataFrame,
+        entity_table_name: str | Sequence[str],
+        entity_column: str,
+        target_column: str,
+        time_column: str | None = None,
+    ) -> None:
+
+        task_type = TaskType(task_type)
+        if task_type not in {  # Currently supported task types:
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+                TaskType.REGRESSION,
+                TaskType.TEMPORAL_LINK_PREDICTION,
+        }:
+            raise ValueError  # TODO
+        self._task_type = task_type
+
+        # TODO Check dfs (unify from local table)
+        self._context_df = context_df.copy(deep=False)
+        self._pred_df = pred_df.copy(deep=False)
+
+        self._dtype_dict: dict[str, Dtype] = {
+            column_name: infer_dtype(context_df[column_name])
+            for column_name in context_df.columns
+        }
+
+        self._entity_table_names: tuple[str] | tuple[str, str]
+        if isinstance(entity_table_name, str):
+            self._entity_table_names = (entity_table_name, )
+        elif len(entity_table_name) == 1:
+            self._entity_table_names = (entity_table_name[0], )
+        elif len(entity_table_name) == 2:
+            self._entity_table_names = (
+                entity_table_name[0],
+                entity_table_name[1],
+            )
+        else:
+            raise ValueError  # TODO
+
+        self._entity_column: str = ''
+        self._target_column: str = ''
+        self._time_column: str | None = None
+
+        self.entity_column = entity_column
+        self.target_column = target_column
+        if time_column is not None:
+            self.time_column = time_column
+
+    @property
+    def task_type(self) -> TaskType:
+        r"""The task type."""
+        return self._task_type
+
+    def narrow(self, start: int, length: int) -> Self:
+        r"""Returns a new :class:`TaskTable` that holds a narrowed version of
+        prediction examples.
+
+        Args:
+            start: Index of the prediction examples to start narrowing.
+            length: Length of the prediction examples.
+        """
+        out = copy.copy(self)
+        df = out._pred_df.iloc[start:start + length].reset_index(drop=True)
+        out._pred_df = df
+        return out
+
+    # Entity column ###########################################################
+
+    @property
+    def entity_table_name(self) -> str:
+        return self._entity_table_names[0]
+
+    @property
+    def entity_table_names(self) -> tuple[str] | tuple[str, str]:
+        return self._entity_table_names
+
+    @property
+    def entity_column(self) -> Column:
+        return Column(
+            name=self._entity_column,
+            expr=None,
+            dtype=self._dtype_dict[self._entity_column],
+            stype=Stype.ID,
+        )
+
+    @entity_column.setter
+    def entity_column(self, name: str) -> None:
+        if name in {self._target_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not Stype.ID.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._entity_column = name
+
+    # Target column ###########################################################
+
+    @property
+    def _target_stype(self) -> Stype:
+        if self.task_type in {
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+        }:
+            return Stype.categorical
+        if self.task_type in {TaskType.REGRESSION}:
+            return Stype.numerical
+        if self.task_type.is_link_pred:
+            return Stype.multicategorical
+        raise ValueError
+
+    @property
+    def target_column(self) -> Column:
+        return Column(
+            name=self._target_column,
+            expr=None,
+            dtype=self._dtype_dict[self._target_column],
+            stype=self._target_stype,
+        )
+
+    @target_column.setter
+    def target_column(self, name: str) -> None:
+        if name in {self._entity_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if not self._target_stype.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._target_column = name
+
+    # Time column #############################################################
+
+    def has_time_column(self) -> bool:
+        r"""Returns ``True`` if this task has a time column; ``False``
+        otherwise.
+        """
+        return self._time_column is not None
+
+    @property
+    def time_column(self) -> Column | None:
+        r"""The time column of this task.
+
+        The getter returns the time column of this task, or ``None`` if no
+        such time column is present.
+
+        The setter sets a column as a time column for this task, and raises a
+        :class:`ValueError` if the time column has a non-timestamp compatible
+        data type or if the column name does not match a column in the data
+        frame.
+        """
+        if self._time_column is None:
+            return None
+        return Column(
+            name=self._time_column,
+            expr=None,
+            dtype=self._dtype_dict[self._time_column],
+            stype=Stype.timestamp,
+        )
+
+    @time_column.setter
+    def time_column(self, name: str | None) -> None:
+        if name is None:
+            self._time_column = None
+            return
+
+        if name in {self._entity_column, self._target_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not contains_timestamp(
+                ser=self._context_df[name],
+                column_name=name,
+                dtype=self._dtype_dict[name],
+        ):
+            raise ValueError  # TODO
+
+        self._time_column = name
+
+    # Metadata ################################################################
+
+    @property
+    def metadata(self) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def print_metadata(self) -> None:
+        raise NotImplementedError
+
+    # Python builtins #########################################################
+
+    def __hash__(self) -> int:
+        return hash((
+            self.task_type,
+            self.entity_table_names,
+            self._entity_column,
+            self._target_column,
+            self._time_column,
+        ))
+
+    def __repr__(self) -> str:
+        if self.task_type.is_link_pred:
+            entity_table_repr = f'entity_table_names={self.entity_table_names}'
+        else:
+            entity_table_repr = f'entity_table_name={self.entity_table_name}'
+        return (f'{self.__class__.__name__}(\n'
+                f'  task_type={self.task_type},\n'
+                f'  num_context_examples={len(self._context_df)},\n'
+                f'  num_prediction_examples={len(self._pred_df)},\n'
+                f'  {entity_table_repr},\n'
+                f'  entity_column={self._entity_column},\n'
+                f'  target_column={self._target_column},\n'
+                f'  time_column={self._time_column},\n'
+                f')')

{kumoai-2.14.0.dev202512301731.dist-info → kumoai-2.14.0.dev202601041732.dist-info}/METADATA

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

{kumoai-2.14.0.dev202512301731.dist-info → kumoai-2.14.0.dev202601041732.dist-info}/RECORD

@@ -1,6 +1,6 @@
 kumoai/_logging.py,sha256=U2_5ROdyk92P4xO4H2WJV8EC7dr6YxmmnM-b7QX9M7I,886
 kumoai/mixin.py,sha256=MP413xzuCqWhxAPUHmloLA3j4ZyF1tEtfi516b_hOXQ,812
-kumoai/_version.py,sha256=zkmtgpHzS-8suGoRkSmHrktIFS142gX_ptBF0P9S3u4,39
+kumoai/_version.py,sha256=ZgEOgzhMS-JHEcx_StKbib3QyH7C44Vgs2zNQ7IM43A,39
 kumoai/kumolib.cpython-312-darwin.so,sha256=xQvdWHx9xmQ11y3F3ywxJv6A0sDk6D3-2fQbxSdM1z4,232576
 kumoai/__init__.py,sha256=x6Emn6VesHQz0wR7ZnbddPRYO9A5-0JTHDkzJ3Ocq6w,10907
 kumoai/formatting.py,sha256=jA_rLDCGKZI8WWCha-vtuLenVKTZvli99Tqpurz1H84,953
@@ -13,10 +13,11 @@ kumoai/_singleton.py,sha256=UTwrbDkoZSGB8ZelorvprPDDv9uZkUi1q_SrmsyngpQ,836
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kumoai/experimental/rfm/relbench.py,sha256=cVsxxV3TIL3PLEoYb-8tAVW3GSef6NQAd3rxdHJL63I,2276
 kumoai/experimental/rfm/graph.py,sha256=H9lIQLDkL5zJMwEHh7PgruvMUxWsjpynXUT7gnmTTUM,46351
-kumoai/experimental/rfm/__init__.py,sha256=TAy2TntkZdwB82wURsZasUsQ-yi06LEXT2u2qTNCVxc,6965
+kumoai/experimental/rfm/__init__.py,sha256=bW2XyYtkbdiu_iICYFF2Fu1Fx5fyGbqne6m_6c1P-fY,7016
 kumoai/experimental/rfm/sagemaker.py,sha256=6fyXO1Jd_scq-DH7kcv6JcV8QPyTbh4ceqwQDPADlZ0,4963
-kumoai/experimental/rfm/rfm.py,sha256=Qna-oSk5lgzmVC_KPolYo5Y6m81qKpyw9wfrvirT3Oc,49526
+kumoai/experimental/rfm/rfm.py,sha256=M5Q_vasSVrAX-5ucFPrdWoTZdpE4VYxKjKZz7BccITc,49672
 kumoai/experimental/rfm/authenticate.py,sha256=G2RkRWznMVQUzvhvbKhn0bMCY7VmoNYxluz3THRqSdE,18851
+kumoai/experimental/rfm/task_table.py,sha256=SPwkEdKRTwHHFcLqbvC1cDkyLXN2-3DpY5ujAyHRE-Q,8377
 kumoai/experimental/rfm/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kumoai/experimental/rfm/backend/sqlite/__init__.py,sha256=jl-DBbhsqQ-dUXyWhyQTM1AU2qNAtXCmi1mokdhtBTg,902
 kumoai/experimental/rfm/backend/sqlite/table.py,sha256=WqYtd_rwlawItRMXZUfv14qdyU6huQmODuFjDo483dI,6683
@@ -42,7 +43,7 @@ kumoai/experimental/rfm/infer/timestamp.py,sha256=vM9--7eStzaGG13Y-oLYlpNJyhL6f9
 kumoai/experimental/rfm/infer/stype.py,sha256=fu4zsOB-C7jNeMnq6dsK4bOZSewe7PtZe_AkohSRLoM,894
 kumoai/experimental/rfm/base/sql_sampler.py,sha256=qurkEVlMhDZw3d9SM2uGud6TMv_Wx_iqWoCgEKd_g9o,5094
 kumoai/experimental/rfm/base/__init__.py,sha256=rjmMux5lG8srw1bjQGcFQFv6zET9e5riP81nPkw28Jg,724
-kumoai/experimental/rfm/base/table.py,sha256=JWaSOcVYfGveUHFZpu85CUr4trLt1PJmAtgsz3QC8N8,26534
+kumoai/experimental/rfm/base/table.py,sha256=6qZeTMfnQejrn6TwqQeJGzJG7C0dSjJ7-NMLX38dvns,26563
 kumoai/experimental/rfm/base/sampler.py,sha256=tXYnVEyKC5NjSIpe8pNYp0V3Qbg-KbUE_QB0Emy2YiQ,30882
 kumoai/experimental/rfm/base/expression.py,sha256=Y7NtLTnKlx6euG_N3fLTcrFKheB6P5KS_jhCfoXV9DE,1252
 kumoai/experimental/rfm/base/source.py,sha256=bwu3GU2TvIXR2fwKAmJ1-5BDoNXMnI1SU3Fgdk8lWnc,301
@@ -87,7 +88,7 @@ kumoai/connector/bigquery_connector.py,sha256=IkyRqvF8Cg96kApUuuz86eYnl-BqBmDX1f
 kumoai/connector/source_table.py,sha256=QLT8bEYaxeMwy-b168url0VfnkTrs5K6VKLbxTI4hEY,17539
 kumoai/connector/__init__.py,sha256=9g6oNJ0qHWFlL5enTSoK4_SSH_5hP74xUDZx-9SggC4,842
 kumoai/connector/file_upload_connector.py,sha256=swp03HgChOvmNPJetuujBSAqADe7NRmS_T0F3o9it4w,7008
-kumoai/connector/utils.py,sha256=wlqQxMmPvnFNoCcczGkKYjSu05h8OhWh4fhTzQm_2bQ,64694
+kumoai/connector/utils.py,sha256=sD3_Dmf42FobMfVayzMVkDHIfXzPN-htD3RHd6Kw8hQ,65055
 kumoai/connector/s3_connector.py,sha256=3kbv-h7DwD8O260Q0h1GPm5wwQpLt-Tb3d_CBSaie44,10155
 kumoai/connector/base.py,sha256=cujXSZF3zAfuxNuEw54DSL1T7XCuR4t0shSMDuPUagQ,5291
 kumoai/pquery/__init__.py,sha256=uTXr7t1eXcVfM-ETaM_1ImfEqhrmaj8BjiIvy1YZTL8,533
@@ -114,8 +115,8 @@ kumoai/trainer/__init__.py,sha256=zUdFl-f-sBWmm2x8R-rdVzPBeU2FaMzUY5mkcgoTa1k,93
 kumoai/trainer/online_serving.py,sha256=9cddb5paeZaCgbUeceQdAOxysCtV5XP-KcsgFz_XR5w,9566
 kumoai/trainer/distilled_trainer.py,sha256=2pPs5clakNxkLfaak7uqPJOrpTWe1RVVM7ztDSqQZvU,6484
 kumoai/trainer/trainer.py,sha256=hBXO7gwpo3t59zKFTeIkK65B8QRmWCwO33sbDuEAPlY,20133
-kumoai-2.14.0.dev202512301731.dist-info/RECORD,,
-kumoai-2.14.0.dev202512301731.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
-kumoai-2.14.0.dev202512301731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.14.0.dev202512301731.dist-info/METADATA,sha256=XW8jzm0aptnoLAkWA04ZBBd_H9QnrcVQLUO5ZaF_HJk,2557
-kumoai-2.14.0.dev202512301731.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
+kumoai-2.14.0.dev202601041732.dist-info/RECORD,,
+kumoai-2.14.0.dev202601041732.dist-info/WHEEL,sha256=V1loQ6TpxABu1APUg0MoTRBOzSKT5xVc3skizX-ovCU,136
+kumoai-2.14.0.dev202601041732.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.14.0.dev202601041732.dist-info/METADATA,sha256=dmLN7vMtkp6iM92XWX0BtKZCF8yU3RmoAKhJ-gwc4ME,2557
+kumoai-2.14.0.dev202601041732.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102