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/compile/sqlglot/sql/test_base.py

@@ -14,6 +14,7 @@
 
 import datetime
 import decimal
+import re
 
 import numpy as np
 import pandas as pd
@@ -59,9 +60,6 @@
             "CAST('2025-01-02T03:45:06.789123+00:00' AS TIMESTAMP)",
             id="timestamp",
         ),
-        pytest.param(
-            shapely.geometry.Point(0, 1), "ST_GEOGFROMTEXT('POINT (0 1)')", id="geo"
-        ),
         pytest.param(np.int64(123), "123", id="np_int64"),
         pytest.param(np.float64(123.75), "123.75", id="np_float64"),
         pytest.param(float("inf"), "CAST('Infinity' AS FLOAT64)", id="inf"),
@@ -77,6 +75,13 @@ def test_literal(value, expected_pattern):
     assert got == expected_pattern
 
 
+def test_literal_for_geo():
+    value = shapely.geometry.Point(0, 1)
+    expected_pattern = r"ST_GEOGFROMTEXT\('POINT \(0[.]?0* 1[.]?0*\)'\)"
+    got = sql.to_sql(sql.literal(value))
+    assert re.match(expected_pattern, got) is not None
+
+
 @pytest.mark.parametrize(
     ("value", "dtype", "expected"),
     (

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`))