ai_predict_class function

Applies to: Databricks SQL Databricks Runtime 19.3 and above

Beta

This feature is in Beta. Workspace admins can control access to this feature from the Previews page. See Manage Databricks previews.

ai_predict_class() trains a classification model from rows where the target column is not NULL and predicts the class for rows where it is NULL. The function returns only the rows it scores and appends prediction and confidence columns to the input columns.

Requirements

• A Pro or Serverless SQL warehouse, or a cluster running Databricks Runtime 19.3 or above
• Enroll your workspace in the Predictive AI Functions preview. See Manage Databricks previews.

Syntax

ai_predict_class(
  input => TABLE(input),
  target_col => target_col,
  feature_cols => feature_cols
  [, seed => seed]
)

Arguments

Pass all scalar arguments by name. You can pass input by name or position.

• input is the table-valued input that contains the training rows and the rows to score. Rows where the target column is not NULL train the model. Rows where the target column is NULL are scored.
• target_col is a constant STRING that names the target column. The target must be a top-level numeric, Boolean, or string column and must not appear in feature_cols.
• feature_cols is a non-empty constant ARRAY<STRING> that names the feature columns. Each feature must be a top-level numeric, Boolean, or string column.
• seed (optional) is an integral value that initializes pseudorandom operations used during model training. The default is 0.

Returns

The function returns a table containing the input rows where the target column is NULL. The output retains all input columns and appends the following columns:

• <target_col>_prediction: The predicted class, with the same data type as the target column.
• <target_col>_confidence: A DOUBLE between 0 and 1 that contains the model probability for the predicted class.

If no input rows have a NULL target, the function returns an empty table.

Example

The following example trains on customers with known churn outcomes and predicts whether customers with missing outcomes will churn:

SQL

WITH customers AS (
  SELECT * FROM VALUES
    (1, 3, 'monthly', true),
    (2, 24, 'annual', false),
    (3, 5, 'monthly', true),
    (4, 30, 'annual', false),
    (5, 8, 'monthly', true),
    (6, 36, 'annual', false),
    (7, 6, 'monthly', NULL),
    (8, 28, 'annual', NULL)
  AS customers(customer_id, tenure_months, plan, churned)
)
SELECT customer_id, churned_prediction, churned_confidence
FROM ai_predict_class(
  input => TABLE(customers),
  target_col => 'churned',
  feature_cols => ARRAY('tenure_months', 'plan'),
  seed => 42
)
ORDER BY customer_id;

Limitations

• Each invocation supports one target column and trains a new model. The function does not persist the model for reuse.
• Target and feature columns must be top-level columns. Nested fields are not supported.
• The training rows must contain at least two distinct, non-NULL target values.
• The input must not already contain columns named <target_col>_prediction or <target_col>_confidence.

Error conditions

For invalid columns, types, or arguments, see AI_PREDICT_INVALID_PARAMETER error condition. For inadequate training data, see AI_PREDICT_INSUFFICIENT_TRAINING_DATA error condition.

Related functions

To predict a numeric value, use ai_predict_value function. To classify text against labels that you provide without training on tabular rows, use ai_classify function. For other AI Functions, see Transform unstructured data using AI Functions.