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

chalk/importer.py

@@ -17,8 +17,9 @@ from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Ty
 
 import pyarrow as pa
 
+import chalk.functions as F
 from chalk._lsp.error_builder import DiagnosticBuilder, LSPErrorBuilder
-from chalk.features import Feature, Features, FeatureSetBase, Filter, unwrap_feature
+from chalk.features import Feature, Features, FeatureSetBase, Filter, Vector, unwrap_feature
 from chalk.features.feature_field import WindowConfigResolved
 from chalk.features.pseudofeatures import Now
 
@@ -101,6 +102,8 @@ supported_aggs = (
     "sum",
     "var",
     "var_sample",
+    "vector_sum",
+    "vector_mean",
 )
 
 
@@ -169,7 +172,16 @@ def _check_types(
         return
 
     joined_annotation = joined_feature.typ.parsed_annotation
-    if aggregation not in {"count", "approx_count_distinct", "approx_top_k", "min_by_n", "max_by_n", "array_agg"}:
+    if aggregation not in {
+        "count",
+        "approx_count_distinct",
+        "approx_top_k",
+        "min_by_n",
+        "max_by_n",
+        "array_agg",
+        "vector_sum",
+        "vector_mean",
+    }:
         _validate_types(
             annotation=joined_annotation,
             permitted_types=(int, float),
@@ -195,7 +207,6 @@ def _check_types(
             joined=False,
             feature_name=feature_name,
         )
-
     elif aggregation == "min" or aggregation == "max":
         if _get_underlying_type(this_annotation, feature_name) != _get_underlying_type(
             joined_annotation, joined_feature.name
@@ -293,6 +304,19 @@ def _parse_agg_function_call(expr: Underscore | None) -> Tuple[str, Underscore,
                 f"expecting 'int' type argument for 'k', but received arg of type '{type(call_expr._chalk__kwargs.get('k'))}'"
             )
         opts = FrozenOrderedSet(call_expr._chalk__kwargs.items())
+    elif aggregation == "approx_percentile":
+        if len(call_expr._chalk__args) > 0:
+            raise ChalkParseError("should not have any positional arguments")
+        elif {"quantile"} != call_expr._chalk__kwargs.keys():
+            raise ChalkParseError("expecting exactly one required keyword argument 'quantile'")
+        elif not isinstance(call_expr._chalk__kwargs.get("quantile"), float):
+            raise ChalkParseError(
+                f"expecting 'float' type argument for 'quantile', but received arg of type '{type(call_expr._chalk__kwargs.get('quantile'))}'"
+            )
+        # TODO: expand proto definition to accept kwargs that are not necessarily `k`
+        quantile = call_expr._chalk__kwargs["quantile"]
+        nano_quantile = int(round(quantile * 1_000_000_000))
+        opts = FrozenOrderedSet([("k", nano_quantile)])
     elif aggregation in ("min_by_n", "max_by_n"):
         if len(call_expr._chalk__kwargs) > 0:
             raise ChalkParseError("should not have any keyword arguments")
@@ -422,8 +446,6 @@ def run_post_import_fixups():
             #       "1m", "2m", materialization={...},
             #       expression=_.transactions[_.amount].sum(),
             #   )
-            assert f.underscore_expression is not None
-            assert f.window_materialization is not None
 
             try:
                 f.window_materialization_parsed = parse_windowed_materialization(f=f)
@@ -561,39 +583,51 @@ def parse_grouped_window(f: Feature) -> WindowConfigResolved:
         aggregation_kwargs=aggregation_kwargs,
         pyarrow_dtype=pyarrow_dtype,
         filters=parsed_filters,
-        backfill_resolver=_try_parse_resolver_fqn(
-            "backfill_resolver",
-            f.window_materialization.get("backfill_resolver", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_schedule=f.window_materialization.get("backfill_schedule", None)
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_lookback_duration_seconds=_try_parse_duration(
-            "backfill_lookback_duration",
-            f.window_materialization.get("backfill_lookback_duration", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_start_time=_try_parse_datetime(
-            "backfill_start_time",
-            f.window_materialization.get("backfill_start_time", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        continuous_resolver=_try_parse_resolver_fqn(
-            "continuous_resolver",
-            f.window_materialization.get("continuous_resolver", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        continuous_buffer_duration_seconds=_try_parse_duration(
-            "continuous_buffer_duration",
-            f.window_materialization.get("continuous_buffer_duration", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
+        backfill_resolver=(
+            _try_parse_resolver_fqn(
+                "backfill_resolver",
+                f.window_materialization.get("backfill_resolver", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_schedule=(
+            f.window_materialization.get("backfill_schedule", None)
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_lookback_duration_seconds=(
+            _try_parse_duration(
+                "backfill_lookback_duration",
+                f.window_materialization.get("backfill_lookback_duration", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_start_time=(
+            _try_parse_datetime(
+                "backfill_start_time",
+                f.window_materialization.get("backfill_start_time", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        continuous_resolver=(
+            _try_parse_resolver_fqn(
+                "continuous_resolver",
+                f.window_materialization.get("continuous_resolver", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        continuous_buffer_duration_seconds=(
+            _try_parse_duration(
+                "continuous_buffer_duration",
+                f.window_materialization.get("continuous_buffer_duration", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
     )
 
     return cfg
@@ -710,6 +744,14 @@ def parse_windowed_materialization(f: Feature) -> WindowConfigResolved | None:
         aggregated_feature_name=aggregated_value,
     )
 
+    if aggregation == "sum" or aggregation == "mean":
+        try:
+            if issubclass(f.typ.parsed_annotation, Vector):
+                aggregation = f"vector_{aggregation}"
+        except TypeError:
+            # Not a class so not a Vector, skip
+            pass
+
     _check_types(
         feature_name=f.window_stem,
         aggregation=aggregation,
@@ -781,39 +823,51 @@ def parse_windowed_materialization(f: Feature) -> WindowConfigResolved | None:
         aggregation_kwargs=aggregation_kwargs,
         pyarrow_dtype=f.converter.pyarrow_dtype,
         filters=parsed_filters,
-        backfill_resolver=_try_parse_resolver_fqn(
-            "backfill_resolver",
-            f.window_materialization.get("backfill_resolver", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_schedule=f.window_materialization.get("backfill_schedule", None)
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_lookback_duration_seconds=_try_parse_duration(
-            "backfill_lookback_duration",
-            f.window_materialization.get("backfill_lookback_duration", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        backfill_start_time=_try_parse_datetime(
-            "backfill_start_time",
-            f.window_materialization.get("backfill_start_time", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        continuous_resolver=_try_parse_resolver_fqn(
-            "continuous_resolver",
-            f.window_materialization.get("continuous_resolver", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
-        continuous_buffer_duration_seconds=_try_parse_duration(
-            "continuous_buffer_duration",
-            f.window_materialization.get("continuous_buffer_duration", None),
-        )
-        if isinstance(f.window_materialization, dict)
-        else None,
+        backfill_resolver=(
+            _try_parse_resolver_fqn(
+                "backfill_resolver",
+                f.window_materialization.get("backfill_resolver", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_schedule=(
+            f.window_materialization.get("backfill_schedule", None)
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_lookback_duration_seconds=(
+            _try_parse_duration(
+                "backfill_lookback_duration",
+                f.window_materialization.get("backfill_lookback_duration", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        backfill_start_time=(
+            _try_parse_datetime(
+                "backfill_start_time",
+                f.window_materialization.get("backfill_start_time", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        continuous_resolver=(
+            _try_parse_resolver_fqn(
+                "continuous_resolver",
+                f.window_materialization.get("continuous_resolver", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
+        continuous_buffer_duration_seconds=(
+            _try_parse_duration(
+                "continuous_buffer_duration",
+                f.window_materialization.get("continuous_buffer_duration", None),
+            )
+            if isinstance(f.window_materialization, dict)
+            else None
+        ),
     )
 
 
@@ -991,6 +1045,33 @@ class _UnderscoreValidationError(ValueError):
     ...
 
 
+def _has_group_by_in_parent_chain(underscore: Underscore) -> bool:
+    """
+    Traverse parent chain to check if .group_by() exists before .agg().
+
+    For valid group_by_windowed: _.x.group_by(_.y).agg(_.z.sum())
+    - Looks for: UnderscoreCall -> UnderscoreAttr("group_by")
+
+    Returns True if .group_by() found, False otherwise.
+    """
+    current: Optional[Any] = underscore
+
+    while current is not None:
+        # Check if current is a .group_by() call
+        if isinstance(current, UnderscoreCall):
+            parent = current._chalk__parent
+            if isinstance(parent, UnderscoreAttr) and parent._chalk__attr == "group_by":
+                return True
+
+        # Move to parent
+        if hasattr(current, "_chalk__parent"):
+            current = current._chalk__parent
+        else:
+            break
+
+    return False
+
+
 class ChalkImporter:
     def __init__(self):
         super().__init__()
@@ -1092,6 +1173,9 @@ class ChalkImporter:
         for feature_class in FeatureSetBase.registry.values():
             # Iterate through every class, to find every underscore definition.
             for f in feature_class.features:
+                if f.is_windowed_pseudofeature is True:
+                    # need one LSP just for the base
+                    continue
                 if f.underscore_expression is not None:
                     # Validate that the underscore expression is well-formed.
                     # If it is not well-formed, then an `_UnderscoreValidationError` will
@@ -1424,15 +1508,26 @@ def _supplemental_validate_underscore_expression(
 
     # TODO: Dominic - impl for UnderscoreCall args (we need some special casing for aggregate functions that take in UnderscoreItems)
     if isinstance(underscore, UnderscoreCall):
+        if not isinstance(underscore._chalk__parent, UnderscoreAttr):
+            # we only support calls on attrs, ie _.a.some_attr(*args, **kwargs)
+            raise _UnderscoreValidationError(f"Cannot call non-attribute {underscore._chalk__parent}.")
         caller = underscore._chalk__parent._chalk__parent
+        op_name = underscore._chalk__parent._chalk__attr
+
+        if (op := getattr(F, op_name, None)) is not None:
+            if getattr(op, "_chalk__method_chaining_predicate", lambda _: True)(underscore):
+                return _supplemental_validate_underscore_expression(
+                    state,
+                    class_namespace=class_namespace,
+                    underscore=op(caller, *underscore._chalk__args, **underscore._chalk__kwargs),
+                )
+
         maybe_parent_result = _supplemental_validate_underscore_expression(
             state=state,
             class_namespace=class_namespace,
             underscore=caller,
         )
-        if not isinstance(underscore._chalk__parent, UnderscoreAttr):
-            return None  # TODO: Dominic - is this ever valid?
-        if underscore._chalk__parent._chalk__attr == "where":
+        if op_name == "where":
             if maybe_parent_result is None:
                 return None
             if not isinstance(maybe_parent_result, _HasManyNamespaceExpr) or not isinstance(caller, UnderscoreItem):
@@ -1457,9 +1552,26 @@ def _supplemental_validate_underscore_expression(
                     raise _UnderscoreValidationError(
                         f"the input '{arg!r}' is a feature namespace '{expr.namespace}' which cannot be used as a scalar value"
                     )
+            return None
+
+        # Validate .agg() usage (addressing TODO at line 1522)
+        if op_name == "agg":
+            if not _has_group_by_in_parent_chain(caller):
+                raise _UnderscoreValidationError(
+                    "'.agg()' can only be used with '.group_by()' for group_by_windowed features. "
+                    + "For windowed features, use direct aggregation methods instead. "
+                    + "For example, instead of using '.agg(_.field.method())', use '.field.method()' directly on the filtered DataFrame"
+                )
 
         return None
 
+        # TODO: check that op_name is a supported agg or .agg/.group_by/etc
+        # if op_name in supported_aggs:
+        #     # TODO: typechecking for agg fns
+        #     return None
+        #
+        # raise _UnderscoreValidationError(f"unrecognized function '{op_name}' in expression '{underscore}'")
+
     if isinstance(underscore, UnderscoreItem):
         parent_result = _supplemental_validate_underscore_expression(
             state=state,

chalk/ml/__init__.py

@@ -1,16 +1,20 @@
 from __future__ import annotations
 
-from chalk.ml.model_file_transfer import HFSourceConfig, LocalSourceConfig, S3SourceConfig, SourceConfig
+from chalk.ml.model_file_transfer import FileInfo, HFSourceConfig, LocalSourceConfig, S3SourceConfig, SourceConfig
 from chalk.ml.model_reference import ModelReference
-from chalk.ml.utils import ModelEncoding, ModelRunCriterion, ModelType
+from chalk.ml.model_version import ModelVersion
+from chalk.ml.utils import ModelClass, ModelEncoding, ModelRunCriterion, ModelType
 
 __all__ = (
     "ModelType",
+    "ModelClass",
     "ModelEncoding",
     "ModelReference",
+    "ModelVersion",
     "SourceConfig",
     "LocalSourceConfig",
     "S3SourceConfig",
     "HFSourceConfig",
     "ModelRunCriterion",
+    "FileInfo",
 )