chore: add scalars_types_df for compiler tests (#1831)

Author: chelsea-lin

bigframes/testing/mocks.py

@@ -41,6 +41,7 @@ def create_bigquery_session(
     bqclient: Optional[mock.Mock] = None,
     session_id: str = "abcxyz",
     table_schema: Sequence[google.cloud.bigquery.SchemaField] = TEST_SCHEMA,
+    table_name: str = "test_table",
     anonymous_dataset: Optional[google.cloud.bigquery.DatasetReference] = None,
     location: str = "test-region",
     ordering_mode: Literal["strict", "partial"] = "partial",
@@ -76,7 +77,7 @@ def create_bigquery_session(
         type(table).schema = mock.PropertyMock(return_value=table_schema)
         type(table).project = anonymous_dataset.project
         type(table).dataset_id = anonymous_dataset.dataset_id
-        type(table).table_id = "test_table"
+        type(table).table_id = table_name
         type(table).num_rows = mock.PropertyMock(return_value=1000000000)
         bqclient.get_table.return_value = table
 
@@ -94,7 +95,7 @@ def query_mock(
         query_job = mock.create_autospec(google.cloud.bigquery.QueryJob, instance=True)
         query_job._properties = {}
         type(query_job).destination = mock.PropertyMock(
-            return_value=anonymous_dataset.table("test_table"),
+            return_value=anonymous_dataset.table(table_name),
         )
         type(query_job).statement_type = mock.PropertyMock(return_value="SELECT")
 

tests/unit/core/compile/sqlglot/conftest.py

@@ -21,6 +21,7 @@
 import pytest
 
 from bigframes import dtypes
+import bigframes.pandas as bpd
 import bigframes.testing.mocks as mocks
 import bigframes.testing.utils
 
@@ -29,26 +30,52 @@
 
 
 @pytest.fixture(scope="session")
-def compiler_session(basic_types_table_schema):
+def compiler_session(scalars_types_table_schema):
     from bigframes.testing import compiler_session
 
     # TODO: Check if ordering mode is needed for the tests.
-    session = mocks.create_bigquery_session(table_schema=basic_types_table_schema)
+    table_name = "scalar_types"
+    anonymous_dataset = bigquery.DatasetReference.from_string(
+        "bigframes-dev.sqlglot_test"
+    )
+    session = mocks.create_bigquery_session(
+        table_name=table_name,
+        table_schema=scalars_types_table_schema,
+        anonymous_dataset=anonymous_dataset,
+    )
     session._executor = compiler_session.SQLCompilerExecutor()
     return session
 
 
 @pytest.fixture(scope="session")
-def basic_types_table_schema() -> typing.Sequence[bigquery.SchemaField]:
+def scalars_types_table_schema() -> typing.Sequence[bigquery.SchemaField]:
     return [
-        bigquery.SchemaField("rowindex", "INTEGER"),
+        bigquery.SchemaField("bool_col", "BOOLEAN"),
+        bigquery.SchemaField("bytes_col", "BYTES"),
+        bigquery.SchemaField("date_col", "DATE"),
+        bigquery.SchemaField("datetime_col", "DATETIME"),
+        bigquery.SchemaField("geography_col", "GEOGRAPHY"),
         bigquery.SchemaField("int64_col", "INTEGER"),
-        bigquery.SchemaField("string_col", "STRING"),
+        bigquery.SchemaField("int64_too", "INTEGER"),
+        bigquery.SchemaField("numeric_col", "NUMERIC"),
         bigquery.SchemaField("float64_col", "FLOAT"),
-        bigquery.SchemaField("bool_col", "BOOLEAN"),
+        bigquery.SchemaField("rowindex", "INTEGER"),
+        bigquery.SchemaField("rowindex_2", "INTEGER"),
+        bigquery.SchemaField("string_col", "STRING"),
+        bigquery.SchemaField("time_col", "TIME"),
+        bigquery.SchemaField("timestamp_col", "TIMESTAMP"),
     ]
 
 
+@pytest.fixture(scope="session")
+def scalars_types_df(compiler_session) -> bpd.DataFrame:
+    """Returns a BigFrames DataFrame containing all scalar types and using the `rowindex`
+    column as the index."""
+    bf_df = compiler_session.read_gbq_table("bigframes-dev.sqlglot_test.scalar_types")
+    bf_df = bf_df.set_index("rowindex", drop=False)
+    return bf_df
+
+
 @pytest.fixture(scope="session")
 def scalars_types_pandas_df() -> pd.DataFrame:
     """Returns a pandas DataFrame containing all scalar types and using the `rowindex`

tests/unit/core/compile/sqlglot/snapshots/test_compile_projection/test_compile_projection/out.sql

@@ -1,33 +1,26 @@
 WITH `bfcte_0` AS (
   SELECT
-    `rowindex` AS `bfcol_0`,
-    `int64_col` AS `bfcol_1`,
-    `string_col` AS `bfcol_2`,
-    `float64_col` AS `bfcol_3`,
-    `bool_col` AS `bfcol_4`
-  FROM `test-project`.`test_dataset`.`test_table`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    *,
-    `bfcol_0` AS `bfcol_5`,
-    `bfcol_2` AS `bfcol_6`,
-    `bfcol_3` AS `bfcol_7`,
-    `bfcol_4` AS `bfcol_8`,
-    `bfcol_1` + 1 AS `bfcol_9`
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
   FROM `bfcte_0`
 ), `bfcte_2` AS (
   SELECT
-    `bfcol_5` AS `bfcol_10`,
-    `bfcol_9` AS `bfcol_11`,
-    `bfcol_6` AS `bfcol_12`,
-    `bfcol_7` AS `bfcol_13`,
-    `bfcol_8` AS `bfcol_14`
+    *,
+    `bfcol_2` AS `bfcol_4`,
+    `bfcol_3` + 1 AS `bfcol_5`
   FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    `bfcol_4` AS `bfcol_6`,
+    `bfcol_5` AS `bfcol_7`
+  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_10` AS `rowindex`,
-  `bfcol_11` AS `int64_col`,
-  `bfcol_12` AS `string_col`,
-  `bfcol_13` AS `float64_col`,
-  `bfcol_14` AS `bool_col`
-FROM `bfcte_2`
+  `bfcol_6` AS `rowindex`,
+  `bfcol_7` AS `int64_col`
+FROM `bfcte_3`

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql

@@ -1,16 +1,53 @@
 WITH `bfcte_0` AS (
   SELECT
-    `rowindex` AS `bfcol_0`,
-    `int64_col` AS `bfcol_1`,
-    `string_col` AS `bfcol_2`,
-    `float64_col` AS `bfcol_3`,
-    `bool_col` AS `bfcol_4`
-  FROM `test-project`.`test_dataset`.`test_table`
+    `bool_col` AS `bfcol_0`,
+    `bytes_col` AS `bfcol_1`,
+    `date_col` AS `bfcol_2`,
+    `datetime_col` AS `bfcol_3`,
+    `geography_col` AS `bfcol_4`,
+    `int64_col` AS `bfcol_5`,
+    `int64_too` AS `bfcol_6`,
+    `numeric_col` AS `bfcol_7`,
+    `float64_col` AS `bfcol_8`,
+    `rowindex` AS `bfcol_9`,
+    `rowindex_2` AS `bfcol_10`,
+    `string_col` AS `bfcol_11`,
+    `time_col` AS `bfcol_12`,
+    `timestamp_col` AS `bfcol_13`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_9` AS `bfcol_14`,
+    `bfcol_0` AS `bfcol_15`,
+    `bfcol_1` AS `bfcol_16`,
+    `bfcol_2` AS `bfcol_17`,
+    `bfcol_3` AS `bfcol_18`,
+    `bfcol_4` AS `bfcol_19`,
+    `bfcol_5` AS `bfcol_20`,
+    `bfcol_6` AS `bfcol_21`,
+    `bfcol_7` AS `bfcol_22`,
+    `bfcol_8` AS `bfcol_23`,
+    `bfcol_9` AS `bfcol_24`,
+    `bfcol_10` AS `bfcol_25`,
+    `bfcol_11` AS `bfcol_26`,
+    `bfcol_12` AS `bfcol_27`,
+    `bfcol_13` AS `bfcol_28`
+  FROM `bfcte_0`
 )
 SELECT
-  `bfcol_0` AS `rowindex`,
-  `bfcol_1` AS `int64_col`,
-  `bfcol_2` AS `string_col`,
-  `bfcol_3` AS `float64_col`,
-  `bfcol_4` AS `bool_col`
-FROM `bfcte_0`
+  `bfcol_14` AS `rowindex`,
+  `bfcol_15` AS `bool_col`,
+  `bfcol_16` AS `bytes_col`,
+  `bfcol_17` AS `date_col`,
+  `bfcol_18` AS `datetime_col`,
+  `bfcol_19` AS `geography_col`,
+  `bfcol_20` AS `int64_col`,
+  `bfcol_21` AS `int64_too`,
+  `bfcol_22` AS `numeric_col`,
+  `bfcol_23` AS `float64_col`,
+  `bfcol_24` AS `rowindex_1`,
+  `bfcol_25` AS `rowindex_2`,
+  `bfcol_26` AS `string_col`,
+  `bfcol_27` AS `time_col`,
+  `bfcol_28` AS `timestamp_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_limit/out.sql

@@ -1,19 +1,18 @@
 WITH `bfcte_0` AS (
   SELECT
-    `rowindex` AS `bfcol_0`,
-    `int64_col` AS `bfcol_1`,
-    `string_col` AS `bfcol_2`,
-    `float64_col` AS `bfcol_3`,
-    `bool_col` AS `bfcol_4`
-  FROM `test-project`.`test_dataset`.`test_table`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
+  FROM `bfcte_0`
 )
 SELECT
-  `bfcol_0` AS `rowindex`,
-  `bfcol_1` AS `int64_col`,
-  `bfcol_2` AS `string_col`,
-  `bfcol_3` AS `float64_col`,
-  `bfcol_4` AS `bool_col`
-FROM `bfcte_0`
+  `bfcol_2` AS `rowindex`,
+  `bfcol_3` AS `int64_col`
+FROM `bfcte_1`
 ORDER BY
-  `bfcol_1` ASC NULLS LAST
+  `bfcol_2` ASC NULLS LAST
 LIMIT 10

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_ordering/out.sql

@@ -1,18 +1,17 @@
 WITH `bfcte_0` AS (
   SELECT
-    `rowindex` AS `bfcol_0`,
-    `int64_col` AS `bfcol_1`,
-    `string_col` AS `bfcol_2`,
-    `float64_col` AS `bfcol_3`,
-    `bool_col` AS `bfcol_4`
-  FROM `test-project`.`test_dataset`.`test_table`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
+  FROM `bfcte_0`
 )
 SELECT
-  `bfcol_0` AS `rowindex`,
-  `bfcol_1` AS `int64_col`,
-  `bfcol_2` AS `string_col`,
-  `bfcol_3` AS `float64_col`,
-  `bfcol_4` AS `bool_col`
-FROM `bfcte_0`
+  `bfcol_2` AS `rowindex`,
+  `bfcol_3` AS `int64_col`
+FROM `bfcte_1`
 ORDER BY
-  `bfcol_0` ASC NULLS LAST
+  `bfcol_3` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_scalar_expr/test_compile_numerical_add/out.sql

@@ -1,33 +1,26 @@
 WITH `bfcte_0` AS (
   SELECT
-    `rowindex` AS `bfcol_0`,
-    `int64_col` AS `bfcol_1`,
-    `string_col` AS `bfcol_2`,
-    `float64_col` AS `bfcol_3`,
-    `bool_col` AS `bfcol_4`
-  FROM `test-project`.`test_dataset`.`test_table`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    *,
-    `bfcol_0` AS `bfcol_5`,
-    `bfcol_2` AS `bfcol_6`,
-    `bfcol_3` AS `bfcol_7`,
-    `bfcol_4` AS `bfcol_8`,
-    `bfcol_1` + `bfcol_1` AS `bfcol_9`
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
   FROM `bfcte_0`
 ), `bfcte_2` AS (
   SELECT
-    `bfcol_5` AS `bfcol_10`,
-    `bfcol_9` AS `bfcol_11`,
-    `bfcol_6` AS `bfcol_12`,
-    `bfcol_7` AS `bfcol_13`,
-    `bfcol_8` AS `bfcol_14`
+    *,
+    `bfcol_2` AS `bfcol_4`,
+    `bfcol_3` + `bfcol_3` AS `bfcol_5`
   FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    `bfcol_4` AS `bfcol_6`,
+    `bfcol_5` AS `bfcol_7`
+  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_10` AS `rowindex`,
-  `bfcol_11` AS `int64_col`,
-  `bfcol_12` AS `string_col`,
-  `bfcol_13` AS `float64_col`,
-  `bfcol_14` AS `bool_col`
-FROM `bfcte_2`
+  `bfcol_6` AS `rowindex`,
+  `bfcol_7` AS `int64_col`
+FROM `bfcte_3`

tests/unit/core/compile/sqlglot/snapshots/test_compile_scalar_expr/test_compile_numerical_add_w_scalar/out.sql

@@ -0,0 +1,26 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    *,
+    `bfcol_2` AS `bfcol_4`,
+    `bfcol_3` + 1 AS `bfcol_5`
+  FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    `bfcol_4` AS `bfcol_6`,
+    `bfcol_5` AS `bfcol_7`
+  FROM `bfcte_2`
+)
+SELECT
+  `bfcol_6` AS `rowindex`,
+  `bfcol_7` AS `int64_col`
+FROM `bfcte_3`