Problem

Using CASE WHEN in TRANSFORM in a bigquery ML model, it breaks on deployment to vertexAI.

It is also broken when exporting the BQML model and loading it using tensorflow.

In particular, CASE WHEN can only handle 2 "WHEN" statements. If the 3rd "WHEN" would apply, it goes straight to the "ELSE" case instead (see Steps to reproduce)

What you expected to happen:

To be able to deploy a CASE WHEN with 3 or more "WHEN" cases. This is expected based on your docs here: https://cloud.google.com/bigquery/docs/exporting-models#export-transform-functions

Steps to reproduce:

Run this query to create a transform-only model (note the option to put it in vertex model registry):

CREATE OR REPLACE MODEL `my_project.my_dataset.transform_only_bug_report`
TRANSFORM(
-- comparing column to a number, will be tested using column1=3 later when deployed:
CASE
    WHEN column1 = 1 then 1 
    WHEN column1 = 2 then 2
    WHEN column1 = 3 then 3
    ELSE 0
END case_when_3_items_buggy -- this works in bigquery, but when deployed in vertex, the case column1==3 returns 0 instead

-- same but 2 WHENs, will be tested using column1=3 later when deployed:
, CASE
    WHEN column1 = 1 then 1 
    WHEN column1 = 3 then 3
    ELSE 0
END case_when_2_items_works -- this works in bigquery and when deployed

-- comparison of columns has the same problem, will be tested using column1=3 and column3=3
, CASE
WHEN column1 = 0 then 'zero'
WHEN column1 = column3 then 'equal'
ELSE 'unequal'
END comparison_2_items_works -- this works in bigquery and when deployed

, CASE
    WHEN column1 = 0 then 'zero'
    WHEN column3 = 0 then 'zero2'
    WHEN column1 = column3 then 'equal'
    ELSE 'unequal'
END comparison_3_items_buggy -- this works in bigquery, but when deployed in vertex, the case column1==column3 returns 0 instead

-- main use case for CASE WHEN is grouping of categoricals, will be tested using column2='string5':
, CASE 
    WHEN column2 in ('string1','string2') THEN 'groupA'
    WHEN column2 in ('string5','string6') THEN 'groupC'
    ELSE 'groupOther'
END grouping_2_items_works

, CASE 
    WHEN column2 in ('string1','string2') THEN 'groupA'
    WHEN column2 in ('string3','string4') THEN 'groupB'
    WHEN column2 in ('string5','string6') THEN 'groupC'
    ELSE 'groupOther'
END grouping_3_items_buggy
)
OPTIONS(
    model_type='TRANSFORM_ONLY'
   ,MODEL_REGISTRY = 'VERTEX_AI' 
   ,VERTEX_AI_MODEL_ID = 'trafo_bug_demo' 
   ,VERTEX_AI_MODEL_VERSION_ALIASES = ['bug_report']
   ) As

(
WITH test_data AS (
  SELECT * FROM UNNEST([
    STRUCT(1 AS column1, 'string1' AS column2, 11 AS column3),
    STRUCT(2 AS column1, 'string2' AS column2, 22 AS column3),
    STRUCT(3 AS column1, 'string3' AS column2, 33 AS column3),
    STRUCT(4 AS column1, 'string4' AS column2, 44 AS column3),
    STRUCT(5 AS column1, 'string5' AS column2, 55 AS column3),
    STRUCT(6 AS column1, 'string5' AS column2, 66 AS column3),
    STRUCT(7 AS column1, 'string5' AS column2, 77 AS column3)
  ])
)
SELECT * FROM test_data)

After creating the model, use vertex UI to create an endpoint. Then I used curl to send this data (following instructions in vertex UI):

{"instances": {"column1": 3, "column2": "string5", "column3": 3}, "parameters": {}}

Result is:

"predictions": [
    {
      "grouping_3_items_buggy": "groupOther",
      "case_when_2_items_works": 3,
      "comparison_2_items_works": "equal",
      "case_when_3_items_buggy": 0,
      "grouping_2_items_works": "groupC",
      "comparison_3_items_buggy": "unequal"
    }
]

Should be:

"predictions(expected)": [
    {
      "grouping_3_items_buggy": "groupC",
      "case_when_2_items_works": 3,
      "comparison_2_items_works": "equal",
      "case_when_3_items_buggy": 3,
      "grouping_2_items_works": "groupC",
      "comparison_3_items_buggy": "equal"
    }
]

Other information:

Same problem happens when exporting and running in a vertexAI workbench Details:

EXPORT MODEL `my_project.my_dataset.transform_only_bug_report_v2`
OPTIONS(URI = 'gs://{BUCKET_NAME}/model_exports/transform_only')

Then download to vertexAI workbench, load with tensorflow.saved_model.load() Then adjust inputs using signature of model. Exactly the same results. I also paid attention that the tensorflow version matches the model's "tensorflow_version". I used tensorflow-cpu.

SQL to run it in Bigquery (works there as expected):

WITH a AS (SELECT 3 AS column1, 'string5' AS column2, 3 AS column3)
SELECT *
FROM
  ML.TRANSFORM(
    MODEL `insparx-test.art_adhoc.transform_only_bug_report_v2`,
    (SELECT * FROM a))

Use case of "TRANSFORM" for our company:

We have analysts that are very good at creating models and doing feature engineering but know too little python to write deployable preprocessing in python. The "TRANSFORM" in BQML would be a huge benefit to our team!