chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/features/feature_field.py

@@ -4,11 +4,12 @@ import copy
 import dataclasses
 import functools
 import inspect
+import itertools
 import os
 import re
 import weakref
 from collections.abc import Mapping, MutableMapping
-from datetime import timedelta, datetime
+from datetime import datetime, timedelta
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -28,7 +29,6 @@ from typing import (
     cast,
 )
 
-import itertools
 import numpy as np
 import pyarrow as pa
 
@@ -39,11 +39,11 @@ from chalk.features._encoding.converter import FeatureConverter, JSONCodec, TDec
 from chalk.features._encoding.primitive import TPrimitive
 from chalk.features.feature_set import CURRENT_FEATURE_REGISTRY, FeatureRegistryProtocol
 from chalk.features.feature_wrapper import FeatureWrapper, NearestNeighborException
-from chalk.features.filter import Filter, TimeDelta, ClauseJoinWithAndException
+from chalk.features.filter import ClauseJoinWithAndException, Filter, TimeDelta
 from chalk.features.tag import Tags
 from chalk.features.underscore import Underscore
 from chalk.serialization.parsed_annotation import ParsedAnnotation
-from chalk.utils.collections import ensure_tuple, get_unique_item, FrozenOrderedSet, OrderedSet
+from chalk.utils.collections import FrozenOrderedSet, OrderedSet, ensure_tuple, get_unique_item
 from chalk.utils.duration import CHALK_MAX_TIMEDELTA, Duration, parse_chalk_duration
 from chalk.utils.import_utils import get_type_checking_imports
 from chalk.utils.json import JSON, pyarrow_json_type
@@ -54,7 +54,7 @@ from chalk.utils.pydanticutil.pydantic_compat import (
     is_pydantic_basemodel_instance,
     parse_pydantic_model,
 )
-from chalk.utils.string import to_snake_case, oxford_comma_list
+from chalk.utils.string import oxford_comma_list, to_snake_case
 
 if TYPE_CHECKING:
     from google.protobuf.message import Message as ProtobufMessage
@@ -68,9 +68,9 @@ _TPrim = TypeVar("_TPrim", bound=TPrimitive)
 _logger = get_logger(__name__)
 
 from chalk.features.feature_cache_strategy import (
-    CacheStrategy,
-    CacheNullsType,
     CacheDefaultsType,
+    CacheNullsType,
+    CacheStrategy,
     get_cache_strategy_from_cache_settings,
 )
 
@@ -252,6 +252,7 @@ class Feature(Generic[_TPrim, _TRich]):
         "tags",
         "underlying",
         "underscore_expression",
+        "offline_underscore_expression",
         "version",
         "window_duration",
         "window_durations",
@@ -276,6 +277,7 @@ class Feature(Generic[_TPrim, _TRich]):
         primary: bool | None = None,
         default: _TRich | ellipsis = ...,
         underscore_expression: Underscore | None = None,
+        offline_underscore_expression: Underscore | None = None,
         max_staleness: Duration | None | ellipsis = ...,
         online_store_max_items: int | None = None,
         cache_strategy: CacheStrategy = CacheStrategy.ALL_WITH_BOTH_UNSET,
@@ -386,6 +388,7 @@ class Feature(Generic[_TPrim, _TRich]):
         self._primary = primary
         self._primary_feature: Optional[Feature] = None
         self.underscore_expression: Underscore | None = underscore_expression
+        self.offline_underscore_expression: Underscore | None = offline_underscore_expression
         self.is_distance_pseudofeature = is_distance_pseudofeature
 
         self._raw_max_staleness = max_staleness
@@ -503,7 +506,7 @@ class Feature(Generic[_TPrim, _TRich]):
 
     @property
     def converter(self) -> FeatureConverter:
-        from chalk.features import DataFrame, Vector, Tensor
+        from chalk.features import DataFrame, Tensor, Vector
 
         self._converter_entered += 1
 
@@ -923,7 +926,7 @@ class Feature(Generic[_TPrim, _TRich]):
 
     def __repr__(self):
         try:
-            root_fqn=self.root_fqn
+            root_fqn = self.root_fqn
         except:
             # self.root_fqn is a property, if it failed then just return the object repr
             return object.__repr__(self)
@@ -1367,7 +1370,7 @@ class Feature(Generic[_TPrim, _TRich]):
                 ),
                 label="invalid nearest neighbor join",
                 range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                      or self.lsp_error_builder.property_range(self.attribute_name),
+                or self.lsp_error_builder.property_range(self.attribute_name),
                 code="32",
                 raise_error=TypeError,
             )
@@ -1387,7 +1390,7 @@ class Feature(Generic[_TPrim, _TRich]):
                     ),
                     label="invalid join",
                     range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                          or self.lsp_error_builder.property_range(self.attribute_name),
+                    or self.lsp_error_builder.property_range(self.attribute_name),
                     code="32",
                     raise_error=TypeError,
                 )
@@ -1400,7 +1403,7 @@ class Feature(Generic[_TPrim, _TRich]):
                     ),
                     label="invalid join",
                     range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                          or self.lsp_error_builder.property_range(self.attribute_name),
+                    or self.lsp_error_builder.property_range(self.attribute_name),
                     code="32",
                     raise_error=TypeError,
                 )
@@ -1424,18 +1427,21 @@ class Feature(Generic[_TPrim, _TRich]):
             )
 
         if not self.is_has_many and not self.is_has_one:
-            assert self.features_cls is not None
-            self.lsp_error_builder.add_diagnostic(
-                message=(
-                    f"The attribute '{self.features_cls.__name__}.{self.attribute_name}' "
-                    f"has a join filter ({join}) but its type annotation '{self.typ}' is not a feature class or DataFrame."
-                ),
-                label="invalid join",
-                range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                or self.lsp_error_builder.property_range(self.attribute_name),
-                code="37",
-                raise_error=TypeError,
-            )
+            # Check if user tried to use DataFrame (even if validation failed)
+            # Use is_dataframe_annotation() to detect DataFrame types without triggering validation errors
+            if not self.typ.is_dataframe_annotation():
+                assert self.features_cls is not None
+                self.lsp_error_builder.add_diagnostic(
+                    message=(
+                        f"The attribute '{self.features_cls.__name__}.{self.attribute_name}' "
+                        f"has a join filter ({join}) but its type annotation '{self.typ}' is not a feature class or DataFrame that links to another feature class."
+                    ),
+                    label="invalid join",
+                    range=self.lsp_error_builder.property_value_range(self.attribute_name)
+                    or self.lsp_error_builder.property_range(self.attribute_name),
+                    code="37",
+                    raise_error=TypeError,
+                )
         if self._join_type == "has_one":
             if self.is_has_many:
                 assert self.features_cls is not None
@@ -1748,6 +1754,7 @@ def feature(
     default: Union[_TRich, ellipsis] = ...,
     underscore: Optional[Underscore] = None,  # Deprecated. Prefer `expression`.
     expression: Optional[Underscore] = None,
+    offline_expression: Optional[Underscore] = None,
     offline_ttl: Optional[Union[ellipsis, Duration]] = ...,
     deprecated: bool = False,
     store_online: bool = True,
@@ -1811,6 +1818,8 @@ def feature(
         ...     total: int = feature(expression=_.subtotal + _.tax, default=0)
 
         See more at https://docs.chalk.ai/docs/expression
+    offline_expression
+        Defines an alternate expression to compute the feature during offline queries.
     dtype
         The backing `pyarrow.DataType` for the feature. This parameter can
         be used to control the storage format of data. For example, if you
@@ -2058,10 +2067,7 @@ def feature(
             if not isinstance(value, Feature):  # pyright: ignore[reportUnnecessaryIsInstance]
                 raise ValueError(f"When `versions` is provided, the values must be features, but `{value}` was given.")
 
-    cache_strategy = get_cache_strategy_from_cache_settings(
-        cache_nulls=cache_nulls,
-        cache_defaults=cache_defaults
-    )
+    cache_strategy = get_cache_strategy_from_cache_settings(cache_nulls=cache_nulls, cache_defaults=cache_defaults)
 
     return cast(
         _TRich,
@@ -2108,6 +2114,7 @@ def feature(
             ),
             default=default,
             underscore_expression=expression if expression is not None else underscore,
+            offline_underscore_expression=offline_expression,
             offline_ttl=offline_ttl,
             is_deprecated=deprecated,
             store_online=store_online,
@@ -2173,7 +2180,7 @@ def has_one(f: Callable[[], Any]) -> Any:
 def has_many(
     f: Callable[[], Any],
     max_staleness: Union[Duration, None, ellipsis] = ...,
-    online_store_max_items: int | None = None
+    online_store_max_items: int | None = None,
 ) -> Any:
     """Specify a feature that represents a one-to-many relationship.
 
@@ -2192,6 +2199,7 @@ def has_many(
         The maximum number of items to cache for the joined feature. The
         items in the joined feature aggregate, storing the latest values
         of the joined feature for each primary key in the joined feature.
+
     Examples
     --------
     >>> from chalk.features import DataFrame, features, has_many
@@ -2207,7 +2215,9 @@ def has_many(
     ...         lambda: User.id == Card.user_id
     ...     )
     """
-    return Feature(join=f, max_staleness=max_staleness,online_store_max_items=online_store_max_items, join_type="has_many")
+    return Feature(
+        join=f, max_staleness=max_staleness, online_store_max_items=online_store_max_items, join_type="has_many"
+    )
 
 
 __all__ = (

chalk/features/feature_set.py

@@ -19,11 +19,10 @@ from typing import (
     Set,
     Tuple,
     Type,
+    TypeGuard,
     cast,
 )
 
-from typing_extensions import TypeGuard
-
 from chalk._lsp._class_finder import get_class_ast
 from chalk.features.feature_wrapper import FeatureWrapper, unwrap_feature
 from chalk.utils import notebook

chalk/features/feature_set_decorator.py

@@ -141,6 +141,7 @@ def features(
         The `cache_nulls` and `cache_defaults` options can be used together on the same feature with the
         following exceptions: if `cache_nulls=False`, then `cache_defaults` cannot be `"evict_defaults"`, and if
         `cache_nulls="evict_defaults"`, then `cache_defaults` cannot be `False`.
+
     Other Parameters
     ----------------
     cls

chalk/features/feature_wrapper.py

@@ -11,6 +11,7 @@ from chalk.features._chalkop import op, Aggregation
 from chalk.features.filter import Filter
 from chalk.serialization.parsed_annotation import ParsedAnnotation
 from chalk.utils.collections import ensure_tuple
+from chalk.utils.notebook import is_notebook
 
 if TYPE_CHECKING:
     from chalk.features.feature_field import Feature
@@ -22,16 +23,43 @@ class NearestNeighborException(ValueError):
     ...
 
 
+class UnresolvedFeature:
+    """Fallback for features that can't be resolved in notebook environments.
+
+    This allows notebooks to work even when the feature registry is stale or incomplete.
+    The server will validate the feature exists when the query is executed.
+    """
+    __slots__ = ("fqn",)
+
+    def __init__(self, fqn: str):
+        self.fqn = fqn
+        super().__init__()
+
+    def __str__(self):
+        return self.fqn
+
+    def __repr__(self):
+        return f"UnresolvedFeature({self.fqn!r})"
+
+    def __hash__(self):
+        return hash(self.fqn)
+
+    def __eq__(self, other: object):
+        if isinstance(other, UnresolvedFeature):
+            return self.fqn == other.fqn
+        return False
+
+
 class _MarkedUnderlyingFeature:
     __slots__ = ("_fn", "_source", "_debug_info")
 
-    def __init__(self, fn: Callable[[], Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation],
+    def __init__(self, fn: Callable[[], Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation | UnresolvedFeature],
                  debug_info: Any = None) -> None:
         super().__init__()
         self._fn = fn
         self._debug_info = debug_info
 
-    def __call__(self, *args: Any, **kwds: Any) -> Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation:
+    def __call__(self, *args: Any, **kwds: Any) -> Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation | UnresolvedFeature:
         return self._fn()
 
 
@@ -51,7 +79,7 @@ class FeatureWrapper:
         super().__init__()
         self._chalk_underlying = underlying
 
-    def _chalk_get_underlying(self) -> Feature | Aggregation | Filter | type[DataFrame]:
+    def _chalk_get_underlying(self) -> Feature | Aggregation | Filter | type[DataFrame] | UnresolvedFeature:
         if isinstance(self._chalk_underlying, _MarkedUnderlyingFeature):
             self._chalk_underlying = self._chalk_underlying()
         if isinstance(self._chalk_underlying, FeatureWrapper):
@@ -303,6 +331,12 @@ class FeatureWrapper:
                     if f.attribute_name == item:
                         return FeatureWrapper(underlying.copy_with_path(f))
 
+                if is_notebook():
+                    # Construct FQN by preserving the path from the underlying feature
+                    # If underlying has a path, we need to include it in the FQN
+                    fqn = f"{underlying.root_fqn}.{item}"
+                    return UnresolvedFeature(fqn)
+
                 assert underlying.features_cls is not None
                 underlying.features_cls.__chalk_error_builder__.invalid_attribute(
                     root_feature_str=joined_class.namespace,
@@ -314,6 +348,11 @@ class FeatureWrapper:
                 )
                 assert False, "unreachable"
 
+            # If in notebook, fallback to constructing FQN string instead of raising error
+            if is_notebook():
+                fqn = f"{underlying.fqn}.{item}"
+                return UnresolvedFeature(fqn)
+
             assert underlying.features_cls is not None
             underlying.features_cls.__chalk_error_builder__.invalid_attribute(
                 root_feature_str=underlying.fqn,

chalk/features/hooks.py

@@ -4,14 +4,15 @@ import asyncio
 import inspect
 import logging
 import time  # Added for measuring duration
-from typing import Any, Callable, Iterable, Optional, Set, Tuple, Union, overload
-
-from typing_extensions import TypeAlias
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Set, Tuple, TypeAlias, Union, overload
 
 from chalk.features.tag import Environments
 from chalk.utils.collections import ensure_tuple
 from chalk.utils.log_with_context import get_logger
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 HookFn: TypeAlias = Callable[[], Any]
 
 
@@ -45,21 +46,31 @@ class Hook:
     venv: Optional[str]
     fn: HookFn
     filename: str
+    resource_hint: Optional["ResourceHint"]
+    resource_group: Optional[str]
 
     def __init__(
-        self, fn: HookFn, filename: str, environment: Optional[Environments] = None, venv: Optional[str] = None
+        self,
+        fn: HookFn,
+        filename: str,
+        environment: Optional[Environments] = None,
+        venv: Optional[str] = None,
+        resource_hint: Optional["ResourceHint"] = None,
+        resource_group: Optional[str] = None,
     ):
         super().__init__()
         self.fn = fn
         self.filename = filename
         self.environment = None if environment is None else ensure_tuple(environment)
         self.venv = venv
+        self.resource_hint = resource_hint
+        self.resource_group = resource_group
 
     def __call__(self):
         return self.fn()
 
     def __repr__(self):
-        return f'Hook(filename={self.filename}, fn={self.fn.__name__}", environment={str(self.environment)}, venv={self.venv})'
+        return f'Hook(filename={self.filename}, fn={self.fn.__name__}", environment={str(self.environment)}, venv={self.venv}, resource_hint={self.resource_hint}, resource_group={self.resource_group})'
 
     @classmethod
     async def async_run_all_before_all(cls, environment: str, venv: Optional[str] = None) -> None:
@@ -77,17 +88,46 @@ def before_all(fn: HookFn, /) -> Hook:
 
 @overload
 def before_all(
-    fn: None = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: HookFn,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
+) -> Hook:
+    ...
+
+
+@overload
+def before_all(
+    fn: None = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Callable[[HookFn], Hook]:
     ...
 
 
 def before_all(
-    fn: Optional[HookFn] = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: Optional[HookFn] = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Union[Hook, Callable[[HookFn], Hook]]:
     def decorator(f: HookFn):
         caller_filename = inspect.getsourcefile(f) or "unknown_file"
-        hook = Hook(fn=f, filename=caller_filename, environment=environment, venv=venv)
+        hook = Hook(
+            fn=f,
+            filename=caller_filename,
+            environment=environment,
+            venv=venv,
+            resource_hint=resource_hint,
+            resource_group=resource_group,
+        )
         Hook.before_all.add(hook)
         return hook
 
@@ -95,23 +135,52 @@ def before_all(
 
 
 @overload
-def after_all(fn: HookFn, /, environment: Optional[Environments] = None, venv: Optional[str] = None) -> Hook:
+def after_all(fn: HookFn, /) -> Hook:
+    ...
+
+
+@overload
+def after_all(
+    fn: HookFn,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
+) -> Hook:
     ...
 
 
 @overload
 def after_all(
-    fn: None = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: None = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Callable[[HookFn], Hook]:
     ...
 
 
 def after_all(
-    fn: Optional[HookFn] = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: Optional[HookFn] = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Union[Hook, Callable[[HookFn], Hook]]:
     def decorator(f: HookFn):
         caller_filename = inspect.getsourcefile(f) or "unknown_file"
-        hook = Hook(fn=f, filename=caller_filename, environment=environment, venv=venv)
+        hook = Hook(
+            fn=f,
+            filename=caller_filename,
+            environment=environment,
+            venv=venv,
+            resource_hint=resource_hint,
+            resource_group=resource_group,
+        )
         Hook.after_all.add(hook)
         return hook
 

chalk/features/inference.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Callable, Optional
 
 from chalk._lsp.error_builder import get_resolver_error_builder
 from chalk.features import DataFrame
@@ -11,8 +11,62 @@ from chalk.ml.model_version import ModelVersion
 from chalk.utils.collections import ensure_tuple
 
 
+def build_inference_function(
+    model_version: ModelVersion, pkey: Feature, output_features: Optional[Feature | list[Feature]] = None
+) -> Callable[[DataFrame], DataFrame]:
+    """Build the core inference function that takes a DataFrame and returns predictions.
+
+    Uses ModelInference.prepare_input() and extract_output() for model-specific logic.
+
+    Parameters
+    ----------
+    model_version
+        The model version to use for prediction
+    pkey
+        The primary key feature to exclude from predictions
+    output_features
+        Optional output feature(s) to add predictions to the DataFrame.
+        Can be a single Feature or a list of Features for multi-output models.
+
+    Returns
+    -------
+    Callable[[DataFrame], DataFrame]
+        Function that takes a DataFrame and returns predictions
+    """
+    # For all other models, use the ModelInference prepare_input/extract_output methods
+    pkey_string = str(pkey)
+
+    def fn(inp: DataFrame):
+        # Get features (excluding primary key) as PyArrow table
+        feature_table = inp[[c for c in inp.columns if c != pkey_string]].to_pyarrow()
+
+        # Use model-specific input preparation (default: __array__(), ONNX: struct array)
+        model_input = model_version.predictor.prepare_input(feature_table)
+
+        # Run prediction
+        result = model_version.predict(model_input)
+
+        if output_features is not None:
+            # Normalize to list for uniform processing
+            features_list = output_features if isinstance(output_features, list) else [output_features]
+
+            # Extract output for each feature and build columns dict
+            columns_dict = {}
+            for output_feature in features_list:
+                # Use model-specific output extraction (default: identity, ONNX: extract field)
+                output_feature_name = str(output_feature).split(".")[-1]
+                result_data = model_version.predictor.extract_output(result, output_feature_name)
+                columns_dict[output_feature] = result_data
+
+            return inp[pkey_string].with_columns(columns_dict)
+
+        return result
+
+    return fn
+
+
 def generate_inference_resolver(
-    inputs: list[Underscore], model_version: ModelVersion, resource_hint: Optional[ResourceHint] = None
+    inputs: list[Underscore] | Underscore, model_version: ModelVersion, resource_hint: Optional[ResourceHint] = None
 ) -> Feature:
     output_feature = Feature()
     previous_hook = output_feature.hook
@@ -22,20 +76,21 @@ def generate_inference_resolver(
             previous_hook(features)
 
         pkey = features.__chalk_primary__
-        pkey_string = str(features.__chalk_primary__)
+        if pkey is None:
+            raise ValueError(f"Feature class {features} does not have a primary key defined")
 
         def resolver_factory():
-            def fn(inp: DataFrame):
-                result = model_version.predict(inp[[c for c in inp.columns if c != pkey_string]].to_pyarrow())
-                return inp[pkey_string].with_columns({output_feature: result})
-
+            # Use the extracted build_inference_function
             cleaned_inputs = []
-            for i in inputs:
+            inputs_list = inputs if isinstance(inputs, list) else [inputs]
+            for i in inputs_list:
                 try:
                     cleaned_inputs.append(Feature.from_root_fqn(output_feature.namespace + str(i)[1:]))
                 except Exception as e:
                     raise ValueError(f"Could not find feature for input {i}: {e}")
 
+            fn = build_inference_function(model_version, pkey, output_feature)
+
             identifier = model_version.identifier or ""
             model_reference = MODEL_REFERENCE_REGISTRY.get((model_version.name, identifier), None)
             if model_reference is not None:
@@ -44,7 +99,7 @@ def generate_inference_resolver(
             return OnlineResolver(
                 function_definition="",
                 filename="",
-                fqn=f"{model_version.name}__{output_feature.namespace}.{output_feature.name}",
+                fqn=f"{model_version.name}__{output_feature.namespace}_{output_feature.name}",
                 doc=None,
                 inputs=[DataFrame[[pkey, *ensure_tuple(cleaned_inputs)]]],
                 state=None,
@@ -58,7 +113,7 @@ def generate_inference_resolver(
                 when=None,
                 tags=None,
                 owner=None,
-                resource_hint=resource_hint,
+                resource_hint=resource_hint or model_version.resource_hint,
                 data_sources=None,
                 is_sql_file_resolver=False,
                 source_line=None,