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

chalk/importer.py

@@ -304,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")
@@ -433,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)
@@ -572,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
@@ -800,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
+        ),
     )
 
 
@@ -1010,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__()
@@ -1111,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
@@ -1489,6 +1554,15 @@ def _supplemental_validate_underscore_expression(
                     )
             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

chalk/ml/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from chalk.ml.model_file_transfer import FileInfo, HFSourceConfig, LocalSourceConfig, S3SourceConfig, SourceConfig
 from chalk.ml.model_reference import ModelReference
+from chalk.ml.model_version import ModelVersion
 from chalk.ml.utils import ModelClass, ModelEncoding, ModelRunCriterion, ModelType
 
 __all__ = (
@@ -9,6 +10,7 @@ __all__ = (
     "ModelClass",
     "ModelEncoding",
     "ModelReference",
+    "ModelVersion",
     "SourceConfig",
     "LocalSourceConfig",
     "S3SourceConfig",

chalk/ml/model_hooks.py

@@ -1,12 +1,15 @@
-from typing import Any, Dict, Optional, Protocol, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Tuple
 
 from chalk.ml.utils import ModelClass, ModelEncoding, ModelType
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 
 class ModelInference(Protocol):
     """Abstract base class for model loading and inference."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         """Load a model from the given path."""
         pass
 
@@ -14,11 +17,27 @@ class ModelInference(Protocol):
         """Run inference on the model with input X."""
         pass
 
+    def prepare_input(self, feature_table: Any) -> Any:
+        """Convert PyArrow table to model input format.
+
+        Default implementation converts to numpy array via __array__().
+        Override for model-specific input formats (e.g., ONNX struct arrays).
+        """
+        return feature_table.__array__()
+
+    def extract_output(self, result: Any, output_feature_name: str) -> Any:
+        """Extract single output from model result.
+
+        Default implementation returns result as-is (for single outputs).
+        Override for models with structured outputs (e.g., ONNX struct arrays).
+        """
+        return result
+
 
 class XGBoostClassifierInference(ModelInference):
     """Model inference for XGBoost classifiers."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import xgboost  # pyright: ignore[reportMissingImports]
 
         model = xgboost.XGBClassifier()
@@ -32,7 +51,7 @@ class XGBoostClassifierInference(ModelInference):
 class XGBoostRegressorInference(ModelInference):
     """Model inference for XGBoost regressors."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import xgboost  # pyright: ignore[reportMissingImports]
 
         model = xgboost.XGBRegressor()
@@ -46,17 +65,27 @@ class XGBoostRegressorInference(ModelInference):
 class PyTorchInference(ModelInference):
     """Model inference for PyTorch models."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import torch  # pyright: ignore[reportMissingImports]
 
         torch.set_grad_enabled(False)
+
+        # Load the model
         model = torch.jit.load(path)
-        model.input_to_tensor = lambda X: torch.from_numpy(X).float()
+
+        # If resource_hint is "gpu", move model to GPU
+        if resource_hint == "gpu" and torch.cuda.is_available():
+            device = torch.device("cuda")
+            model = model.to(device)
+            model.input_to_tensor = lambda X: torch.from_numpy(X).float().to(device)
+        else:
+            model.input_to_tensor = lambda X: torch.from_numpy(X).float()
+
         return model
 
     def predict(self, model: Any, X: Any) -> Any:
         outputs = model(model.input_to_tensor(X))
-        result = outputs.detach().numpy().astype("float64")
+        result = outputs.detach().cpu().numpy().astype("float64")
         result = result.squeeze()
 
         # Convert 0-dimensional array to scalar, or ensure we have a proper 1D array
@@ -69,7 +98,7 @@ class PyTorchInference(ModelInference):
 class SklearnInference(ModelInference):
     """Model inference for scikit-learn models."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import joblib  # pyright: ignore[reportMissingImports]
 
         return joblib.load(path)
@@ -81,7 +110,7 @@ class SklearnInference(ModelInference):
 class TensorFlowInference(ModelInference):
     """Model inference for TensorFlow models."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import tensorflow  # pyright: ignore[reportMissingImports]
 
         return tensorflow.keras.models.load_model(path)
@@ -93,7 +122,7 @@ class TensorFlowInference(ModelInference):
 class LightGBMInference(ModelInference):
     """Model inference for LightGBM models."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import lightgbm  # pyright: ignore[reportMissingImports]
 
         return lightgbm.Booster(model_file=path)
@@ -105,7 +134,7 @@ class LightGBMInference(ModelInference):
 class CatBoostInference(ModelInference):
     """Model inference for CatBoost models."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import catboost  # pyright: ignore[reportMissingImports]
 
         return catboost.CatBoost().load_model(path)
@@ -115,31 +144,170 @@ class CatBoostInference(ModelInference):
 
 
 class ONNXInference(ModelInference):
-    """Model inference for ONNX models."""
+    """Model inference for ONNX models with struct input/output support."""
 
-    def load_model(self, path: str) -> Any:
+    def load_model(self, path: str, resource_hint: Optional["ResourceHint"] = None) -> Any:
         import onnxruntime  # pyright: ignore[reportMissingImports]
 
-        return onnxruntime.InferenceSession(path)
+        # Conditionally add CUDAExecutionProvider based on resource_hint
+        providers = (
+            ["CUDAExecutionProvider", "CPUExecutionProvider"] if resource_hint == "gpu" else ["CPUExecutionProvider"]
+        )
+        return onnxruntime.InferenceSession(path, providers=providers)
+
+    def prepare_input(self, feature_table: Any) -> Any:
+        """Convert PyArrow table to struct array for ONNX models."""
+        import pyarrow as pa
+
+        # Get arrays for each column, combining chunks if necessary
+        arrays = []
+        for i in range(feature_table.num_columns):
+            col = feature_table.column(i)
+            if isinstance(col, pa.ChunkedArray):
+                arrays.append(col.combine_chunks())
+            else:
+                arrays.append(col)
+
+        # Create fields from schema, preserving original field names
+        # Field names should match ONNX input names exactly
+        fields = []
+        for field in feature_table.schema:
+            fields.append(pa.field(field.name, field.type))
+
+        # Create struct array where each row is a struct with named fields
+        return pa.StructArray.from_arrays(arrays, fields=fields)
+
+    def extract_output(self, result: Any, output_feature_name: str) -> Any:
+        """Extract single field from ONNX struct output."""
+        import pyarrow as pa
+
+        if not isinstance(result, (pa.StructArray, pa.ChunkedArray)):
+            return result
+
+        struct_type = result.type if isinstance(result, pa.StructArray) else result.chunk(0).type
+
+        # Find matching field by name, or use first field
+        field_index = None
+        for i, field in enumerate(struct_type):
+            if field.name == output_feature_name:
+                field_index = i
+                break
+
+        return result.field(field_index if field_index is not None else 0)
 
     def predict(self, model: Any, X: Any) -> Any:
+        """Run ONNX inference with struct input/output."""
+        # Get ONNX model input/output names
+        input_names = [inp.name for inp in model.get_inputs()]
+        output_names = [out.name for out in model.get_outputs()]
+
+        # Convert struct input to ONNX input dict
+        input_dict = self._struct_to_inputs(X, input_names)
+
+        # Run ONNX inference
+        outputs = model.run(output_names, input_dict)
+
+        # Always return outputs as struct array
+        return self._outputs_to_struct(output_names, outputs)
+
+    def _struct_to_inputs(self, struct_array: Any, input_names: list) -> dict:
+        """Extract ONNX inputs from struct array by matching field names.
+
+        Struct field names must match ONNX input names (supports list/Tensor types).
+        If ONNX expects a single input but struct has multiple scalar fields,
+        stack them into a 2D array.
+        """
         import numpy as np
+        import pyarrow as pa
+
+        if isinstance(struct_array, pa.ChunkedArray):
+            struct_array = struct_array.combine_chunks()
+
+        input_dict = {}
+        struct_fields = {field.name: i for i, field in enumerate(struct_array.type)}
+
+        # Check if struct field names match ONNX input names
+        fields_match = all(input_name in struct_fields for input_name in input_names)
+
+        if not fields_match:
+            # Special case 1: ONNX expects single input and struct has single field
+            # Use that field regardless of name mismatch
+            if len(input_names) == 1 and len(struct_fields) == 1:
+                field_data = struct_array.field(0)
+                input_dict[input_names[0]] = self._arrow_to_numpy(field_data)
+                return input_dict
+
+            # Special case 2: ONNX expects single input, but struct has multiple scalar fields
+            # Stack them into a 2D array [batch_size, num_fields]
+            if len(input_names) == 1 and len(struct_fields) > 1:
+                # Check if all fields are scalar (not nested lists)
+                all_scalar = all(
+                    not pa.types.is_list(struct_array.type[i].type)
+                    and not pa.types.is_large_list(struct_array.type[i].type)
+                    for i in range(len(struct_array.type))
+                )
+
+                if all_scalar:
+                    # Stack all fields into a single 2D array
+                    columns = []
+                    for i in range(len(struct_array.type)):
+                        field_data = struct_array.field(i)
+                        col_array = self._arrow_to_numpy(field_data)
+                        columns.append(col_array)
+
+                    # Stack columns horizontally to create [batch_size, num_features]
+                    stacked = np.column_stack(columns)
+                    input_dict[input_names[0]] = stacked
+                    return input_dict
+
+            raise ValueError(
+                f"ONNX inputs {input_names} not found in struct fields {list(struct_fields.keys())}. "
+                + "Struct field names must match ONNX input names."
+            )
+
+        # Direct mapping: struct fields match ONNX inputs (for Tensor/list types or named inputs)
+        for input_name in input_names:
+            field_data = struct_array.field(struct_fields[input_name])
+            input_dict[input_name] = self._arrow_to_numpy(field_data)
+
+        return input_dict
+
+    def _arrow_to_numpy(self, arrow_array: Any) -> Any:
+        """Convert Arrow array (including nested lists) to dense numpy array."""
+        import numpy as np
+        import pyarrow as pa
 
-        # Get input names from the model metadata
-        input_names = [inp.name for inp in model.get_inputs()]
+        if isinstance(arrow_array, pa.ChunkedArray):
+            arrow_array = arrow_array.combine_chunks()
 
-        # Convert X to float32 if needed
-        X_float32 = X.astype("float32") if hasattr(X, "astype") else np.array(X, dtype="float32")
+        # Convert to Python list, then numpy - handles all cases (nested lists, flat arrays, etc.)
+        return np.array(arrow_array.to_pylist(), dtype=np.float32)
 
-        # If there's only one input, use it directly
-        if len(input_names) == 1:
-            input_dict = {input_names[0]: X_float32}
-        else:
-            # For multiple inputs, we'd need additional logic
-            # For now, assume the first input is the main one
-            input_dict = {input_names[0]: X_float32}
+    def _outputs_to_struct(self, output_names: list, outputs: list) -> Any:
+        """Convert ONNX outputs to PyArrow struct array."""
+        import pyarrow as pa
+
+        if not outputs:
+            raise ValueError("ONNX model returned no outputs")
+
+        # Convert each output to Arrow array with proper type
+        fields = []
+        arrays = []
+
+        for name, output_array in zip(output_names, outputs):
+            arrow_array = self._numpy_to_arrow_array(output_array)
+            fields.append(pa.field(name, arrow_array.type))
+            arrays.append(arrow_array)
+
+        return pa.StructArray.from_arrays(arrays, fields=fields)
+
+    def _numpy_to_arrow_array(self, arr: Any) -> Any:
+        """Convert numpy array to PyArrow array (possibly nested list)."""
+        import pyarrow as pa
 
-        return model.run(None, input_dict)[0]
+        # PyArrow can infer the correct nested list type from Python lists
+        # Shape (batch, dim1, dim2, ...) -> list[list[...]]
+        return pa.array(arr.tolist())
 
 
 class ModelInferenceRegistry: