kumoai 2.13.0.dev202512040649__cp313-cp313-win_amd64.whl → 2.14.0.dev202601081732__cp313-cp313-win_amd64.whl

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

@@ -1,13 +1,22 @@
 import re
-import warnings
-from typing import List, Optional, Sequence
+from collections import Counter
+from collections.abc import Sequence
+from typing import cast
 
 import pandas as pd
+from kumoapi.model_plan import MissingType
 from kumoapi.typing import Dtype
 
 from kumoai.experimental.rfm.backend.sqlite import Connection
-from kumoai.experimental.rfm.base import SourceColumn, SourceForeignKey, Table
-from kumoai.experimental.rfm.infer import infer_dtype
+from kumoai.experimental.rfm.base import (
+    ColumnSpec,
+    ColumnSpecType,
+    DataBackend,
+    SourceColumn,
+    SourceForeignKey,
+    Table,
+)
+from kumoai.utils import quote_ident
 
 
 class SQLiteTable(Table):
@@ -16,6 +25,8 @@ class SQLiteTable(Table):
     Args:
         connection: The connection to a :class:`sqlite` database.
         name: The name of this table.
+        source_name: The source name of this table. If set to ``None``,
+            ``name`` is being used.
         columns: The selected columns of this table.
         primary_key: The name of the primary key of this table, if it exists.
         time_column: The name of the time column of this table, if it exists.
@@ -26,76 +37,148 @@ class SQLiteTable(Table):
         self,
         connection: Connection,
         name: str,
-        columns: Optional[Sequence[str]] = None,
-        primary_key: Optional[str] = None,
-        time_column: Optional[str] = None,
-        end_time_column: Optional[str] = None,
+        source_name: str | None = None,
+        columns: Sequence[ColumnSpecType] | None = None,
+        primary_key: MissingType | str | None = MissingType.VALUE,
+        time_column: str | None = None,
+        end_time_column: str | None = None,
     ) -> None:
 
         self._connection = connection
 
         super().__init__(
             name=name,
+            source_name=source_name,
             columns=columns,
             primary_key=primary_key,
             time_column=time_column,
             end_time_column=end_time_column,
         )
 
-    def _get_source_columns(self) -> List[SourceColumn]:
-        source_columns: List[SourceColumn] = []
+    @property
+    def backend(self) -> DataBackend:
+        return cast(DataBackend, DataBackend.SQLITE)
+
+    def _get_source_columns(self) -> list[SourceColumn]:
+        source_columns: list[SourceColumn] = []
         with self._connection.cursor() as cursor:
-            cursor.execute(f"PRAGMA table_info({self.name})")
-            rows = cursor.fetchall()
+            sql = f"PRAGMA table_info({self._quoted_source_name})"
+            cursor.execute(sql)
+            columns = cursor.fetchall()
+
+            if len(columns) == 0:
+                raise ValueError(f"Table '{self.source_name}' does not exist "
+                                 f"in the SQLite database")
+
+            unique_keys: set[str] = set()
+            sql = f"PRAGMA index_list({self._quoted_source_name})"
+            cursor.execute(sql)
+            for _, index_name, is_unique, *_ in cursor.fetchall():
+                if bool(is_unique):
+                    sql = f"PRAGMA index_info({quote_ident(index_name)})"
+                    cursor.execute(sql)
+                    index = cursor.fetchall()
+                    if len(index) == 1:
+                        unique_keys.add(index[0][2])
 
-            if len(rows) == 0:
-                raise ValueError(f"Table '{self.name}' does not exist")
-
-            for _, column, type, _, _, is_pkey in rows:
-                # Determine column affinity:
-                type = type.strip().upper()
-                if re.search('INT', type):
-                    dtype = Dtype.int
-                elif re.search('TEXT|CHAR|CLOB', type):
-                    dtype = Dtype.string
-                elif re.search('REAL|FLOA|DOUB', type):
-                    dtype = Dtype.float
-                else:  # NUMERIC affinity.
-                    ser = self._sample_df[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
+            # Special SQLite case that creates a rowid alias for
+            # `INTEGER PRIMARY KEY` annotated columns:
+            rowid_candidates = [
+                column for _, column, dtype, _, _, is_pkey in columns
+                if bool(is_pkey) and dtype.strip().upper() == 'INTEGER'
+            ]
+            if len(rowid_candidates) == 1:
+                unique_keys.add(rowid_candidates[0])
 
+            for _, column, dtype, notnull, _, is_pkey in columns:
                 source_column = SourceColumn(
                     name=column,
-                    dtype=dtype,
+                    dtype=self._to_dtype(dtype),
                     is_primary_key=bool(is_pkey),
-                    is_unique_key=False,
+                    is_unique_key=column in unique_keys,
+                    is_nullable=not bool(is_pkey) and not bool(notnull),
                 )
                 source_columns.append(source_column)
 
         return source_columns
 
-    def _get_source_foreign_keys(self) -> List[SourceForeignKey]:
-        source_fkeys: List[SourceForeignKey] = []
+    def _get_source_foreign_keys(self) -> list[SourceForeignKey]:
+        source_foreign_keys: list[SourceForeignKey] = []
         with self._connection.cursor() as cursor:
-            cursor.execute(f"PRAGMA foreign_key_list({self.name})")
-            for _, _, dst_table, fkey, pkey, _, _, _ in cursor.fetchall():
-                source_fkeys.append(SourceForeignKey(fkey, dst_table, pkey))
-        return source_fkeys
+            sql = f"PRAGMA foreign_key_list({self._quoted_source_name})"
+            cursor.execute(sql)
+            rows = cursor.fetchall()
+            counts = Counter(row[0] for row in rows)
+            for idx, _, dst_table, foreign_key, primary_key, *_ in rows:
+                if counts[idx] == 1:
+                    source_foreign_key = SourceForeignKey(
+                        name=foreign_key,
+                        dst_table=dst_table,
+                        primary_key=primary_key,
+                    )
+                    source_foreign_keys.append(source_foreign_key)
+        return source_foreign_keys
 
-    def _get_sample_df(self) -> pd.DataFrame:
+    def _get_source_sample_df(self) -> pd.DataFrame:
         with self._connection.cursor() as cursor:
-            cursor.execute(f"SELECT * FROM {self.name} "
-                           f"ORDER BY rowid LIMIT 1000")
+            columns = [quote_ident(col) for col in self._source_column_dict]
+            sql = (f"SELECT {', '.join(columns)} "
+                   f"FROM {self._quoted_source_name} "
+                   f"ORDER BY rowid "
+                   f"LIMIT {self._NUM_SAMPLE_ROWS}")
+            cursor.execute(sql)
             table = cursor.fetch_arrow_table()
-            return table.to_pandas(types_mapper=pd.ArrowDtype)
 
-    def _get_num_rows(self) -> Optional[int]:
+        if len(table) == 0:
+            raise RuntimeError(f"Table '{self.source_name}' is empty")
+
+        return self._sanitize(
+            df=table.to_pandas(types_mapper=pd.ArrowDtype),
+            dtype_dict={
+                column.name: column.dtype
+                for column in self._source_column_dict.values()
+            },
+            stype_dict=None,
+        )
+
+    def _get_num_rows(self) -> int | None:
         return None
+
+    def _get_expr_sample_df(
+        self,
+        columns: Sequence[ColumnSpec],
+    ) -> pd.DataFrame:
+        with self._connection.cursor() as cursor:
+            projections = [
+                f"{column.expr} AS {quote_ident(column.name)}"
+                for column in columns
+            ]
+            sql = (f"SELECT {', '.join(projections)} "
+                   f"FROM {self._quoted_source_name} "
+                   f"ORDER BY rowid "
+                   f"LIMIT {self._NUM_SAMPLE_ROWS}")
+            cursor.execute(sql)
+            table = cursor.fetch_arrow_table()
+
+        if len(table) == 0:
+            raise RuntimeError(f"Table '{self.source_name}' is empty")
+
+        return self._sanitize(
+            df=table.to_pandas(types_mapper=pd.ArrowDtype),
+            dtype_dict={column.name: column.dtype
+                        for column in columns},
+            stype_dict=None,
+        )
+
+    @staticmethod
+    def _to_dtype(dtype: str | None) -> Dtype | None:
+        if dtype is None:
+            return None
+        dtype = dtype.strip().upper()
+        if re.search('INT', dtype):
+            return Dtype.int
+        if re.search('TEXT|CHAR|CLOB', dtype):
+            return Dtype.string
+        if re.search('REAL|FLOA|DOUB', dtype):
+            return Dtype.float
+        return None  # NUMERIC affinity.

kumoai/experimental/rfm/base/__init__.py

@@ -1,10 +1,30 @@
-from .source import SourceColumn, SourceForeignKey
-from .column import Column
-from .table import Table
+from kumoapi.common import StrEnum
+
+
+class DataBackend(StrEnum):
+    LOCAL = 'local'
+    SQLITE = 'sqlite'
+    SNOWFLAKE = 'snowflake'
+
+
+from .source import SourceColumn, SourceForeignKey  # noqa: E402
+from .expression import Expression, LocalExpression  # noqa: E402
+from .column import ColumnSpec, ColumnSpecType, Column  # noqa: E402
+from .table import Table  # noqa: E402
+from .sampler import SamplerOutput, Sampler  # noqa: E402
+from .sql_sampler import SQLSampler  # noqa: E402
 
 __all__ = [
+    'DataBackend',
     'SourceColumn',
     'SourceForeignKey',
+    'Expression',
+    'LocalExpression',
+    'ColumnSpec',
+    'ColumnSpecType',
     'Column',
     'Table',
+    'SamplerOutput',
+    'Sampler',
+    'SQLSampler',
 ]

kumoai/experimental/rfm/base/column.py

@@ -1,37 +1,119 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Mapping, TypeAlias
 
 from kumoapi.typing import Dtype, Stype
+from typing_extensions import Self
+
+from kumoai.experimental.rfm.base import Expression
+from kumoai.mixin import CastMixin
+
+
+@dataclass(init=False)
+class ColumnSpec(CastMixin):
+    r"""A column specification for adding a column to a table.
+
+    A column specification can either refer to a physical column present in
+    the data source, or be defined logically via an expression.
+
+    Args:
+        name: The name of the column.
+        expr: A column expression to define logical columns.
+        dtype: The data type of the column.
+    """
+    def __init__(
+        self,
+        name: str,
+        expr: Expression | Mapping[str, str] | str | None = None,
+        dtype: Dtype | str | None = None,
+        stype: Stype | str | None = None,
+    ) -> None:
+
+        self.name = name
+        self.expr = Expression.coerce(expr)
+        self.dtype = Dtype(dtype) if dtype is not None else None
+        self.stype = Stype(dtype) if stype is not None else None
+
+    @classmethod
+    def coerce(cls, spec: ColumnSpec | Mapping[str, Any] | str) -> Self:
+        r"""Coerces a column specification into a :class:`ColumnSpec`."""
+        if isinstance(spec, cls):
+            return spec
+        if isinstance(spec, str):
+            return cls(name=spec)
+        if isinstance(spec, Mapping):
+            try:
+                return cls(**spec)
+            except TypeError:
+                pass
+        raise TypeError(f"Unable to coerce 'ColumnSpec' from '{spec}'")
+
+    @property
+    def is_source(self) -> bool:
+        r"""Whether the column specification refers to a phyiscal column
+        present in the data source.
+        """
+        return self.expr is None
+
+
+ColumnSpecType: TypeAlias = ColumnSpec | Mapping[str, Any] | str
 
 
 @dataclass(init=False, repr=False, eq=False)
 class Column:
+    r"""Column-level metadata information.
+
+    A column can either refer to a physical column present in the data source,
+    or be defined logically via an expression.
+
+    Args:
+        name: The name of the column.
+        expr: A column expression to define logical columns.
+        dtype: The data type of the column.
+        stype: The semantic type of the column.
+    """
     stype: Stype
 
     def __init__(
         self,
         name: str,
+        expr: Expression | None,
         dtype: Dtype,
         stype: Stype,
-        is_primary_key: bool = False,
-        is_time_column: bool = False,
-        is_end_time_column: bool = False,
     ) -> None:
         self._name = name
+        self._expr = expr
         self._dtype = Dtype(dtype)
-        self._is_primary_key = is_primary_key
-        self._is_time_column = is_time_column
-        self._is_end_time_column = is_end_time_column
+
+        self._is_primary_key = False
+        self._is_time_column = False
+        self._is_end_time_column = False
+
         self.stype = Stype(stype)
 
     @property
     def name(self) -> str:
+        r"""The name of the column."""
         return self._name
 
+    @property
+    def expr(self) -> Expression | None:
+        r"""The expression of column (if logically)."""
+        return self._expr
+
     @property
     def dtype(self) -> Dtype:
+        r"""The data type of the column."""
         return self._dtype
 
+    @property
+    def is_source(self) -> bool:
+        r"""Whether the column refers to a phyiscal column present in the data
+        source.
+        """
+        return self.expr is None
+
     def __setattr__(self, key: str, val: Any) -> None:
         if key == 'stype':
             if isinstance(val, str):
@@ -54,7 +136,7 @@ class Column:
         super().__setattr__(key, val)
 
     def __hash__(self) -> int:
-        return hash((self.name, self.stype, self.dtype))
+        return hash((self.name, self.expr, self.dtype, self.stype))
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, Column):
@@ -62,5 +144,9 @@ class Column:
         return hash(self) == hash(other)
 
     def __repr__(self) -> str:
-        return (f'{self.__class__.__name__}(name={self.name}, '
-                f'stype={self.stype}, dtype={self.dtype})')
+        parts = [f'name={self.name}']
+        if self.expr is not None:
+            parts.append(f'expr={self.expr}')
+        parts.append(f'dtype={self.dtype}')
+        parts.append(f'stype={self.stype}')
+        return f"{self.__class__.__name__}({', '.join(parts)})"

kumoai/experimental/rfm/base/expression.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass
+from typing import Mapping
+
+
+class Expression(ABC):
+    """A base expression to define logical columns."""
+    @classmethod
+    def coerce(
+        cls,
+        spec: Expression | Mapping[str, str] | str | None,
+    ) -> Expression | None:
+        r"""Coerces an expression specification into an :class:`Expression`, if
+        possible.
+        """
+        if spec is None:
+            return None
+        if isinstance(spec, Expression):
+            return spec
+        if isinstance(spec, str):
+            return LocalExpression(spec)
+        if isinstance(spec, Mapping):
+            for sub_cls in (LocalExpression, ):
+                try:
+                    return sub_cls(**spec)
+                except TypeError:
+                    pass
+        raise TypeError(f"Unable to coerce 'Expression' from '{spec}'")
+
+
+@dataclass(frozen=True, repr=False)
+class LocalExpression(Expression):
+    r"""A local expression to define a row-level logical attribute based on
+    physical columns of the data source in the same row.
+
+    Args:
+        value: The value of the expression.
+    """
+    value: str
+
+    def __repr__(self) -> str:
+        return self.value