chalkruntime 3.41.0__tar.gz → 3.41.2__tar.gz

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/PKG-INFO

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

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/graph/convert_chalkpy_underscore.py

@@ -3050,6 +3050,7 @@ def underscore_to_online_resolver(
         incremental_settings=None,
         is_externally_defined=False,
         output_row_order=None,
+        max_batch_size=None,
         venv=None,
         postprocessing=None,
     )
@@ -3107,6 +3108,7 @@ def underscore_to_offline_resolver(
         incremental_settings=None,
         is_externally_defined=False,
         output_row_order=None,
+        max_batch_size=None,
         venv=None,
         postprocessing=None,
     )

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/graph/dataframe_resolver_correlation.py

@@ -27,6 +27,14 @@ def dataframe_resolver_correlation_mode(resolver: ResolverParsed) -> DataFrameRe
     output_has_index = dataframe_type_has_index(resolver.unique_df_output, resolver.graph)
 
     if not indexed_inputs and not output_has_index:
+        if resolver.max_batch_size is not None:
+            # Legacy positional correlation cannot safely reassemble chunked outputs if the
+            # resolver filters or reorders rows, so chunking requires Index correlation.
+            raise ValueError(
+                f"Resolver '{resolver.fqn}' sets max_batch_size={resolver.max_batch_size}, but max_batch_size requires"
+                + " Index row correlation. Declare Index in the resolver's input DataFrame and output DataFrame, or"
+                + " remove max_batch_size."
+            )
         return DataFrameResolverCorrelationMode.LEGACY
 
     if len(indexed_inputs) != 1 or not output_has_index:

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/graph/protograph_deserializer.py

@@ -2294,6 +2294,7 @@ class FromProtoConverter:
             is_externally_defined=False,
             feature_expressions=feature_expressions,
             output_row_order=None,
+            max_batch_size=None,
             venv=None,
             stream_sink_config=sink_config,
             postprocessing=None,
@@ -2487,6 +2488,8 @@ class FromProtoConverter:
         if output_row_order not in (None, "one-to-one"):
             raise ValueError(f"output_row_order must be either None or 'one-to-one', received {r.output_row_order}")
 
+        max_batch_size = r.max_batch_size if r.HasField("max_batch_size") and r.max_batch_size > 0 else None
+
         postprocessing_expr: Underscore | None = None
         if r.WhichOneof("postprocessing") == "underscore_expr":
             if not r.HasField("sql_settings"):
@@ -2544,6 +2547,7 @@ class FromProtoConverter:
             else None,
             is_externally_defined=is_externally_defined,
             output_row_order=output_row_order,  # ty:ignore[invalid-argument-type]
+            max_batch_size=max_batch_size,
             venv=r.venv if r.HasField("venv") else None,
             postprocessing=postprocessing_expr,
         )

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/graph/resolver.py

@@ -301,6 +301,12 @@ class ResolverParsed:
     If "one-to-one", verify that this is a DF -> DF resolver, and treat it as a bulk scalar operation.
     """
 
+    max_batch_size: int | None = dataclasses.field(compare=False, repr=False)
+    """
+    Max rows per invocation for Index-correlated DF -> DF resolvers; the invoker slices the
+    input and invokes the resolver once per chunk. None = unlimited (single invocation).
+    """
+
     venv: str | None = dataclasses.field(compare=False, repr=False)
     """
     The venv that the resolver should be run in.
@@ -574,6 +580,17 @@ class ResolverParsed:
                     f"Resolver '{self.fqn}' has output_row_order='one-to-one', but it does not take in and return the pkey column."
                 )
 
+        if self.max_batch_size is not None:
+            if self.max_batch_size < 1:
+                raise ValueError(
+                    f"Resolver '{self.fqn}' has max_batch_size={self.max_batch_size}, but it must be >= 1."
+                )
+            if not self.input_is_df or not self.output_is_df:
+                raise ValueError(
+                    f"Resolver '{self.fqn}' has max_batch_size={self.max_batch_size}, but max_batch_size only applies"
+                    + " to DataFrame -> DataFrame resolvers."
+                )
+
     inputs_include_now_pseudofeature: bool = dataclasses.field(init=False, compare=False, repr=False)
 
     def get_fn(self):

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/invoker/bound_invoker.py

@@ -71,6 +71,8 @@ class PlanTimeInvokerConfig:
     # should correspond to the planner option value
     allow_planner_postponed_has_many_sampling: bool
     invoker_uses_strict_timeouts: bool
+    # should correspond to the planner option value; None or <=0 = unlimited
+    df_resolver_default_max_batch_size: int | None
 
     def __init__(
         self,
@@ -97,6 +99,7 @@ class PlanTimeInvokerConfig:
         oom_max_num_inputs_per_hm_sample: int,
         allow_planner_postponed_has_many_sampling: bool,
         invoker_uses_strict_timeouts: bool,
+        df_resolver_default_max_batch_size: int | None,
         _freeze: bool = True,
     ):
         super().__init__()
@@ -123,6 +126,7 @@ class PlanTimeInvokerConfig:
         self.oom_max_num_inputs_per_hm_sample = int(oom_max_num_inputs_per_hm_sample)
         self.allow_planner_postponed_has_many_sampling = allow_planner_postponed_has_many_sampling
         self.invoker_uses_strict_timeouts = invoker_uses_strict_timeouts
+        self.df_resolver_default_max_batch_size = _optional_int(df_resolver_default_max_batch_size)
         self._frozen = _freeze
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -154,6 +158,7 @@ class PlanTimeInvokerConfig:
             "oom_max_num_inputs_per_hm_sample": self.oom_max_num_inputs_per_hm_sample,
             "allow_planner_postponed_has_many_sampling": self.allow_planner_postponed_has_many_sampling,
             "invoker_uses_strict_timeouts": self.invoker_uses_strict_timeouts,
+            "df_resolver_default_max_batch_size": self.df_resolver_default_max_batch_size,
         }
         return _drop_none(data) if exclude_none else data
 
@@ -209,6 +214,7 @@ class InvokerConfig(PlanTimeInvokerConfig):
         oom_max_num_inputs_per_hm_sample: int,
         allow_planner_postponed_has_many_sampling: bool,
         invoker_uses_strict_timeouts: bool,
+        df_resolver_default_max_batch_size: int | None,
         operation_id: str,
         correlation_id: str | None = None,
         query_context: ContextJsonDict | None = None,
@@ -236,6 +242,7 @@ class InvokerConfig(PlanTimeInvokerConfig):
             oom_max_num_inputs_per_hm_sample=oom_max_num_inputs_per_hm_sample,
             allow_planner_postponed_has_many_sampling=allow_planner_postponed_has_many_sampling,
             invoker_uses_strict_timeouts=invoker_uses_strict_timeouts,
+            df_resolver_default_max_batch_size=df_resolver_default_max_batch_size,
             _freeze=False,
         )
         self.operation_id = operation_id
@@ -267,6 +274,7 @@ class InvokerConfig(PlanTimeInvokerConfig):
             "oom_max_num_inputs_per_hm_sample": self.oom_max_num_inputs_per_hm_sample,
             "allow_planner_postponed_has_many_sampling": self.allow_planner_postponed_has_many_sampling,
             "invoker_uses_strict_timeouts": self.invoker_uses_strict_timeouts,
+            "df_resolver_default_max_batch_size": self.df_resolver_default_max_batch_size,
             "operation_id": self.operation_id,
             "correlation_id": self.correlation_id,
             "query_context": self.query_context,

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/invoker/general_bound_invoker.py

@@ -682,8 +682,29 @@ class GeneralBoundInvoker(BoundInvokerProtocol):
                 "There should not be any scalarish features if the resolver is dataframe-taking or Now -> DataFrame"
             )
             assert batch.observed_at_fallback is not None
-            # Resolver takes only DFs or nothing. Should be called once.
-            yield arg_constructor(batch.observed_at_fallback, table_feature_to_results)
+            sample = arg_constructor(batch.observed_at_fallback, table_feature_to_results)
+            resolver_args = sample.resolver_args_or_error
+            max_batch_size = 0
+            if self._resolver.has_index_correlated_df_io and isinstance(resolver_args, tuple):
+                # Chunking is only safe under Index correlation: the operator joins outputs back to
+                # inputs by the Index column, so N invocations concatenate without positional
+                # assumptions. Legacy (hstack) correlation must keep the single invocation.
+                max_batch_size = self._resolver.max_batch_size or config.df_resolver_default_max_batch_size or 0
+            if max_batch_size > 0 and isinstance(resolver_args, tuple) and len(resolver_args[0]) > max_batch_size:
+                # has_index_correlated_df_io implies exactly one DataFrame input. Each chunk runs as
+                # its own invocation: the resolver timeout applies per chunk, and a chunk that raises
+                # produces a ChalkError for that chunk while the remaining chunks still emit data.
+                (df,) = resolver_args
+                for offset in range(0, len(df), max_batch_size):
+                    yield Sample(
+                        index=None,
+                        pkey=None,
+                        execution_ts=sample.execution_ts,
+                        resolver_args_or_error=(df.slice(offset, max_batch_size),),
+                    )
+            else:
+                # Resolver takes only DFs or nothing. Should be called once.
+                yield sample
         else:
             assert batch.scalarish_features is not None
             assert self._resolver.cron_filter is None, "Resolvers with filters are not supported"

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/invoker/one_to_one_invoker.py

@@ -663,6 +663,7 @@ class ParallelResolverInvoker(BoundInvokerProtocol):
                 incremental_settings=None,
                 is_externally_defined=False,
                 output_row_order=None,
+                max_batch_size=None,
                 venv=None,
                 postprocessing=None,
             )

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime/loader/converter.py

@@ -1105,6 +1105,7 @@ class RegistryConverter:
             is_externally_defined=False,
             feature_expressions=parsed_feature_expressions,
             output_row_order=None,
+            max_batch_size=None,
             venv=None,
             stream_sink_config=sink_config,
             postprocessing=r.postprocessing,
@@ -1167,6 +1168,7 @@ class RegistryConverter:
             incremental_settings=None,
             is_externally_defined=False,
             output_row_order=None,
+            max_batch_size=None,
             venv=None,
             postprocessing=r.postprocessing,
             runtime_contract=r.runtime_contract,
@@ -1267,6 +1269,7 @@ class RegistryConverter:
             incremental_settings=self._convert_incremental_config(r.incremental_settings),
             is_externally_defined=False,
             output_row_order=r.output_row_order,  # pyright: ignore[reportArgumentType]
+            max_batch_size=r.max_batch_size,
             venv=r.venv,
             postprocessing=r.postprocessing,
             handle_duplicate_outputs=cast(HandleDuplicateOutputs | None, r.handle_duplicate_outputs or None),
@@ -1333,6 +1336,7 @@ class RegistryConverter:
             incremental_settings=self._convert_incremental_config(r.incremental_settings),
             is_externally_defined=False,
             output_row_order=r.output_row_order,  # pyright: ignore[reportArgumentType]
+            max_batch_size=r.max_batch_size,
             venv=r.venv,
             postprocessing=r.postprocessing,
             handle_duplicate_outputs=cast(HandleDuplicateOutputs | None, r.handle_duplicate_outputs or None),

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/chalkruntime.egg-info/PKG-INFO

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

{chalkruntime-3.41.0 → chalkruntime-3.41.2}/pyproject.toml

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