snowflake.ml.model.custom_model.partitioned_api

snowflake.ml.model.custom_model.partitioned_api(func: Callable[[Concatenate[CustomModelType, DataFrame, InferenceParams]], DataFrame]) → Callable[[Concatenate[CustomModelType, DataFrame, InferenceParams]], DataFrame]

Decorator to mark a method as a partitioned inference API in a CustomModel.

Methods decorated with @partitioned_api are exposed as partitioned inference endpoints, enabling efficient batch processing where the model processes data partitions independently. This is useful for models that can benefit from parallel execution across data partitions.

The decorated method must accept a pandas DataFrame as its first argument (after self) and return a pandas DataFrame.

Parameters:

func – The method to decorate.

Returns:

The decorated function with partitioned API metadata.

Example:

class MyPartitionedModel(CustomModel):
    @partitioned_api
    def predict(self, input_df: pd.DataFrame) -> pd.DataFrame:
        # Process each partition independently
        return pd.DataFrame({"output": input_df["feature"] * 2})
Copy
Language: English