chalkruntime 3.34.2__tar.gz → 3.34.4__tar.gz

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkruntime
-Version: 3.34.2
+Version: 3.34.4
 Summary: Runtime support library for Chalk AI
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/chalkruntime/graph/convert_chalkpy_underscore.py

@@ -1108,6 +1108,30 @@ def _convert_underscore_expression(
                 if not isinstance(n_arg, UnderscoreConstant):
                     raise ValueError(f"2nd argument to of '{operation_name}' must be constant")
                 additional_args.append(n_arg)
+        elif operation_name == "approx_top_k" and "by" in underscore._chalk__kwargs:
+            # `approx_top_k(k=..., by=_.weight_col)` is a weighted top-k. Lift the `by=` kwarg
+            # into a positional operand so it shares the same shape as `max_by`/`min_by` for
+            # downstream has-many handling (e.g. `has_aggregate_input` / `_process_underscore_hm`).
+            if not isinstance(parent, UnderscoreItemParsed):
+                raise ValueError(f"'{operation_name}' with 'by=' must be called on a column")
+            if parent.feature_keys[0].root_namespace is None:
+                raise ValueError(
+                    f"'{operation_name}' with 'by=' must be called on a feature with a valid parent namespace"
+                )
+            by_arg = _convert_underscore_expression(
+                underscore=underscore._chalk__kwargs["by"],
+                namespace=parent.feature_keys[0].root_namespace,
+                graph=graph,
+                for_feature=for_feature,
+                local_scope=local_scope,
+                parent_op=parent_op,
+                context=context,
+            )
+            if not isinstance(by_arg, UnderscoreValue):
+                raise ValueError(
+                    f"Expected 'by' argument to '{operation_name}' to be an UnderscoreValue, received '{type(by_arg)}'"
+                )
+            additional_args: list[UnderscoreValue | UnderscoreConstant] = [by_arg]
         else:
             additional_args: list[UnderscoreValue | UnderscoreConstant] = []
 
@@ -1675,6 +1699,20 @@ def _convert_underscore_materialized_aggregation(
         first_arg = resolved_underscore.operands.positional_items[1]
         if not isinstance(first_arg, UnderscoreFeature):
             raise ValueError(f"The second argument to '{window_materialization.aggregation}' should be a feature")
+    elif window_materialization.aggregation == MaterializedAggregationType.ApproxTopK:
+        # `approx_top_k` is invoked as either `approx_top_k(_.col, k=...)` (1 operand) or
+        # `approx_top_k(_.col, by=_.weight, k=...)` (2 operands - the parser lifts `by=` into
+        # a positional operand). Either is valid.
+        if len(resolved_underscore.operands.positional_items) not in (1, 2):
+            raise ValueError(
+                f"The windowed aggregation operation '{window_materialization.aggregation}' accepts 1 or 2 positional arguments, but received {len(resolved_underscore.operands.positional_items)} arguments"
+            )
+        if len(resolved_underscore.operands.positional_items) == 2:
+            second_arg = resolved_underscore.operands.positional_items[1]
+            if not isinstance(second_arg, UnderscoreFeature):
+                raise ValueError(
+                    f"The second (`by=`) argument to '{window_materialization.aggregation}' should be a feature"
+                )
     elif not resolved_underscore.operands.has_exactly_n_and_only_positional_items(1):
         raise ValueError(
             f"The windowed aggregation operation '{window_materialization.aggregation}' only accepts a single positional argument, but received {len(resolved_underscore.operands.positional_items)} arguments"

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/chalkruntime/graph/underscore.py

@@ -785,7 +785,11 @@ class UnderscoreValidationError(Exception):
 
 OperationSupplier = Callable[[str], "UnderscoreOperation"]
 
-by_aggregate_fns = ("max_by", "min_by", "max_by_n", "min_by_n")
+by_aggregate_fns = ("max_by", "min_by", "max_by_n", "min_by_n", "approx_top_k")
+# `approx_top_k` is in `by_aggregate_fns` because it can be invoked with a `by=` operand
+# (lifted into a positional operand at parse time). Without `by=`, it has only a single
+# positional operand and is treated as a regular has-many aggregation.
+_always_by_aggregate_fns = frozenset(("max_by", "min_by", "max_by_n", "min_by_n"))
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -828,10 +832,24 @@ class UnderscoreOperationExpression(UnderscoreValue):
                 f"Expected kwargs to be a 'FrozenOrderedSet', but is actually a '{type(self.kwargs).__name__}'."
             )
 
+    def is_by_aggregate_invocation(self) -> bool:
+        """True iff this operation is invoked with a `by=` (sort) operand.
+
+        `max_by`/`min_by`/`max_by_n`/`min_by_n` always require a `by=` operand. `approx_top_k`
+        accepts an optional `by=` kwarg that the parser lifts into a 2nd positional operand;
+        when called without `by=`, it has only one positional operand (the has-many column)
+        and behaves like a normal aggregation.
+        """
+        if self.operation_name in _always_by_aggregate_fns:
+            return True
+        if self.operation_name == "approx_top_k":
+            return len(self.operands.positional_items) >= 2
+        return False
+
     @cached_property
     def has_aggregate_input(self):
         # Used to parse max_by and min_by aggregations
-        if self.operation_name in by_aggregate_fns and len(self.operands.positional_items) not in (2, 3):  # max/min_by
+        if not self.is_by_aggregate_invocation() or len(self.operands.positional_items) not in (2, 3):
             raise ValueError(
                 f"Expected 2 or 3 operands while converting expression to DataFrame filter; got: {self.operands.pprint()}"
             )
@@ -886,8 +904,10 @@ class UnderscoreOperationExpression(UnderscoreValue):
 
     @cached_property
     def root_namespace(self) -> str | None:
-        if self.operation_name in by_aggregate_fns:
-            # These functions are special
+        if self.is_by_aggregate_invocation():
+            # `by=` aggregations span two namespaces (the has-many's parent and the has-many's
+            # child, where the `by` column lives). Pick the first non-pseudonamespace as the
+            # canonical root rather than rejecting the multi-namespace case.
             special_namespaces = [
                 op.root_namespace
                 for _, op in MaybeNamedCollection.enumerate(self.operands)
@@ -936,7 +956,7 @@ class UnderscoreOperationExpression(UnderscoreValue):
     @cached_property
     def inputs(self) -> tuple[UnderscoreInputFeatureType, ...]:
         all_inputs: list[UnderscoreInputFeatureType] = []
-        if self.operation_name in by_aggregate_fns:
+        if self.is_by_aggregate_invocation():
             if isinstance(self.operands.positional_items[0], UnderscoreValue):
                 all_inputs.append(self.has_aggregate_input)
                 return tuple(all_inputs)

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/chalkruntime/incrementalization/incrementalizer.py

@@ -167,9 +167,14 @@ class Incrementalizer(QueryRewriter):
                             sqlglot.column(settings.incremental_column)
                             < sqlglot_expressions.Placeholder(this=END_TS_PARAM)
                         )
+                    params = {**chalk_query.params}
+                    if apply_lower_bound:
+                        params[START_TS_PARAM] = start_ts
+                    if apply_upper_bound:
+                        params[END_TS_PARAM] = end_ts
                     return FinalizedChalkQuery(
                         query=text(sqlglot_query.sql(dialect=chalk_query.source.get_sqlglot_dialect())),
-                        params={**chalk_query.params, START_TS_PARAM: start_ts, END_TS_PARAM: end_ts},
+                        params=params,
                         finalizer=chalk_query.finalizer,
                         incremental_settings=None,
                         source=chalk_query.source,

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/chalkruntime/invoker/vectorized_hasmany_sampler.py

@@ -19,7 +19,7 @@ from chalk.features.dataframe._validation import (
 )
 from chalk.utils.collections import FrozenOrderedSet, OrderedSet, unwrap_optional
 from chalk.utils.df_utils import pa_cast, pa_table_to_pl_df
-from chalk.utils.pl_helpers import is_new_polars, pl_is_uniquable_on
+from chalk.utils.pl_helpers import is_new_polars, pl_is_uniquable_on, str_json_decode_compat
 from chalk.utils.tracing import safe_trace
 
 from chalkruntime.constants import (
@@ -1660,7 +1660,7 @@ def _validate_df_schema(underlying: Union[pl.DataFrame, pl.LazyFrame], graph: Re
         elif (
             isinstance(expected_dtype, pl.List) and actual_dtype == pl.Utf8  # pyright: ignore[reportUnnecessaryComparison]
         ):
-            col = pl.col(root_fqn).str.json_extract(expected_dtype)
+            col = str_json_decode_compat(pl.col(root_fqn), expected_dtype)
             try:
                 underlying = underlying.with_columns(col.cast(expected_dtype))
             except (Exception, PolarsPanicErrorCompat) as e:

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/chalkruntime.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkruntime
-Version: 3.34.2
+Version: 3.34.4
 Summary: Runtime support library for Chalk AI
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{chalkruntime-3.34.2 → chalkruntime-3.34.4}/pyproject.toml

@@ -44,7 +44,7 @@ name = "chalkruntime"
 description = "Runtime support library for Chalk AI"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "3.34.2"
+version = "3.34.4"
 
 
 [tool.deptry]