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

chalk/client/dataset.py

@@ -48,6 +48,7 @@ from chalk.integrations.catalogs.base_catalog import BaseCatalog
 from chalk.utils.df_utils import read_parquet
 from chalk.utils.log_with_context import get_logger
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import apply_compat, polars_group_by_instead_of_groupby
 from chalk.utils.threading import DEFAULT_IO_EXECUTOR
 
 if TYPE_CHECKING:
@@ -541,14 +542,19 @@ def _extract_df_columns(
             for fqn in unique_features
         ]
 
-        df = df.groupby("pkey").agg(cols)
+        if polars_group_by_instead_of_groupby:
+            df = df.group_by("pkey").agg(cols)
+        else:
+            df = df.groupby("pkey").agg(cols)  # pyright: ignore
         decoded_stmts: List[pl.Expr] = []
         for col in df.columns:
             if col == "pkey":
                 continue
             else:
                 decoded_stmts.append(
-                    pl.col(col).apply(_json_decode, return_dtype=Feature.from_root_fqn(col).converter.polars_dtype)
+                    apply_compat(
+                        pl.col(col), _json_decode, return_dtype=Feature.from_root_fqn(col).converter.polars_dtype
+                    )
                 )
         df = df.select(decoded_stmts)
         # it might be a good idea to remember that we used to rename this __id__ column to the primary key
@@ -560,7 +566,13 @@ def _extract_df_columns(
 
     decoded_stmts: List[pl.Expr] = []
     feature_name_to_metadata = None if column_metadata is None else {x.feature_fqn: x for x in column_metadata}
-    for col, dtype in zip(df.columns, df.dtypes):
+    # Use collect_schema().dtypes() for newer Polars versions to avoid performance warning
+    # Fall back to df.dtypes for older versions
+    try:
+        dtypes = df.collect_schema().dtypes()
+    except AttributeError:
+        dtypes = df.dtypes
+    for col, dtype in zip(df.columns, dtypes):
         if version in (
             DatasetVersion.BIGQUERY_JOB_WITH_B32_ENCODED_COLNAMES,
             DatasetVersion.BIGQUERY_JOB_WITH_B32_ENCODED_COLNAMES_V2,
@@ -571,7 +583,7 @@ def _extract_df_columns(
                 # Assuming that the only datetime column is for timestamps
                 decoded_stmts.append(to_utc(df, col, pl.col(col)))
             else:
-                decoded_stmts.append(pl.col(col).apply(_json_decode, return_dtype=dtype))
+                decoded_stmts.append(apply_compat(pl.col(col), _json_decode, return_dtype=dtype))
         elif version in (DatasetVersion.NATIVE_DTYPES, DatasetVersion.NATIVE_COLUMN_NAMES):
             # We already decoded the column names so matching against the fqn
             if col == CHALK_TS_FEATURE or col == OBSERVED_AT_FEATURE:
@@ -1338,12 +1350,14 @@ class DatasetRevisionImpl(DatasetRevision):
                         actual_args.append(
                             DataFrame(
                                 pl.DataFrame([pl.Series(col_name, [], dtype=raw_input_df.schema[col_name])])
-                                .explode(pl.all())
+                                .explode(col_name)
                                 .unnest(col_name)
                             )
                         )
                     else:
-                        actual_args.append(DataFrame(has_many_input_df.explode(pl.all()).unnest(col_name)))
+                        actual_args.append(
+                            DataFrame(has_many_input_df.explode(has_many_input_df.columns).unnest(col_name))
+                        )
                 else:
                     value = args[i]
                     if isinstance(input, Feature):
@@ -1363,6 +1377,13 @@ This occurred during the actual execution of resolver {resolver.fqn}.
                     )
                     raise e
             print(f"resolver_replay: {resolver.fqn} returned {output}")
+            if isinstance(output, DataFrame):
+                try:
+                    output = output.to_polars().collect().rows(named=True)
+                except Exception as e:
+                    raise RuntimeError(
+                        f'Failed to convert DataFrame output from resolver "{resolver.fqn}" during resolver replay'
+                    ) from e
             output_col.append(output)
         return raw_input_df.with_columns(pl.Series(name="__resolver_replay_output__", values=output_col))
 

chalk/client/models.py

@@ -7,10 +7,9 @@ import traceback
 import uuid
 from datetime import datetime, timedelta
 from enum import Enum, IntEnum
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, TypeAlias, Union
 
 import numpy as np
-from typing_extensions import TypeAlias
 
 from chalk.byte_transmit.model import ByteBaseModel, ByteDict
 from chalk.client._internal_models.models import OfflineQueryGivensVersion
@@ -461,6 +460,15 @@ class OfflineQueryInput(BaseModel):
     values: List[List[Any]]  # Values should be of type TJSON
 
 
+class OfflineQueryInputSql(BaseModel):
+    """Input to an offline query specified as a ChalkSQL query instead
+    of literal data.
+
+    Alternative to OfflineQueryInput or OfflineQueryInputUri."""
+
+    input_sql: str
+
+
 class OnlineQueryRequest(BaseModel):
     inputs: Mapping[str, Any]  # Values should be of type TJSON
     outputs: List[str]
@@ -839,6 +847,7 @@ class CreateOfflineQueryJobRequest(BaseModel):
         None,
         UploadedParquetShardedOfflineQueryInput,
         OfflineQueryInputUri,
+        OfflineQueryInputSql,
     ] = None
     """Any givens"""
 
@@ -1659,6 +1668,7 @@ class PlanQueryResponse(BaseModel):
     output_schema: List[FeatureSchema]
     errors: List[ChalkError]
     structured_plan: Optional[str] = None
+    serialized_plan_proto_bytes: Optional[str] = None
 
 
 class IngestDatasetRequest(BaseModel):
@@ -1783,3 +1793,90 @@ class GetRegisteredModelVersionResponse(BaseModel):
 
 class CreateModelTrainingJobResponse(BaseModel):
     success: bool
+
+
+class ScheduledQueryRunStatus(str, Enum):
+    """Status of a scheduled query run."""
+
+    UNSPECIFIED = "UNSPECIFIED"
+    INITIALIZING = "INITIALIZING"
+    INIT_FAILED = "INIT_FAILED"
+    SKIPPED = "SKIPPED"
+    QUEUED = "QUEUED"
+    WORKING = "WORKING"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELED = "CANCELED"
+
+
+@dataclasses.dataclass
+class ScheduledQueryRun:
+    """A single scheduled query run."""
+
+    id: int
+    environment_id: str
+    deployment_id: str
+    run_id: str
+    cron_query_id: int
+    cron_query_schedule_id: int
+    cron_name: str
+    gcr_execution_id: str
+    gcr_job_name: str
+    offline_query_id: str
+    created_at: datetime
+    updated_at: datetime
+    status: ScheduledQueryRunStatus
+    blocker_operation_id: str
+
+    @staticmethod
+    def from_proto(proto_run: Any) -> "ScheduledQueryRun":
+        """Convert a proto ScheduledQueryRun to the dataclass version."""
+        from datetime import timezone
+
+        # Map proto status enum to our enum
+        status_map = {
+            0: ScheduledQueryRunStatus.UNSPECIFIED,
+            1: ScheduledQueryRunStatus.INITIALIZING,
+            2: ScheduledQueryRunStatus.INIT_FAILED,
+            3: ScheduledQueryRunStatus.SKIPPED,
+            4: ScheduledQueryRunStatus.QUEUED,
+            5: ScheduledQueryRunStatus.WORKING,
+            6: ScheduledQueryRunStatus.COMPLETED,
+            7: ScheduledQueryRunStatus.FAILED,
+            8: ScheduledQueryRunStatus.CANCELED,
+        }
+
+        # Helper to convert proto Timestamp to datetime
+        def _timestamp_to_datetime(ts: Any) -> datetime:
+            return datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=timezone.utc)
+
+        return ScheduledQueryRun(
+            id=proto_run.id,
+            environment_id=proto_run.environment_id,
+            deployment_id=proto_run.deployment_id,
+            run_id=proto_run.run_id,
+            cron_query_id=proto_run.cron_query_id,
+            cron_query_schedule_id=proto_run.cron_query_schedule_id,
+            cron_name=proto_run.cron_name,
+            gcr_execution_id=proto_run.gcr_execution_id,
+            gcr_job_name=proto_run.gcr_job_name,
+            offline_query_id=proto_run.offline_query_id,
+            created_at=_timestamp_to_datetime(proto_run.created_at),
+            updated_at=_timestamp_to_datetime(proto_run.updated_at),
+            status=status_map.get(proto_run.status, ScheduledQueryRunStatus.UNSPECIFIED),
+            blocker_operation_id=proto_run.blocker_operation_id,
+        )
+
+
+@dataclasses.dataclass
+class ManualTriggerScheduledQueryResponse:
+    """Response from manually triggering a scheduled query."""
+
+    scheduled_query_run: ScheduledQueryRun
+
+    @staticmethod
+    def from_proto(proto_response: Any) -> "ManualTriggerScheduledQueryResponse":
+        """Convert a proto ManualTriggerScheduledQueryResponse to the dataclass version."""
+        return ManualTriggerScheduledQueryResponse(
+            scheduled_query_run=ScheduledQueryRun.from_proto(proto_response.scheduled_query_run),
+        )

chalk/client/serialization/model_serialization.py

@@ -72,11 +72,20 @@ MODEL_SERIALIZERS = {
         filename="model.cbm",
         encoding=ModelEncoding.CBM,
         serialize_fn=lambda model, path: model.save_model(path),
+        schema_fn=lambda model: ModelAttributeExtractor.infer_catboost_schemas(model),
     ),
     ModelType.ONNX: ModelSerializationConfig(
         filename="model.onnx",
         encoding=ModelEncoding.PROTOBUF,
-        serialize_fn=lambda model, path: model.save_model(path),
+        serialize_fn=lambda model, path: ModelSerializer.with_import(
+            "onnx",
+            lambda onnx: onnx.save(
+                # Unwrap model if it has a _model attribute (e.g., wrapped ONNX models)
+                model._model if hasattr(model, "_model") else model,
+                path,
+            ),
+            "Please install onnx to save ONNX models.",
+        ),
     ),
 }
 
@@ -85,7 +94,13 @@ class ModelSerializer:
     def __init__(self, model: Any, model_type: Optional[ModelType]):
         self._temp_files: List[str] = []
         self.model = model
-        self.model_type = model_type if model_type is not None else ModelAttributeExtractor.infer_model_type(model)
+        if model_type is not None:
+            self.model_type = model_type
+            self.model_class = None
+        else:
+            model_type, model_class = ModelAttributeExtractor.infer_model_type(model)
+            self.model_type = model_type
+            self.model_class = model_class
 
         if self.model_type is None:
             raise ValueError("Unable to infer model type from object and no type given.")
@@ -125,6 +140,32 @@ class ModelSerializer:
         self._temp_files.append(model_path)
         return model_path, serializer_config.encoding
 
+    def serialize_to_path(self, path: str, cleanup: bool = True) -> Tuple[str, ModelEncoding]:
+        assert self.model_type is not None, "Could not determine model type. Please set parameter: model_type."
+        return self.serialize_model_to_path(self.model, self.model_type, path, cleanup)
+
+    def serialize_model_to_path(
+        self,
+        model: Any,
+        model_type: ModelType,
+        path: str,
+        cleanup: bool = True,
+    ) -> tuple[str, ModelEncoding]:
+        if model_type not in MODEL_SERIALIZERS:
+            raise NotImplementedError(f"Unsupported model type: {model_type}")
+        dir = path
+
+        serializer_config = MODEL_SERIALIZERS[model_type]
+
+        file_name = serializer_config.filename
+
+        model_path = os.path.join(dir, file_name)
+
+        serializer_config.serialize_fn(model, model_path)
+        if cleanup:
+            self._temp_files.append(model_path)
+        return model_path, serializer_config.encoding
+
     @staticmethod
     def with_import(module_name: str, func: Callable[[Any], Any], error_msg: str) -> Any:
         try:
@@ -248,7 +289,15 @@ class ModelSerializer:
         tensor_schema = _model_artifact_pb2.TensorSchema()
 
         for shape, dtype in tensor_specs:
-            if not isinstance(dtype, pa.DataType):
+            # Handle Chalk Tensor types
+            if hasattr(dtype, "__mro__") and any("Tensor" in base.__name__ for base in dtype.__mro__):
+                # Extract shape and dtype from Tensor type
+                if hasattr(dtype, "shape") and hasattr(dtype, "dtype"):
+                    shape = dtype.shape
+                    pa_dtype = dtype.dtype
+                else:
+                    raise ValueError(f"Tensor type is missing shape or dtype attributes")
+            elif not isinstance(dtype, pa.DataType):
                 if dtype == str:
                     pa_dtype = pa.string()
                 elif dtype == int:
@@ -272,12 +321,73 @@ class ModelSerializer:
 
         return tensor_schema
 
+    @staticmethod
+    def convert_onnx_list_schema_to_dict(schema: Any, model: Any, is_input: bool = True) -> Any:
+        """Convert list-based schema to dict-based schema for ONNX models.
+
+        Args:
+            schema: The schema (list or dict)
+            model: The ONNX model (ModelProto or wrapped)
+            is_input: True for input schema, False for output schema
+
+        Returns:
+            Dict-based schema with field names from ONNX model
+        """
+        if not isinstance(schema, list):
+            return schema
+
+        try:
+            import onnx  # type: ignore[reportMissingImports]
+        except ImportError:
+            raise ValueError("onnx package is required to convert list schemas for ONNX models")
+
+        # Unwrap model if needed
+        onnx_model = model._model if hasattr(model, "_model") else model
+
+        if not isinstance(onnx_model, onnx.ModelProto):
+            raise ValueError(
+                f"ONNX models must be registered with tabular schema (dict format). "
+                + f"Use dict format like {{'input': Tensor[...]}} instead of list format."
+            )
+
+        # Get input/output names from ONNX model
+        if is_input:
+            names = [inp.name for inp in onnx_model.graph.input]
+            schema_type = "input"
+        else:
+            names = [out.name for out in onnx_model.graph.output]
+            schema_type = "output"
+
+        if len(names) != len(schema):
+            raise ValueError(f"ONNX model has {len(names)} {schema_type}s but schema has {len(schema)} entries")
+
+        # Convert to dict format
+        return {name: spec for name, spec in zip(names, schema)}
+
     @staticmethod
     def convert_schema(schema: Any) -> Optional[_model_artifact_pb2.ModelSchema]:
         model_schema = _model_artifact_pb2.ModelSchema()
         if schema is not None:
             if isinstance(schema, dict):
-                model_schema.tabular.CopyFrom(ModelSerializer.build_tabular_schema(schema))
+                # Convert Tensor/Vector types to their PyArrow types for tabular schema
+                converted_schema = {}
+                for col_name, dtype in schema.items():
+                    if hasattr(dtype, "__mro__") and any("Tensor" in base.__name__ for base in dtype.__mro__):
+                        # Use Tensor's to_pyarrow_dtype() method to convert to Arrow type
+                        if hasattr(dtype, "to_pyarrow_dtype"):
+                            converted_schema[col_name] = dtype.to_pyarrow_dtype()
+                        else:
+                            raise ValueError(f"Tensor type for '{col_name}' is missing to_pyarrow_dtype method")
+                    elif hasattr(dtype, "__mro__") and any("Vector" in base.__name__ for base in dtype.__mro__):
+                        # Vector already has a .dtype attribute that's a PyArrow type
+                        if hasattr(dtype, "dtype"):
+                            converted_schema[col_name] = dtype.dtype
+                        else:
+                            raise ValueError(f"Vector type for '{col_name}' is missing dtype attribute")
+                    else:
+                        converted_schema[col_name] = dtype
+
+                model_schema.tabular.CopyFrom(ModelSerializer.build_tabular_schema(converted_schema))
             elif isinstance(schema, list):
                 model_schema.tensor.CopyFrom(ModelSerializer.build_tensor_schema(schema))
             else:
@@ -289,21 +399,27 @@ class ModelSerializer:
 
     @staticmethod
     def convert_run_criterion_to_proto(
-        run_name: Optional[str], criterion: Optional[ModelRunCriterion]
+        run_id: Optional[str] = None, run_name: Optional[str] = None, criterion: Optional[ModelRunCriterion] = None
     ) -> Optional[RunCriterion]:
-        if run_name is None:
-            return None
+        if run_id is None and run_name is None:
+            raise ValueError("Please specify either run_id or run_name.")
 
         if criterion is None:
-            return RunCriterion(run_id=run_name)
+            return RunCriterion(run_id=run_id, run_name=run_name)
 
         if criterion.direction == "max":
             return RunCriterion(
-                run_id=run_name, metric=criterion.metric, direction=RunCriterionDirection.RUN_CRITERION_DIRECTION_MAX
+                run_id=run_id,
+                run_name=run_name,
+                metric=criterion.metric,
+                direction=RunCriterionDirection.RUN_CRITERION_DIRECTION_MAX,
             )
         elif criterion.direction == "min":
             return RunCriterion(
-                run_id=run_name, metric=criterion.metric, direction=RunCriterionDirection.RUN_CRITERION_DIRECTION_MIN
+                run_id=run_id,
+                run_name=run_name,
+                metric=criterion.metric,
+                direction=RunCriterionDirection.RUN_CRITERION_DIRECTION_MIN,
             )
         else:
             raise ValueError(

chalk/config/project_config.py

@@ -184,7 +184,7 @@ def _load_project_config_at_path(filename: Path) -> Optional[ProjectSettings]:
     except OSError:
         return None
     except ValueError as e:
-        raise ValueError(f"Failed to load project config (chalkpy=={__version__}): {e}")
+        raise ValueError(f"Failed to load project config from {filename} (chalkpy=={__version__}): {e}") from e
 
 
 def load_project_config() -> Optional[ProjectSettings]: