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

chalk/client/dataset.py

@@ -566,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,

chalk/client/models.py

@@ -460,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]
@@ -838,6 +847,7 @@ class CreateOfflineQueryJobRequest(BaseModel):
         None,
         UploadedParquetShardedOfflineQueryInput,
         OfflineQueryInputUri,
+        OfflineQueryInputSql,
     ] = None
     """Any givens"""
 
@@ -1658,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):
@@ -1782,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

@@ -77,7 +77,15 @@ MODEL_SERIALIZERS = {
     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.",
+        ),
     ),
 }
 
@@ -281,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:
@@ -305,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:
@@ -322,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(