Accelerate performance for queries with AI functions

AI function acceleration lets you optimize performance when you run queries that use AI functions. It provides the same or better performance than array-based functions, without requiring you to write complex queries. For more information about AI functions, see AI functions overview.

AI function acceleration supports the following AI functions:

  • ai.if: Acceleration applies only to SELECT queries where the function is in the WHERE clause. It supports a single ai.if function per scan, and works with sequential, index, and bitmap heap scan. If a query contains multiple table scans—for example, on different tables, relations, or Common Table Expression (CTEs), then you can accelerate one ai.if function per scan.
  • ai.rank: Acceleration applies only to SELECT queries where the function is in the ORDER BY clause. It supports a single ai.rank function per scan. If a query contains multiple table scans—for example, on different tables, relations, or CTEs, then you can accelerate one ai.rank function per scan.

Enable AI function acceleration

To enable AI function acceleration, set the google_ml_integration.enable_ai_function_acceleration flag to on. This flag is disabled by default and controls whether the query executor is allowed to offload eligible AI operations to the AI function apply node for acceleration. When left unset or set to off, all query operations are executed by the standard PostgreSQL executor.

You can configure this flag at either the session or instance level.

To enable AI function acceleration at the instance level, use the gcloud alloydb instances update command:

gcloud alloydb instances update INSTANCE_ID \
  --database-flags google_ml_integration.enable_ai_function_acceleration=on \
  --region=REGION_ID \
  --cluster=CLUSTER_ID \
  --project=PROJECT_ID

Replace the following:

  • INSTANCE_ID: The ID of the instance you want to update.
  • REGION_ID: The ID of the region where the cluster is located. For more information, see Supported regions.
  • CLUSTER_ID: The ID of the cluster.
  • PROJECT_ID: The ID of the project.

For more information about how to set an instance-level flag, see Configure database flags.

Verify AI function acceleration

To verify that AI function acceleration is enabled, you can use the EXPLAIN statement to analyze your query.

Verify AI function acceleration for ai.if queries

The following example shows how to enable AI function acceleration for the current session and execute a query using the ai.if function:

SET google_ml_integration.enable_ai_function_acceleration = on;

EXPLAIN (COSTS OFF) SELECT r.name
FROM restaurant_reviews r
WHERE ai.if('Is the following a positive review? Review: '||r.review) AND r.city = 'Los Angeles'
GROUP BY r.name
HAVING count(*) > 20
ORDER BY r.name;

When you use the EXPLAIN statement to analyze your query, the query plan shows an AI Function Apply node:

                                       QUERY PLAN                                       
----------------------------------------------------------------------------------------
 GroupAggregate
   Group Key: name
   Filter: (count(*) > 20)
   ->  Sort
         Sort Key: name
         ->  AI Function Apply
               Filter: ai.if((('Is the following a positive review? Review: '::text || review)), NULL::character varying)
               ->  Index Scan using idx_restaurant_reviews_city on restaurant_reviews r
                     Index Cond: ((city)::text = 'Los Angeles'::text)

Verify AI function acceleration for ai.rank queries

The following example shows how to enable AI function acceleration for the current session and execute a query using the ai.rank function:

SET google_ml_integration.enable_ai_function_acceleration = on;

EXPLAIN (COSTS OFF) SELECT r.name, r.review
FROM restaurant_reviews r
WHERE r.city = 'Los Angeles'
ORDER BY ai.rank('Rank these reviews based on how much they emphasize high-quality ingredients. Review: ' || r.review) DESC
LIMIT 50;

When you use the EXPLAIN statement to analyze your query, the query plan shows an AI Function Apply node:

                                       QUERY PLAN                                       
----------------------------------------------------------------------------------------
 Limit
  ->  Sort
        Sort Key: (ai.rank(('Rank these reviews based on how much they emphasize high-quality ingredients. Review: ' || r.review'), NULL)) DESC
        ->  AI Function Apply
              ->  Bitmap Heap Scan on restaurant_reviews r
                    Recheck Cond: (city = 'Los Angeles')
                    ->  Bitmap Index Scan on idx_restaurant_reviews_city
                          Index Cond: (city = 'Los Angeles')

The AIFunctionApply node in the query plan indicates that the query used AI function acceleration. If this node isn't present, the query used standard PostgreSQL execution.

Limitations

  • Only basic arguments to the functions are supported. For ai.if and ai.rank, the prompt argument must be either a column reference or a string literal concatenated with a column reference. Any other arguments must be simple constants. Examples of supported prompt arguments include the following:
    • ai.if(r.review)
    • ai.if('Is this true? : ' || r.review)
    • ai.rank(r.review)
    • ai.rank('Rate this review: ' || r.review)
  • Only SELECT queries with ai.if in WHERE clauses or ai.rank in ORDER BY clauses are supported.
  • Only one AI function per table scan is supported.
  • Only the default Gemini models are supported. For more information, see Gemini models. To use a different model, you can set the google_ml_integration.default_llm_model flag to the model ID you want to use.
  • The feature is available only in regions that support the gemini-2.5-flash-lite model or the regions for the LLM model that you set using the google_ml_integration.default_llm_model flag. For more information, see Supported regions.

What's next