fix unit tests

Author: chelsea-lin

bigframes/dtypes.py

@@ -724,10 +724,6 @@ def infer_literal_type(literal) -> typing.Optional[Dtype]:
     # Maybe also normalize literal to canonical python representation to remove this burden from compilers?
     if isinstance(literal, pa.Scalar):
         return arrow_dtype_to_bigframes_dtype(literal.type)
-    if pd.api.types.is_list_like(literal):
-        element_types = [infer_literal_type(i) for i in literal]
-        common_type = lcd_type(*element_types)
-        return list_type(common_type)
     if pd.api.types.is_dict_like(literal):
         fields = []
         for key in literal.keys():
@@ -738,6 +734,10 @@ def infer_literal_type(literal) -> typing.Optional[Dtype]:
                 pa.field(key, field_type, nullable=(not pa.types.is_list(field_type)))
             )
         return pd.ArrowDtype(pa.struct(fields))
+    if pd.api.types.is_list_like(literal):
+        element_types = [infer_literal_type(i) for i in literal]
+        common_type = lcd_type(*element_types)
+        return list_type(common_type)
     if pd.isna(literal):
         return None  # Null value without a definite type
     # Make sure to check datetime before date as datetimes are also dates

tests/unit/bigquery/test_ai.py

@@ -91,7 +91,7 @@ def test_generate_embedding_with_dataframe(mock_dataframe, mock_session):
     expected_part_1 = "SELECT * FROM AI.GENERATE_EMBEDDING("
     expected_part_2 = f"MODEL `{model_name}`,"
     expected_part_3 = "(SELECT * FROM my_table),"
-    expected_part_4 = "STRUCT(256 AS OUTPUT_DIMENSIONALITY)"
+    expected_part_4 = "STRUCT(256 AS `OUTPUT_DIMENSIONALITY`)"
 
     assert expected_part_1 in query
     assert expected_part_2 in query
@@ -117,7 +117,7 @@ def test_generate_embedding_with_series(mock_embedding_series, mock_session):
     assert f"MODEL `{model_name}`" in query
     assert "(SELECT my_col AS content FROM my_table)" in query
     assert (
-        "STRUCT(0.0 AS START_SECOND, 10.0 AS END_SECOND, 5.0 AS INTERVAL_SECONDS)"
+        "STRUCT(0.0 AS `START_SECOND`, 10.0 AS `END_SECOND`, 5.0 AS `INTERVAL_SECONDS`)"
         in query
     )
 
@@ -180,7 +180,7 @@ def test_generate_text_with_dataframe(mock_dataframe, mock_session):
     expected_part_1 = "SELECT * FROM AI.GENERATE_TEXT("
     expected_part_2 = f"MODEL `{model_name}`,"
     expected_part_3 = "(SELECT * FROM my_table),"
-    expected_part_4 = "STRUCT(256 AS MAX_OUTPUT_TOKENS)"
+    expected_part_4 = "STRUCT(256 AS `MAX_OUTPUT_TOKENS`)"
 
     assert expected_part_1 in query
     assert expected_part_2 in query
@@ -238,7 +238,7 @@ def test_generate_table_with_dataframe(mock_dataframe, mock_session):
     expected_part_1 = "SELECT * FROM AI.GENERATE_TABLE("
     expected_part_2 = f"MODEL `{model_name}`,"
     expected_part_3 = "(SELECT * FROM my_table),"
-    expected_part_4 = "STRUCT('col1 STRING, col2 INT64' AS output_schema)"
+    expected_part_4 = "STRUCT('col1 STRING, col2 INT64' AS `output_schema`)"
 
     assert expected_part_1 in query
     assert expected_part_2 in query
@@ -264,7 +264,7 @@ def test_generate_table_with_options(mock_dataframe, mock_session):
     assert f"MODEL `{model_name}`" in query
     assert "(SELECT * FROM my_table)" in query
     assert (
-        "STRUCT('col1 STRING' AS output_schema, 0.5 AS temperature, 100 AS max_output_tokens)"
+        "STRUCT('col1 STRING' AS `output_schema`, 0.5 AS `temperature`, 100 AS `max_output_tokens`)"
         in query
     )
 
@@ -287,7 +287,7 @@ def test_generate_table_with_mapping_schema(mock_dataframe, mock_session):
     expected_part_1 = "SELECT * FROM AI.GENERATE_TABLE("
     expected_part_2 = f"MODEL `{model_name}`,"
     expected_part_3 = "(SELECT * FROM my_table),"
-    expected_part_4 = "STRUCT('col1 STRING, col2 INT64' AS output_schema)"
+    expected_part_4 = "STRUCT('col1 STRING, col2 INT64' AS `output_schema`)"
 
     assert expected_part_1 in query
     assert expected_part_2 in query

tests/unit/bigquery/test_ml.py

@@ -167,14 +167,14 @@ def test_generate_text_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mo
     assert "ML.GENERATE_TEXT" in generated_sql
     assert f"MODEL `{MODEL_NAME}`" in generated_sql
     assert "(SELECT * FROM `pandas_df`)" in generated_sql
-    assert "STRUCT(0.5 AS temperature" in generated_sql
-    assert "128 AS max_output_tokens" in generated_sql
-    assert "20 AS top_k" in generated_sql
-    assert "0.9 AS top_p" in generated_sql
-    assert "true AS flatten_json_output" in generated_sql
-    assert "['a', 'b'] AS stop_sequences" in generated_sql
-    assert "true AS ground_with_google_search" in generated_sql
-    assert "'TYPE' AS request_type" in generated_sql
+    assert "STRUCT(\n  0.5 AS `temperature`" in generated_sql
+    assert "128 AS `max_output_tokens`" in generated_sql
+    assert "20 AS `top_k`" in generated_sql
+    assert "0.9 AS `top_p`" in generated_sql
+    assert "TRUE AS `flatten_json_output`" in generated_sql
+    assert "['a', 'b'] AS `stop_sequences`" in generated_sql
+    assert "TRUE AS `ground_with_google_search`" in generated_sql
+    assert "'TYPE' AS `request_type`" in generated_sql
 
 
 @mock.patch("bigframes.pandas.read_gbq_query")
@@ -210,6 +210,6 @@ def test_generate_embedding_with_pandas_dataframe(
     assert "ML.GENERATE_EMBEDDING" in generated_sql
     assert f"MODEL `{MODEL_NAME}`" in generated_sql
     assert "(SELECT * FROM `pandas_df`)" in generated_sql
-    assert "true AS flatten_json_output" in generated_sql
-    assert "'RETRIEVAL_DOCUMENT' AS task_type" in generated_sql
-    assert "256 AS output_dimensionality" in generated_sql
+    assert "STRUCT(\n  TRUE AS `flatten_json_output`" in generated_sql
+    assert "'RETRIEVAL_DOCUMENT' AS `task_type`" in generated_sql
+    assert "256 AS `output_dimensionality`" in generated_sql

tests/unit/core/sql/snapshots/test_ml/test_evaluate_model_with_options/evaluate_model_with_options.sql

@@ -1 +1 @@
-SELECT * FROM ML.EVALUATE(MODEL `my_model`, STRUCT(false AS perform_aggregation, 10 AS horizon, 0.95 AS confidence_level))
+SELECT * FROM ML.EVALUATE(MODEL `my_model`, STRUCT(FALSE AS `perform_aggregation`, 10 AS `horizon`, 0.95 AS `confidence_level`))

tests/unit/core/sql/snapshots/test_ml/test_explain_predict_model_with_options/explain_predict_model_with_options.sql

@@ -1 +1 @@
-SELECT * FROM ML.EXPLAIN_PREDICT(MODEL `my_model`, (SELECT * FROM new_data), STRUCT(5 AS top_k_features))
+SELECT * FROM ML.EXPLAIN_PREDICT(MODEL `my_model`, (SELECT * FROM new_data), STRUCT(5 AS `top_k_features`))

tests/unit/core/sql/snapshots/test_ml/test_generate_embedding_model_with_options/generate_embedding_model_with_options.sql

@@ -1 +1,5 @@
-SELECT * FROM ML.GENERATE_EMBEDDING(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(true AS flatten_json_output, 'RETRIEVAL_DOCUMENT' AS task_type, 256 AS output_dimensionality))
+SELECT * FROM ML.GENERATE_EMBEDDING(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(
+  TRUE AS `flatten_json_output`,
+  'RETRIEVAL_DOCUMENT' AS `task_type`,
+  256 AS `output_dimensionality`
+))

tests/unit/core/sql/snapshots/test_ml/test_generate_text_model_with_options/generate_text_model_with_options.sql

@@ -1 +1,10 @@
-SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(0.5 AS temperature, 128 AS max_output_tokens, 20 AS top_k, 0.9 AS top_p, true AS flatten_json_output, ['a', 'b'] AS stop_sequences, true AS ground_with_google_search, 'TYPE' AS request_type))
+SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(
+  0.5 AS `temperature`,
+  128 AS `max_output_tokens`,
+  20 AS `top_k`,
+  0.9 AS `top_p`,
+  TRUE AS `flatten_json_output`,
+  ['a', 'b'] AS `stop_sequences`,
+  TRUE AS `ground_with_google_search`,
+  'TYPE' AS `request_type`
+))

tests/unit/core/sql/snapshots/test_ml/test_global_explain_model_with_options/global_explain_model_with_options.sql

@@ -1 +1 @@
-SELECT * FROM ML.GLOBAL_EXPLAIN(MODEL `my_model`, STRUCT(true AS class_level_explain))
+SELECT * FROM ML.GLOBAL_EXPLAIN(MODEL `my_model`, STRUCT(TRUE AS `class_level_explain`))

tests/unit/core/sql/snapshots/test_ml/test_predict_model_with_options/predict_model_with_options.sql

@@ -1 +1 @@
-SELECT * FROM ML.PREDICT(MODEL `my_model`, (SELECT * FROM new_data), STRUCT(true AS keep_original_columns))
+SELECT * FROM ML.PREDICT(MODEL `my_model`, (SELECT * FROM new_data), STRUCT(TRUE AS `keep_original_columns`))

tests/unit/ml/test_golden_sql.py

@@ -124,7 +124,7 @@ def test_linear_regression_default_fit(
     model.fit(mock_X, mock_y)
 
     mock_session._start_query_ml_ddl.assert_called_once_with(
-        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LINEAR_REG',\n  data_split_method='NO_SPLIT',\n  optimize_strategy='auto_strategy',\n  fit_intercept=True,\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=False,\n  enable_global_explain=False,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
+        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LINEAR_REG',\n  data_split_method='NO_SPLIT',\n  optimize_strategy='auto_strategy',\n  fit_intercept=TRUE,\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=FALSE,\n  enable_global_explain=FALSE,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
     )
 
 
@@ -134,7 +134,7 @@ def test_linear_regression_params_fit(bqml_model_factory, mock_session, mock_X,
     model.fit(mock_X, mock_y)
 
     mock_session._start_query_ml_ddl.assert_called_once_with(
-        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LINEAR_REG',\n  data_split_method='NO_SPLIT',\n  optimize_strategy='auto_strategy',\n  fit_intercept=False,\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=False,\n  enable_global_explain=False,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
+        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LINEAR_REG',\n  data_split_method='NO_SPLIT',\n  optimize_strategy='auto_strategy',\n  fit_intercept=FALSE,\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=FALSE,\n  enable_global_explain=FALSE,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
     )
 
 
@@ -169,7 +169,7 @@ def test_logistic_regression_default_fit(
     model.fit(mock_X, mock_y)
 
     mock_session._start_query_ml_ddl.assert_called_once_with(
-        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LOGISTIC_REG',\n  data_split_method='NO_SPLIT',\n  fit_intercept=True,\n  auto_class_weights=False,\n  optimize_strategy='auto_strategy',\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=False,\n  enable_global_explain=False,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql",
+        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LOGISTIC_REG',\n  data_split_method='NO_SPLIT',\n  fit_intercept=TRUE,\n  auto_class_weights=FALSE,\n  optimize_strategy='auto_strategy',\n  l2_reg=0.0,\n  max_iterations=20,\n  learn_rate_strategy='line_search',\n  min_rel_progress=0.01,\n  calculate_p_values=FALSE,\n  enable_global_explain=FALSE,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql",
     )
 
 
@@ -191,7 +191,7 @@ def test_logistic_regression_params_fit(
     model.fit(mock_X, mock_y)
 
     mock_session._start_query_ml_ddl.assert_called_once_with(
-        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LOGISTIC_REG',\n  data_split_method='NO_SPLIT',\n  fit_intercept=False,\n  auto_class_weights=True,\n  optimize_strategy='batch_gradient_descent',\n  l2_reg=0.2,\n  max_iterations=30,\n  learn_rate_strategy='constant',\n  min_rel_progress=0.02,\n  calculate_p_values=False,\n  enable_global_explain=False,\n  l1_reg=0.2,\n  learn_rate=0.2,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
+        "CREATE OR REPLACE MODEL `test-project`.`_anon123`.`temp_model_id`\nOPTIONS(\n  model_type='LOGISTIC_REG',\n  data_split_method='NO_SPLIT',\n  fit_intercept=FALSE,\n  auto_class_weights=TRUE,\n  optimize_strategy='batch_gradient_descent',\n  l2_reg=0.2,\n  max_iterations=30,\n  learn_rate_strategy='constant',\n  min_rel_progress=0.02,\n  calculate_p_values=FALSE,\n  enable_global_explain=FALSE,\n  l1_reg=0.2,\n  learn_rate=0.2,\n  INPUT_LABEL_COLS=['input_column_label'])\nAS input_X_y_no_index_sql"
     )