test: add golden sql tests for TPC-H queries

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

@@ -0,0 +1,148 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import unittest.mock as mock
+
+from google.cloud import bigquery
+import pytest
+
+import bigframes.testing.mocks as mocks
+
+freezegun = pytest.importorskip("freezegun")
+
+TPCH_SCHEMAS = {
+    "LINEITEM": [
+        bigquery.SchemaField("L_ORDERKEY", "INTEGER"),
+        bigquery.SchemaField("L_PARTKEY", "INTEGER"),
+        bigquery.SchemaField("L_SUPPKEY", "INTEGER"),
+        bigquery.SchemaField("L_LINENUMBER", "INTEGER"),
+        bigquery.SchemaField("L_QUANTITY", "FLOAT"),
+        bigquery.SchemaField("L_EXTENDEDPRICE", "FLOAT"),
+        bigquery.SchemaField("L_DISCOUNT", "FLOAT"),
+        bigquery.SchemaField("L_TAX", "FLOAT"),
+        bigquery.SchemaField("L_RETURNFLAG", "STRING"),
+        bigquery.SchemaField("L_LINESTATUS", "STRING"),
+        bigquery.SchemaField("L_SHIPDATE", "DATE"),
+        bigquery.SchemaField("L_COMMITDATE", "DATE"),
+        bigquery.SchemaField("L_RECEIPTDATE", "DATE"),
+        bigquery.SchemaField("L_SHIPINSTRUCT", "STRING"),
+        bigquery.SchemaField("L_SHIPMODE", "STRING"),
+        bigquery.SchemaField("L_COMMENT", "STRING"),
+    ],
+    "ORDERS": [
+        bigquery.SchemaField("O_ORDERKEY", "INTEGER"),
+        bigquery.SchemaField("O_CUSTKEY", "INTEGER"),
+        bigquery.SchemaField("O_ORDERSTATUS", "STRING"),
+        bigquery.SchemaField("O_TOTALPRICE", "FLOAT"),
+        bigquery.SchemaField("O_ORDERDATE", "DATE"),
+        bigquery.SchemaField("O_ORDERPRIORITY", "STRING"),
+        bigquery.SchemaField("O_CLERK", "STRING"),
+        bigquery.SchemaField("O_SHIPPRIORITY", "INTEGER"),
+        bigquery.SchemaField("O_COMMENT", "STRING"),
+    ],
+    "PART": [
+        bigquery.SchemaField("P_PARTKEY", "INTEGER"),
+        bigquery.SchemaField("P_NAME", "STRING"),
+        bigquery.SchemaField("P_MFGR", "STRING"),
+        bigquery.SchemaField("P_BRAND", "STRING"),
+        bigquery.SchemaField("P_TYPE", "STRING"),
+        bigquery.SchemaField("P_SIZE", "INTEGER"),
+        bigquery.SchemaField("P_CONTAINER", "STRING"),
+        bigquery.SchemaField("P_RETAILPRICE", "FLOAT"),
+        bigquery.SchemaField("P_COMMENT", "STRING"),
+    ],
+    "SUPPLIER": [
+        bigquery.SchemaField("S_SUPPKEY", "INTEGER"),
+        bigquery.SchemaField("S_NAME", "STRING"),
+        bigquery.SchemaField("S_ADDRESS", "STRING"),
+        bigquery.SchemaField("S_NATIONKEY", "INTEGER"),
+        bigquery.SchemaField("S_PHONE", "STRING"),
+        bigquery.SchemaField("S_ACCTBAL", "FLOAT"),
+        bigquery.SchemaField("S_COMMENT", "STRING"),
+    ],
+    "PARTSUPP": [
+        bigquery.SchemaField("PS_PARTKEY", "INTEGER"),
+        bigquery.SchemaField("PS_SUPPKEY", "INTEGER"),
+        bigquery.SchemaField("PS_AVAILQTY", "INTEGER"),
+        bigquery.SchemaField("PS_SUPPLYCOST", "FLOAT"),
+        bigquery.SchemaField("PS_COMMENT", "STRING"),
+    ],
+    "CUSTOMER": [
+        bigquery.SchemaField("C_CUSTKEY", "INTEGER"),
+        bigquery.SchemaField("C_NAME", "STRING"),
+        bigquery.SchemaField("C_ADDRESS", "STRING"),
+        bigquery.SchemaField("C_NATIONKEY", "INTEGER"),
+        bigquery.SchemaField("C_PHONE", "STRING"),
+        bigquery.SchemaField("C_ACCTBAL", "FLOAT"),
+        bigquery.SchemaField("C_MKTSEGMENT", "STRING"),
+        bigquery.SchemaField("C_COMMENT", "STRING"),
+    ],
+    "NATION": [
+        bigquery.SchemaField("N_NATIONKEY", "INTEGER"),
+        bigquery.SchemaField("N_NAME", "STRING"),
+        bigquery.SchemaField("N_REGIONKEY", "INTEGER"),
+        bigquery.SchemaField("N_COMMENT", "STRING"),
+    ],
+    "REGION": [
+        bigquery.SchemaField("R_REGIONKEY", "INTEGER"),
+        bigquery.SchemaField("R_NAME", "STRING"),
+        bigquery.SchemaField("R_COMMENT", "STRING"),
+    ],
+}
+
+
+@pytest.fixture(scope="session")
+def tpch_session():
+    from bigframes.testing import compiler_session
+
+    anonymous_dataset = bigquery.DatasetReference.from_string("bigframes-dev.tpch")
+    location = "us-central1"
+
+    with freezegun.freeze_time("2026-03-10 18:00:00"):
+        session = mocks.create_bigquery_session(
+            anonymous_dataset=anonymous_dataset,
+            location=location,
+        )
+
+        def get_table_mock(table_ref):
+            if isinstance(table_ref, str):
+                table_ref = bigquery.TableReference.from_string(table_ref)
+
+            table_id = table_ref.table_id
+            schema = TPCH_SCHEMAS.get(table_id, [])
+
+            table = mock.create_autospec(bigquery.Table, instance=True)
+            table._properties = {}
+            # mocks.create_bigquery_session's CURRENT_TIMESTAMP() returns offset-naive datetime.now()
+            # So we should also use offset-naive here to avoid comparison errors.
+            now = datetime.datetime.now()
+            type(table).schema = mock.PropertyMock(return_value=schema)
+            type(table).project = table_ref.project
+            type(table).dataset_id = table_ref.dataset_id
+            type(table).table_id = table_id
+            type(table).num_rows = mock.PropertyMock(return_value=1000000)
+            type(table).num_bytes = mock.PropertyMock(return_value=1000000)
+            type(table).location = mock.PropertyMock(return_value=location)
+            type(table).table_type = mock.PropertyMock(return_value="TABLE")
+            type(table).created = mock.PropertyMock(return_value=now)
+            type(table).modified = mock.PropertyMock(return_value=now)
+            type(table).range_partitioning = mock.PropertyMock(return_value=None)
+            type(table).time_partitioning = mock.PropertyMock(return_value=None)
+            type(table).clustering_fields = mock.PropertyMock(return_value=None)
+            return table
+
+        session.bqclient.get_table.side_effect = get_table_mock
+        session._executor = compiler_session.SQLCompilerExecutor()
+        return session

tests/unit/core/compile/sqlglot/tpch/snapshots/test_tpch/test_tpch_query/1/out.sql

@@ -0,0 +1,77 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `L_QUANTITY`,
+    `L_EXTENDEDPRICE`,
+    `L_DISCOUNT`,
+    `L_TAX`,
+    `L_RETURNFLAG`,
+    `L_LINESTATUS`,
+    `L_SHIPDATE`,
+    `L_QUANTITY` AS `bfcol_7`,
+    `L_EXTENDEDPRICE` AS `bfcol_8`,
+    `L_DISCOUNT` AS `bfcol_9`,
+    `L_TAX` AS `bfcol_10`,
+    `L_RETURNFLAG` AS `bfcol_11`,
+    `L_LINESTATUS` AS `bfcol_12`,
+    `L_SHIPDATE` <= CAST('1998-09-02' AS DATE) AS `bfcol_13`,
+    `L_QUANTITY` AS `bfcol_27`,
+    `L_EXTENDEDPRICE` AS `bfcol_28`,
+    `L_DISCOUNT` AS `bfcol_29`,
+    `L_TAX` AS `bfcol_30`,
+    `L_RETURNFLAG` AS `bfcol_31`,
+    `L_LINESTATUS` AS `bfcol_32`,
+    `L_EXTENDEDPRICE` * (
+      1.0 - `L_DISCOUNT`
+    ) AS `bfcol_33`,
+    `L_QUANTITY` AS `bfcol_41`,
+    `L_EXTENDEDPRICE` AS `bfcol_42`,
+    `L_DISCOUNT` AS `bfcol_43`,
+    `L_RETURNFLAG` AS `bfcol_44`,
+    `L_LINESTATUS` AS `bfcol_45`,
+    `L_EXTENDEDPRICE` * (
+      1.0 - `L_DISCOUNT`
+    ) AS `bfcol_46`,
+    (
+      `L_EXTENDEDPRICE` * (
+        1.0 - `L_DISCOUNT`
+      )
+    ) * (
+      1.0 + `L_TAX`
+    ) AS `bfcol_47`
+  FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
+  WHERE
+    `L_SHIPDATE` <= CAST('1998-09-02' AS DATE)
+), `bfcte_1` AS (
+  SELECT
+    `bfcol_44`,
+    `bfcol_45`,
+    COALESCE(SUM(`bfcol_41`), 0) AS `bfcol_55`,
+    COALESCE(SUM(`bfcol_42`), 0) AS `bfcol_56`,
+    COALESCE(SUM(`bfcol_46`), 0) AS `bfcol_57`,
+    COALESCE(SUM(`bfcol_47`), 0) AS `bfcol_58`,
+    AVG(`bfcol_41`) AS `bfcol_59`,
+    AVG(`bfcol_42`) AS `bfcol_60`,
+    AVG(`bfcol_43`) AS `bfcol_61`,
+    COUNT(`bfcol_41`) AS `bfcol_62`
+  FROM `bfcte_0`
+  WHERE
+    NOT `bfcol_44` IS NULL AND NOT `bfcol_45` IS NULL
+  GROUP BY
+    `bfcol_44`,
+    `bfcol_45`
+)
+SELECT
+  `bfcol_44` AS `L_RETURNFLAG`,
+  `bfcol_45` AS `L_LINESTATUS`,
+  `bfcol_55` AS `SUM_QTY`,
+  `bfcol_56` AS `SUM_BASE_PRICE`,
+  `bfcol_57` AS `SUM_DISC_PRICE`,
+  `bfcol_58` AS `SUM_CHARGE`,
+  `bfcol_59` AS `AVG_QTY`,
+  `bfcol_60` AS `AVG_PRICE`,
+  `bfcol_61` AS `AVG_DISC`,
+  `bfcol_62` AS `COUNT_ORDER`
+FROM `bfcte_1`
+ORDER BY
+  `bfcol_44` ASC NULLS LAST,
+  `bfcol_45` ASC NULLS LAST