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

chalk/streams/_windows.py

@@ -13,13 +13,12 @@ from typing import (
     Set,
     Tuple,
     Type,
+    TypedDict,
     TypeVar,
     Union,
     cast,
 )
 
-from typing_extensions import TypedDict
-
 from chalk._validation.feature_validation import FeatureValidation
 from chalk._validation.validation import Validation
 from chalk.features._encoding.primitive import TPrimitive
@@ -71,6 +70,7 @@ class WindowedMeta(type, Generic[TRich]):
             validations=None,
             offline_ttl=None,
             expression=None,
+            offline_expression=None,
             materialization=None,
         )  # noqa
 
@@ -80,9 +80,14 @@ JsonValue = Any
 
 def get_name_with_duration(name_or_fqn: str, duration: Union[str, int, timedelta]) -> str:
     duration_secs = parse_chalk_duration_s(duration)
+    name_or_fqn_components = name_or_fqn.split("@")
+    assert len(name_or_fqn_components) <= 2, f"Received invalid fqn format.\nfqn={name_or_fqn}"
+    unversioned_fqn = name_or_fqn_components[0]
+    version = None if len(name_or_fqn_components) != 2 else name_or_fqn_components[1]
+
     if duration_secs >= CHALK_MAX_TIMEDELTA.total_seconds():
-        return f"{name_or_fqn}__all__"
-    return f"{name_or_fqn}__{duration_secs}__"
+        return f"{unversioned_fqn}__all__" + ("" if version is None else f"@{version}")
+    return f"{unversioned_fqn}__{duration_secs}__" + ("" if version is None else f"@{version}")
 
 
 if TYPE_CHECKING:
@@ -176,6 +181,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
             window_durations=tuple(self.buckets_seconds) if bucket is None else tuple(),
             window_duration=window_duration,
             underscore_expression=self._expression,
+            offline_underscore_expression=self._offline_expression,
             window_materialization=(
                 MaterializationWindowConfig(bucket_duration=timedelta(seconds=window_duration))
                 if self._materialization is True and window_duration is not None
@@ -209,6 +215,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
         kind: Type[TRich] | None,
         offline_ttl: Duration | ellipsis | None,
         expression: Underscore | None,
+        offline_expression: Underscore | None,
         materialization: MaterializationWindowConfig | Literal[True] | None,
     ):
         super().__init__()
@@ -235,6 +242,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
         self._validations = validations
         self._dtype = dtype
         self._expression = expression
+        self._offline_expression = offline_expression
         self._materialization = materialization
 
 
@@ -575,6 +583,7 @@ def windowed(
     validations: List[Validation] | None = None,
     dtype: pa.DataType | None = None,
     expression: Underscore | None = None,
+    offline_expression: Underscore | None = None,
     materialization: MaterializationWindowConfig | Literal[True] | None = None,
 ) -> Windowed[TRich]:
     """Create a windowed feature.
@@ -651,6 +660,8 @@ def windowed(
         the feature value and will treat it as failed.
     expression
         The expression to compute the feature. This is an underscore expression, like `_.transactions[_.amount].sum()`.
+    offline_expression
+        Defines an alternate expression to compute the feature during offline queries.
     validations
         A list of Validations to apply to this feature.
 
@@ -749,5 +760,6 @@ def windowed(
         validations=validations,
         offline_ttl=offline_ttl,
         expression=expression,
+        offline_expression=offline_expression,
         materialization=materialization,
     )

chalk/streams/types.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, Any, List, Sequence, Set, Type, Union
+from typing import TYPE_CHECKING, Any, List, Sequence, Set, Type, TypeAlias, Union
 
 import pyarrow
-from typing_extensions import TypeAlias
 
 from chalk.utils import AnyDataclass
 

chalk/utils/__init__.py

@@ -1,6 +1,4 @@
-from typing import Any
-
-from typing_extensions import TypeAlias
+from typing import Any, TypeAlias
 
 MachineType: TypeAlias = str
 """The type of machine to use.

chalk/utils/_otel_version.py

@@ -0,0 +1,13 @@
+def _can_use_otel_trace() -> bool:
+    try:
+        from opentelemetry import context  # pyright: ignore[reportUnusedImport]
+        from opentelemetry import trace  # pyright: ignore[reportUnusedImport]
+        from opentelemetry.sdk.resources import Resource  # pyright: ignore[reportUnusedImport]
+        from opentelemetry.sdk.trace import TracerProvider  # pyright: ignore[reportUnusedImport]
+
+        return True
+    except Exception:
+        return False
+
+
+can_use_otel_trace = _can_use_otel_trace()

chalk/utils/async_helpers.py

@@ -6,9 +6,18 @@ import functools
 import queue
 import threading
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import Any, AsyncIterable, AsyncIterator, Callable, Coroutine, Generic, Iterable, TypeVar, overload
-
-from typing_extensions import ParamSpec
+from typing import (
+    Any,
+    AsyncIterable,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Generic,
+    Iterable,
+    ParamSpec,
+    TypeVar,
+    overload,
+)
 
 from chalk.utils.tracing import safe_activate_trace_context, safe_current_trace_context
 
@@ -52,8 +61,8 @@ def run_coroutine_fn_threadsafe(
             @functools.wraps(coro_fn)
             async def wrapped_with_context(*args: P.args, **kwargs: P.kwargs):
                 try:
-                    safe_activate_trace_context(current_trace_context)
-                    res = await coro_fn(*args, **kwargs)
+                    with safe_activate_trace_context(current_trace_context):
+                        res = await coro_fn(*args, **kwargs)
                 except BaseException as exc:
                     ans.set_exception(exc)
                 else:

chalk/utils/df_utils.py

@@ -595,9 +595,9 @@ def _read_parquet(
         if source.startswith("file://"):
             # Polars cannot handle local file uris properly
             source = source[len("file://") :]
-        if source.startswith("gs://"):
+        if source.startswith("gcs://"):
             # Fsspec expects gcs instead of gs
-            source = "gcs://" + source[len("gs://") :]
+            source = "gs://" + source[len("gcs://") :]
 
     try:
         # Polars added a crash bug in polars 0.20 that isn't fixed until polars 1.8.0; this fallback case handles

chalk/utils/duration.py

@@ -3,9 +3,7 @@ from __future__ import annotations
 import re
 from datetime import timedelta
 from re import Pattern
-from typing import Literal, Mapping, Optional, Tuple, Union
-
-from typing_extensions import TypeAlias
+from typing import Literal, Mapping, Optional, Tuple, TypeAlias, Union
 
 Duration: TypeAlias = Union[str, timedelta, Literal["infinity", "all"]]
 """