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

chalk/features/resolver.py

@@ -21,6 +21,7 @@ import random
 import re
 import statistics
 import types
+import typing
 from dataclasses import dataclass, is_dataclass
 from datetime import datetime
 from enum import Enum, IntEnum
@@ -43,13 +44,18 @@ from typing import (
     Literal,
     Mapping,
     Optional,
+    ParamSpec,
     Protocol,
     Sequence,
     Set,
     Type,
+    TypeAlias,
     TypeVar,
     Union,
     cast,
+    final,
+    get_args,
+    get_origin,
     overload,
 )
 
@@ -61,10 +67,10 @@ from google.protobuf import message_factory
 from google.protobuf.descriptor import Descriptor
 from google.protobuf.internal.python_message import GeneratedProtocolMessageType
 from pydantic import BaseModel
-from typing_extensions import ParamSpec, TypeAlias, final, get_args, get_origin
 
 from chalk._lsp._class_finder import get_function_caller_info
 from chalk._lsp.error_builder import FunctionCallErrorBuilder, ResolverErrorBuilder, get_resolver_error_builder
+from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
 from chalk.features._encoding.protobuf import (
     convert_proto_message_type_to_pyarrow_type,
     serialize_message_file_descriptor,
@@ -109,6 +115,7 @@ if TYPE_CHECKING:
 
     from chalk.features import Underscore
     from chalk.features.underscore import UnderscoreAttr, UnderscoreCall, UnderscoreCast, UnderscoreFunction
+    from chalk.ml.model_version import ModelVersion
     from chalk.sql import BaseSQLSourceProtocol, SQLSourceGroup
     from chalk.sql._internal.sql_settings import SQLResolverSettings
     from chalk.sql._internal.sql_source import BaseSQLSource
@@ -550,28 +557,18 @@ class ResolverRegistry:
         short_name = resolver.name
         if short_name in self._short_name_to_resolver:
             if not override and not notebook.is_notebook():
-                if resolver.fqn == self._short_name_to_resolver[short_name]:
-                    # Same resolver was redefined
-                    resolver.lsp_builder.add_diagnostic(
-                        message=f"Duplicate resolver '{resolver.fqn}'. Multiple resolvers cannot have the same name.",
-                        code="71",
-                        label="duplicate name",
-                        range=resolver.lsp_builder.function_name(),
-                        raise_error=ValueError,
-                    )
-                else:
-                    # Same short name was reused
-                    resolver.lsp_builder.add_diagnostic(
-                        message=(
-                            f"Another resolver with the same function name '{resolver.name}' in module "
-                            f"'{self._short_name_to_resolver[short_name].__module__}' exists. "
-                            f"Resolver function names must be unique. Please rename this resolver in module '{resolver.__module__}'."
-                        ),
-                        label="duplicate resolver shortname",
-                        code="71",
-                        range=resolver.lsp_builder.function_name(),
-                        raise_error=None,
-                    )
+                # Same short name was reused
+                resolver.lsp_builder.add_diagnostic(
+                    message=(
+                        f"Another resolver with the same function name '{resolver.name}' in module "
+                        f"'{self._short_name_to_resolver[short_name].__module__}' exists. "
+                        f"Resolver function names must be unique. Please rename this resolver in module '{resolver.__module__}'."
+                    ),
+                    label="duplicate resolver shortname",
+                    code="71",
+                    range=resolver.lsp_builder.function_name(),
+                    raise_error=None,
+                )
                 return
             existing_resolver = self._short_name_to_resolver[short_name]
             # Need to remove the resolver from the typed registry
@@ -644,6 +641,7 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
         output_row_order: Literal["one-to-one"] | None = None,
         venv: str | None = None,
         name: None = None,  # deprecated
+        postprocessing: Underscore | None = None,
     ):
         self._function_definition = ... if function_definition is None else function_definition
         self._function_captured_globals = ... if function_captured_globals is None else function_captured_globals
@@ -687,6 +685,7 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
         self._data_lineage = data_lineage
         self._sql_settings = sql_settings
         self.output_row_order = output_row_order
+        self.postprocessing = postprocessing
         super().__init__()
 
     @property
@@ -844,14 +843,20 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
                 annotation = None
 
             if annotation is not None:
-                if not isinstance(val, DataFrame):
+                if self.static and type(val).__name__ == "DataFrame" and type(val).__module__ == "chalkdf.dataframe":
+                    # No need to wrap this class in DataFrame.
+                    pass
+                elif self.static and isinstance(val, LazyFramePlaceholder):
+                    # No need to wrap this class in DataFrame.
+                    pass
+                elif not isinstance(val, DataFrame):
                     val = DataFrame(val)
 
                 if time_is_frozen():
                     frozen_filter = Filter(lhs=CHALK_TS_FEATURE, operation="<=", rhs=TimeDelta(hours_ago=0))
                     annotation.filters = (frozen_filter, *annotation.filters)
 
-                if annotation.filters and len(annotation.filters) > 0:
+                if annotation.filters and len(annotation.filters) > 0 and not isinstance(val, LazyFramePlaceholder):
                     try:
                         val = val[annotation.filters]
                         val._materialize()  # pyright: ignore[reportPrivateUsage]
@@ -2824,6 +2829,8 @@ class StreamResolver(Resolver[P, T]):
         sql_settings: SQLResolverSettings | None,
         feature_expressions: dict[Feature, Underscore] | None,
         message_producer_parsed: StreamResolverMessageProducerParsed | None,
+        skip_online: bool = False,
+        skip_offline: bool = False,
     ):
         super().__init__(
             function_definition=function_definition,
@@ -2892,6 +2899,8 @@ class StreamResolver(Resolver[P, T]):
 
         self.feature_expressions: dict[Feature, Underscore] | None = feature_expressions
         self.message_producer_parsed: StreamResolverMessageProducerParsed | None = message_producer_parsed
+        self.skip_online = skip_online
+        self.skip_offline = skip_offline
 
     @property
     def output_features(self) -> Sequence[Feature]:
@@ -3175,6 +3184,18 @@ def _parse_stream_resolver_output_features(
     return output_features
 
 
+def _is_valid_stream_message_type(typ: Type) -> bool:
+    from chalk.functions.proto import _is_protobuf_message
+
+    if is_pydantic_basemodel(typ):
+        return True
+    if is_dataclass(typ):
+        return True
+    if _is_protobuf_message(typ):
+        return True
+    return False
+
+
 @dataclass(frozen=True)
 class ParseInfo(Generic[T, V]):
     fn: Callable[[T], V]
@@ -3186,18 +3207,16 @@ class ParseInfo(Generic[T, V]):
 
 
 def _validate_parse_function(
-    stream_fn: Callable[P, T],
+    stream_fqn: str,
     parse_fn: Callable[[T], Any],
     globals: dict[str, Any] | None,
     locals: dict[str, Any] | None,
-    params: Sequence[StreamResolverParam],
-    resolver_error_builder: ResolverErrorBuilder,
+    stream_fn_input_type: Type[Any],
     name: str | None,
 ) -> ParseInfo:
     parse_error_builder = get_resolver_error_builder(parse_fn)
     """We need separate error builders for resolver and parse fn: different AST nodes"""
 
-    stream_fqn = get_resolver_fqn(function=stream_fn, name=name)
     parse_fqn = get_resolver_fqn(function=parse_fn, name=name)
     sig = inspect.signature(parse_fn)
     annotation_parser = ResolverAnnotationParser(parse_fn, globals, locals, parse_error_builder)
@@ -3237,18 +3256,14 @@ def _validate_parse_function(
         raise
     else:
         parse_output = return_annotation
-    if not is_pydantic_basemodel(parse_output):
+    if not _is_valid_stream_message_type(parse_output):
         parse_error_builder.add_diagnostic(
-            message=f"Parse function '{parse_fqn}' return annotation must be of type pydantic.BaseModel",
+            message=f"Parse function '{parse_fqn}' return annotation must be either a pydantic BaseModel, decorated with @dataclass, or a protobuf Message type",
             code="101",
             label="invalid parse function return annotation",
             range=parse_error_builder.function_return_annotation(),
             raise_error=TypeError,
         )
-    stream_fn_input_arg = next(
-        param for param in params if isinstance(param, (StreamResolverParamMessage, StreamResolverParamMessageWindow))
-    )
-    stream_fn_input_type = _get_stream_resolver_input_type(stream_fn_input_arg, stream_fqn, resolver_error_builder)
     if parse_output != stream_fn_input_type:
         parse_error_builder.add_diagnostic(
             message=(
@@ -3655,13 +3670,19 @@ def parse_and_register_stream_resolver(
     )
     parse_info = None
     if parse:
+        stream_fqn = get_resolver_fqn(function=fn, name=name)
+        stream_fn_input_arg = next(
+            param
+            for param in params
+            if isinstance(param, (StreamResolverParamMessage, StreamResolverParamMessageWindow))
+        )
+        stream_fn_input_type = _get_stream_resolver_input_type(stream_fn_input_arg, stream_fqn, error_builder)
         parse_info = _validate_parse_function(
-            stream_fn=fn,
+            stream_fqn=stream_fqn,
             parse_fn=parse,
             globals=caller_globals,
             locals=caller_locals,
-            params=params,
-            resolver_error_builder=error_builder,
+            stream_fn_input_type=stream_fn_input_type,
             name=name,
         )
     if keys is not None:
@@ -3744,18 +3765,55 @@ def parse_and_register_stream_resolver(
     return resolver
 
 
+def _validate_message_type(message_type: Type[Any], allow_lists: bool = True) -> str | None:
+    if hasattr(message_type, "__origin__"):
+        assert hasattr(message_type, "__args__")
+        if message_type.__origin__ not in (list, typing.List, typing.Sequence):
+            return "The only generic type supported is list|List|Sequence."
+        if len(message_type.__args__) != 1:
+            return f"Found {len(message_type.__args__)} type parameters for generic type, only one supported."
+        if not allow_lists:
+            return f"Nested lists (e.g. List[List[MessageType]]) not supported."
+        sub_type = message_type.__args__[0]
+        sub_res = _validate_message_type(sub_type, allow_lists=False)
+        if sub_res is not None:
+            sub_res = f"Found type List[T] with invalid T: {sub_res}"
+        return sub_res
+
+    if message_type in (str, bytes):
+        return None
+
+    if inspect.isclass(message_type):  # pyright: ignore[reportUnnecessaryIsInstance]
+        if issubclass(message_type, BaseModel):
+            return None
+        elif issubclass(message_type, google.protobuf.message.Message):
+            return None
+        elif is_dataclass(message_type):
+            return None
+        else:
+            return "Unsupported type (expected str/bytes, a struct type, or a list[struct])"
+    else:
+        return "message type should be a type"
+
+
+def _is_list_message_type(message_type: Type[Any]):
+    return getattr(message_type, "__origin__", None) in (list, typing.List, typing.Sequence)
+
+
 def make_stream_resolver(
     *,
     name: str,
     source: StreamSource,
     message_type: Type[BaseModel | google.protobuf.message.Message | AnyDataclass | str | bytes],
     output_features: "Mapping[FeatureWrapper, Underscore]",
-    parse: Optional[Underscore] = None,
+    parse: Underscore | Callable[[str | bytes], Any] | None = None,
     environment: Optional[Environments] = None,
     machine_type: Optional[MachineType] = None,
     owner: Optional[str] = None,
     doc: str | None = None,
     sink: Sink | None = None,
+    skip_online: bool = False,
+    skip_offline: bool = False,
 ) -> StreamResolver:
     """Constructs a streaming resolver that, instead of a Python function,
     defines its output features as column projections on an input message.
@@ -3788,12 +3846,21 @@ def make_stream_resolver(
     sink
         An optional message producer configuration that specifies where to send messages.
         Read more at https://docs.chalk.ai/api-docs#Sink
+    skip_online
+        If True, skip online persistence (no writes to Redis/DynamoDB/etc).
+        Results will still be processed but not stored in online stores.
+        Note: Only applies to native streaming. Default: False
+    skip_offline
+        If True, skip offline persistence (no result bus publishing for offline storage).
+        Results will still be processed but not stored in offline stores (S3/BigQuery/etc).
+        Note: Only applies to native streaming. Default: False
 
     Returns
     -------
     StreamResolver
         A configured stream resolver.
     """
+    from chalk.features.underscore import Underscore
 
     # The function "definition" will be the source code of the invocation, for error reporting / LSP highlighting.
     caller_info = get_function_caller_info(frame_offset=1)
@@ -3802,6 +3869,15 @@ def make_stream_resolver(
     caller_lineno = caller_info.lineno
     error_builder = FunctionCallErrorBuilder(caller_info)
 
+    # TODO unify this with the above
+    frame = inspect.currentframe()
+    assert frame is not None
+    caller_frame = frame.f_back
+    assert caller_frame is not None
+    caller_globals = caller_frame.f_globals
+    caller_locals = caller_frame.f_locals
+    del frame
+
     if not isinstance(source, StreamSource):  # pyright: ignore[reportUnnecessaryIsInstance]
         error_builder.add_diagnostic(
             message=(
@@ -3841,27 +3917,26 @@ def make_stream_resolver(
             )
 
     # Validate message_type is one of the allowed types
-    valid_message_type = False
-    if inspect.isclass(message_type):  # pyright: ignore[reportUnnecessaryIsInstance]
-        if issubclass(message_type, BaseModel):
-            valid_message_type = True
-        elif issubclass(message_type, google.protobuf.message.Message):
-            valid_message_type = True
-        elif is_dataclass(message_type):
-            valid_message_type = True
-
-    # Also allow str and bytes types
-    if message_type in (str, bytes):
-        valid_message_type = True
-
-    if not valid_message_type:  # pyright: ignore[reportUnnecessaryIsInstance]
+    message_type_validation_msg = _validate_message_type(message_type)
+    if message_type_validation_msg is not None:
         error_builder.add_diagnostic(
-            message=f"Invalid message_type for stream resolver '{name}': expected Pydantic BaseModel, protobuf Message, dataclass, str, or bytes, got {type(message_type).__name__}",
+            message=f"Invalid message_type for stream resolver '{name}' ({message_type}): {message_type_validation_msg}. Supported message types includes 'str', 'bytes', some struct type (a pydantic BaseModel or dataclass), or a list[T] where T is a string/bytes/struct)",
             code="195",
             label="Invalid message type",
             range=error_builder.function_arg_range_by_name("message_type"),
         )
 
+    if _is_list_message_type(message_type) and parse is None:
+        error_builder.add_diagnostic(
+            message=(
+                f"Found list message_type without a parse function for stream resolver '{name}' ({message_type}): List message types are only supported if a custom parse function is provided. "
+                f"Otherwise, only struct or bytes/string messages are supported."
+            ),
+            code="196",
+            label="List message type without parse function",
+            range=error_builder.function_arg_range_by_name("message_type"),
+        )
+
     from chalk import Features
 
     unwrapped_features: list[Feature] = []
@@ -3891,8 +3966,10 @@ def make_stream_resolver(
             f"Stream resolver '{name}' has expression-based parse function so it can't be called directly."
         )
 
+    params = [StreamResolverParamMessage(typ=message_type, name="message")]
+
     parse_info: Optional[ParseInfo] = None
-    if parse is not None:
+    if isinstance(parse, Underscore):
         parse_info = ParseInfo(
             fn=_dummy_parse_fn,
             input_type=bytes,
@@ -3901,6 +3978,19 @@ def make_stream_resolver(
             parse_function_captured_globals=None,
             parse_expression=parse,
         )
+    elif callable(parse):
+        parse_info = _validate_parse_function(
+            name,
+            parse_fn=parse,
+            globals=caller_globals,
+            locals=caller_locals,
+            stream_fn_input_type=message_type,
+            name=name,
+        )
+        if parse_info.input_type != bytes:
+            raise ValueError(
+                f"Native streaming resolvers only support python parse functions with input bytes 'bytes'. Function {parse} has input type {parse_info.input_type}"
+            )
 
     # Validate and parse sink before creating StreamResolver
     message_producer_parsed: StreamResolverMessageProducerParsed | None = None
@@ -3925,7 +4015,7 @@ def make_stream_resolver(
         message=message_type,
         output=output_type,
         signature=StreamResolverSignature(
-            params=[StreamResolverParamMessage(typ=message_type, name="message")],
+            params=params,
             output_feature_fqns={str(x) for x in output_features.keys()},
         ),
         state=None,
@@ -3941,6 +4031,8 @@ def make_stream_resolver(
         sql_settings=None,
         feature_expressions={unwrap_feature(x): u for x, u in output_features.items()},
         message_producer_parsed=message_producer_parsed,
+        skip_online=skip_online,
+        skip_offline=skip_offline,
     )
     resolver.add_to_registry(override=False)
     return resolver
@@ -4146,6 +4238,10 @@ def is_structured_type(
     if not inspect.isclass(message_type):  # pyright: ignore[reportUnnecessaryIsInstance]
         return False
 
+    if hasattr(message_type, "__origin__"):
+        # It's a generic type like List[T]
+        return False
+
     # Check if it's a Pydantic BaseModel
     if issubclass(message_type, BaseModel):
         return True
@@ -4227,6 +4323,9 @@ def validate_field_chain(
 
     # Base case: if parent is UnderscoreRoot (_), validate field against current_type
     if isinstance(underscore_attr._chalk__parent, UnderscoreRoot):
+        if underscore_attr._chalk__attr == "chalk_now":
+            return datetime
+
         # Check if current_type allows field access
         if current_type in (str, bytes):
             error_builder.add_diagnostic(
@@ -4392,5 +4491,188 @@ def validate_message_attributes(
     name: str,
 ) -> None:
     """Validate that all underscore expressions use valid field names for the message_type."""
+    if _is_list_message_type(message_type):
+        message_type = message_type.__args__[0]  # pyright: ignore[reportAttributeAccessIssue]
     for expression in expressions:
         validate_underscore_expression(expression, message_type, error_builder, name)
+
+
+def make_model_resolver(
+    name: str,
+    model: "ModelVersion",
+    inputs: Dict[Feature, str] | List[Feature],
+    output: Feature | List[Feature] | Dict[Feature, str],
+    feature_class: Optional[type[Features]] = None,
+    resource_group: Optional[str] = None,
+    resource_hint: Optional[ResourceHint] = None,
+) -> OnlineResolver:
+    """
+    Create an online resolver that runs inference on a model.
+
+    This function provides an imperative API for creating model inference resolvers,
+    as an alternative to using F.inference in feature definitions. It uses the same
+    underlying implementation as F.inference but allows you to create resolvers
+    programmatically.
+
+    Parameters
+    ----------
+    name
+        The name of the resolver
+    model
+        A ModelVersion reference to a deployed model
+    inputs
+        Either a dict mapping Feature objects to model input names (strings), or a list of
+        Feature objects. If a dict, the values represent the model's expected input names
+        (for future use). If a list, the features will be passed as a single DataFrame to
+        the model.
+    output
+        The output feature(s) that will contain the predictions.
+        Can be a single Feature, a list of Features, or a dict mapping Feature objects to
+        model output names (strings) for future use with multi-output models.
+    feature_class
+        Optional feature class to use. If not provided, will be inferred from the inputs.
+    resource_group
+        Optional resource group for the resolver
+    resource_hint
+        Optional resource hint for execution (e.g., CPU/GPU preferences)
+
+    Returns
+    -------
+    OnlineResolver
+        The created resolver
+
+    Examples
+    --------
+    >>> from chalk.features import features, feature
+    >>> from chalk.features.resolver import make_model_resolver
+    >>> from chalk.ml import ModelVersion
+    >>>
+    >>> @features
+    ... class User:
+    ...     id: str = feature(primary=True)
+    ...     age: float
+    ...     income: float
+    ...     risk_score: float
+    ...     credit_score: float
+    >>>
+    >>> # Create a model version reference
+    >>> model = ModelVersion(
+    ...     name="risk_model",
+    ...     version=1,
+    ...     model_type="sklearn",
+    ...     model_encoding="pickle",
+    ...     filename="model.pkl"
+    ... )
+    >>>
+    >>> # Create resolver with single output
+    >>> resolver = make_model_resolver(
+    ...     name="risk_model",
+    ...     model=model,
+    ...     inputs=[User.age, User.income],
+    ...     output=User.risk_score,
+    ... )
+    >>>
+    >>> # Create resolver with multiple outputs (list)
+    >>> resolver = make_model_resolver(
+    ...     name="multi_output_model",
+    ...     model=model,
+    ...     inputs=[User.age, User.income],
+    ...     output=[User.risk_score, User.credit_score],
+    ... )
+    >>>
+    >>> # Create resolver with named inputs and outputs (dict)
+    >>> resolver = make_model_resolver(
+    ...     name="named_model",
+    ...     model=model,
+    ...     inputs={User.age: "age_input", User.income: "income_input"},
+    ...     output={User.risk_score: "risk_output", User.credit_score: "credit_output"},
+    ... )
+    """
+    from chalk.features.inference import build_inference_function
+
+    if isinstance(inputs, dict):
+        input_features_raw = list(inputs.keys())
+    else:
+        input_features_raw = inputs
+
+    input_features = [unwrap_feature(f) for f in input_features_raw]
+
+    if isinstance(output, dict):
+        output_features = [unwrap_feature(f) for f in output.keys()]
+    elif isinstance(output, list):
+        output_features = [unwrap_feature(f) for f in output]
+    else:
+        output_features = [unwrap_feature(output)]
+
+    # If feature_class is not provided, try to infer it from the first input feature
+    if feature_class is None:
+        if not input_features:
+            raise ValueError("Cannot infer feature class: no input features provided and feature_class not specified")
+
+        first_input = input_features[0]
+
+        if hasattr(first_input, "features_cls") and first_input.features_cls is not None:
+            feature_class = first_input.features_cls
+        else:
+            raise ValueError(
+                "Cannot infer feature class from inputs. Please provide feature_class parameter explicitly."
+            )
+
+    pkey = feature_class.__chalk_primary__
+    if pkey is None:
+        raise ValueError(f"Feature class {feature_class} does not have a primary key defined")
+
+    first_output = output_features[0]
+
+    output_namespace = (
+        first_output.namespace
+        if hasattr(first_output, "namespace") and first_output.namespace
+        else feature_class.__name__.lower()
+    )
+
+    # Use the same underlying inference function as F.inference
+    # Pass list of outputs if multiple, single if only one
+    output_for_inference = output_features if len(output_features) > 1 else output_features[0]
+    inference_fn = build_inference_function(model, pkey, output_for_inference)
+
+    if len(output_features) == 1:
+        output_names = output_features[0].name
+    else:
+        output_names = "_".join(f.name for f in output_features)
+
+    resolver = OnlineResolver(
+        function_definition="",
+        filename="",
+        fqn=f"{name}__{output_namespace}_{output_names}",
+        doc=None,
+        inputs=[DataFrame[[pkey, *ensure_tuple(input_features)]]],
+        state=None,
+        output=Features[DataFrame[tuple([*output_features, pkey])]],  # type: ignore[misc]
+        fn=inference_fn,
+        environment=None,
+        machine_type=None,
+        default_args=[None],
+        timeout=None,
+        cron=None,
+        when=None,
+        tags=None,
+        owner=None,
+        resource_hint=resource_hint or model.resource_hint,
+        data_sources=None,
+        is_sql_file_resolver=False,
+        source_line=None,
+        lsp_builder=get_resolver_error_builder(inference_fn),
+        parse=None,
+        static=False,
+        total=False,
+        autogenerated=False,
+        unique_on=None,
+        partitioned_by=None,
+        data_lineage=None,
+        sql_settings=None,
+    )
+
+    # Register the resolver
+    RESOLVER_REGISTRY.add_to_registry(resolver, override=False)
+
+    return resolver

chalk/features/tag.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from typing import Sequence, Union
-
-from typing_extensions import TypeAlias
+from typing import Sequence, TypeAlias, Union
 
 Tags: TypeAlias = Union[str, Sequence[str]]
 """

chalk/features/underscore_features.py

@@ -24,6 +24,7 @@ from chalk.features.underscore import (
     UnderscoreRoot,
 )
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import schema_compat
 
 if typing_extensions.TYPE_CHECKING:
     from chalk import Features
@@ -66,7 +67,7 @@ def _parse_underscore_in_context(exp: Any, context: Any, is_pydantic: bool) -> A
         from chalk.features.dataframe import DataFrame
 
         if isinstance(parent_context, DataFrame) and is_pydantic:
-            if attr not in parent_context._underlying.schema:
+            if attr not in schema_compat(parent_context._underlying):
                 warnings.warn(
                     f"Attribute {attr} not found in dataframe schema. Returning None. Found expression {exp}."
                 )