Skip to main content
Last updated on

ai_predict_value function

Applies to: check marked yes Databricks SQL check marked yes 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_value() trains a regression model from rows where the target column is not NULL and predicts the target for rows where it is NULL. The function returns only the rows it scores and appends a prediction column 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_value(
  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 numeric target column. The column must be a top-level 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 <target_col>_prediction as a DOUBLE.

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

Example​

The following example trains on rows with known revenue and predicts revenue for stores where it is missing:

SQL
WITH sales AS (
  SELECT * FROM VALUES
    (1, 10.0, 'east', 120.0),
    (2, 12.0, 'east', 135.0),
    (3, 15.0, 'west', 160.0),
    (4, 18.0, 'west', 190.0),
    (5, 20.0, 'east', 205.0),
    (6, 24.0, 'west', 240.0),
    (7, 16.0, 'east', NULL),
    (8, 22.0, 'west', NULL)
  AS sales(store_id, advertising_spend, region, revenue)
)
SELECT store_id, revenue_prediction
FROM ai_predict_value(
  input => TABLE(sales),
  target_col => 'revenue',
  feature_cols => ARRAY('advertising_spend', 'region'),
  seed => 42
)
ORDER BY store_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 input must contain at least one row with a non-NULL target.
  • The input must not already contain a column named <target_col>_prediction.

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.

To predict a class label, use ai_predict_class function. For other AI Functions, see Transform unstructured data using AI Functions.

On this page