Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,21 @@
# limitations under the License.

import datetime
import functools
import unittest.mock as mock

import pytest
from google.cloud import bigquery

import bigframes.testing.mocks as mocks
from bigframes.testing import compiler_session

freezegun = pytest.importorskip("freezegun")

PROJECT_NAME = "bigframes-dev-perf"
DATASET_NAME = "tpch_0001t"
LOCATION_NAME = "test-region"

TPCH_SCHEMAS = {
"LINEITEM": [
bigquery.SchemaField("L_ORDERKEY", "INTEGER", mode="REQUIRED"),
Expand Down Expand Up @@ -103,46 +109,55 @@
}


def _create_mock_bqclient():
"""Helper function to create a compiler session."""

bqclient = mock.create_autospec(bigquery.Client, instance=True)
bqclient.project = DATASET_NAME
bqclient.location = LOCATION_NAME
table_create_time = datetime.datetime.now()

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 = {}
type(table).created = mock.PropertyMock(return_value=table_create_time)
type(table).location = mock.PropertyMock(return_value=LOCATION_NAME)
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=1000000000)
return table

bqclient.get_table.side_effect = get_table_mock
return bqclient


@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
anonymous_dataset = bigquery.DatasetReference.from_string(
f"{PROJECT_NAME}.{DATASET_NAME}"
)
session = mocks.create_bigquery_session(
bqclient=_create_mock_bqclient(),
anonymous_dataset=anonymous_dataset,
)

# Disable snapshotting for TPC-H tests to keep snapshots clean
original_read_gbq_table = session._loader.read_gbq_table

@functools.wraps(original_read_gbq_table)
def read_gbq_table_no_snapshot(*args, **kwargs):
kwargs["enable_snapshot"] = False
return original_read_gbq_table(*args, **kwargs)

session._loader.read_gbq_table = read_gbq_table_no_snapshot

session._executor = compiler_session.SQLCompilerExecutor()
return session
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ WITH `bfcte_0` AS (
) * (
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'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_0`
WHERE
`L_SHIPDATE` <= CAST('1998-09-02' AS DATE)
), `bfcte_1` AS (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,20 @@ WITH `bfcte_0` AS (
SELECT
`N_NATIONKEY` AS `bfcol_0`,
`N_NAME` AS `bfcol_1`
FROM `bigframes-dev`.`tpch`.`NATION` AS `bft_3` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`NATION` AS `bft_3`
), `bfcte_1` AS (
SELECT
`L_ORDERKEY` AS `bfcol_2`,
`L_EXTENDEDPRICE` AS `bfcol_3`,
`L_DISCOUNT` AS `bfcol_4`,
`L_RETURNFLAG` AS `bfcol_5`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_2` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_2`
), `bfcte_2` AS (
SELECT
`O_ORDERKEY` AS `bfcol_6`,
`O_CUSTKEY` AS `bfcol_7`,
`O_ORDERDATE` AS `bfcol_8`
FROM `bigframes-dev`.`tpch`.`ORDERS` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`ORDERS` AS `bft_1`
), `bfcte_3` AS (
SELECT
`C_CUSTKEY` AS `bfcol_9`,
Expand All @@ -25,7 +25,7 @@ WITH `bfcte_0` AS (
`C_PHONE` AS `bfcol_13`,
`C_ACCTBAL` AS `bfcol_14`,
`C_COMMENT` AS `bfcol_15`
FROM `bigframes-dev`.`tpch`.`CUSTOMER` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`CUSTOMER` AS `bft_0`
), `bfcte_4` AS (
SELECT
`bfcol_9` AS `bfcol_16`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,23 +7,23 @@ WITH `bfcte_0` AS (
`PS_SUPPKEY` AS `bfcol_0`,
`PS_AVAILQTY` AS `bfcol_1`,
`PS_SUPPLYCOST` AS `bfcol_2`
FROM `bigframes-dev`.`tpch`.`PARTSUPP` AS `bft_2` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PARTSUPP` AS `bft_2`
), `bfcte_2` AS (
SELECT
`PS_PARTKEY` AS `bfcol_10`,
`PS_SUPPKEY` AS `bfcol_11`,
`PS_AVAILQTY` AS `bfcol_12`,
`PS_SUPPLYCOST` AS `bfcol_13`
FROM `bigframes-dev`.`tpch`.`PARTSUPP` AS `bft_2` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PARTSUPP` AS `bft_2`
), `bfcte_3` AS (
SELECT
`S_SUPPKEY` AS `bfcol_3`,
`S_NATIONKEY` AS `bfcol_4`
FROM `bigframes-dev`.`tpch`.`SUPPLIER` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`SUPPLIER` AS `bft_1`
), `bfcte_4` AS (
SELECT
`N_NATIONKEY` AS `bfcol_18`
FROM `bigframes-dev`.`tpch`.`NATION` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`NATION` AS `bft_0`
WHERE
`N_NAME` = 'GERMANY'
), `bfcte_5` AS (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,12 @@ WITH `bfcte_0` AS (
`L_COMMITDATE` AS `bfcol_2`,
`L_RECEIPTDATE` AS `bfcol_3`,
`L_SHIPMODE` AS `bfcol_4`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_1`
), `bfcte_1` AS (
SELECT
`O_ORDERKEY` AS `bfcol_5`,
`O_ORDERPRIORITY` AS `bfcol_6`
FROM `bigframes-dev`.`tpch`.`ORDERS` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`ORDERS` AS `bft_0`
), `bfcte_2` AS (
SELECT
`bfcol_5`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@ WITH `bfcte_0` AS (
SELECT
`O_ORDERKEY` AS `bfcol_10`,
`O_CUSTKEY` AS `bfcol_11`
FROM `bigframes-dev`.`tpch`.`ORDERS` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`ORDERS` AS `bft_1`
WHERE
NOT (
REGEXP_CONTAINS(`O_COMMENT`, 'special.*requests')
)
), `bfcte_1` AS (
SELECT
`C_CUSTKEY` AS `bfcol_3`
FROM `bigframes-dev`.`tpch`.`CUSTOMER` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`CUSTOMER` AS `bft_0`
), `bfcte_2` AS (
SELECT
`bfcol_3`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,16 @@ WITH `bfcte_0` AS (
`L_EXTENDEDPRICE` AS `bfcol_1`,
`L_DISCOUNT` AS `bfcol_2`,
`L_SHIPDATE` AS `bfcol_3`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_1`
), `bfcte_3` AS (
SELECT
`P_PARTKEY` AS `bfcol_4`
FROM `bigframes-dev`.`tpch`.`PART` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PART` AS `bft_0`
), `bfcte_4` AS (
SELECT
`P_PARTKEY` AS `bfcol_8`,
`P_TYPE` AS `bfcol_9`
FROM `bigframes-dev`.`tpch`.`PART` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PART` AS `bft_0`
), `bfcte_5` AS (
SELECT
`bfcol_4`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ WITH `bfcte_0` AS (
`L_EXTENDEDPRICE` * (
1 - `L_DISCOUNT`
) AS `bfcol_24`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_1`
WHERE
(
`L_SHIPDATE` >= CAST('1996-01-01' AS DATE)
Expand All @@ -32,14 +32,14 @@ WITH `bfcte_0` AS (
), `bfcte_2` AS (
SELECT
`S_SUPPKEY` AS `bfcol_4`
FROM `bigframes-dev`.`tpch`.`SUPPLIER` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`SUPPLIER` AS `bft_0`
), `bfcte_3` AS (
SELECT
`S_SUPPKEY` AS `bfcol_8`,
`S_NAME` AS `bfcol_9`,
`S_ADDRESS` AS `bfcol_10`,
`S_PHONE` AS `bfcol_11`
FROM `bigframes-dev`.`tpch`.`SUPPLIER` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`SUPPLIER` AS `bft_0`
), `bfcte_4` AS (
SELECT
`bfcol_23`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ WITH `bfcte_0` AS (
NOT (
REGEXP_CONTAINS(`S_COMMENT`, 'Customer.*Complaints')
) AS `bfcol_9`
FROM `bigframes-dev`.`tpch`.`SUPPLIER` AS `bft_2` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`SUPPLIER` AS `bft_2`
WHERE
NOT (
REGEXP_CONTAINS(`S_COMMENT`, 'Customer.*Complaints')
Expand All @@ -15,14 +15,14 @@ WITH `bfcte_0` AS (
SELECT
`PS_PARTKEY` AS `bfcol_2`,
`PS_SUPPKEY` AS `bfcol_3`
FROM `bigframes-dev`.`tpch`.`PARTSUPP` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PARTSUPP` AS `bft_1`
), `bfcte_2` AS (
SELECT
`P_PARTKEY` AS `bfcol_4`,
`P_BRAND` AS `bfcol_5`,
`P_TYPE` AS `bfcol_6`,
`P_SIZE` AS `bfcol_7`
FROM `bigframes-dev`.`tpch`.`PART` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PART` AS `bft_0`
), `bfcte_3` AS (
SELECT
`bfcol_8`
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ WITH `bfcte_0` AS (
), `bfcte_1` AS (
SELECT
`P_PARTKEY` AS `bfcol_15`
FROM `bigframes-dev`.`tpch`.`PART` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PART` AS `bft_1`
WHERE
(
`P_BRAND` = 'Brand#23'
Expand All @@ -17,12 +17,12 @@ WITH `bfcte_0` AS (
`L_PARTKEY` AS `bfcol_3`,
`L_QUANTITY` AS `bfcol_4`,
`L_EXTENDEDPRICE` AS `bfcol_5`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_0`
), `bfcte_3` AS (
SELECT
`L_PARTKEY` AS `bfcol_6`,
`L_QUANTITY` AS `bfcol_7`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_0`
), `bfcte_4` AS (
SELECT
`bfcol_4` AS `bfcol_16`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,19 @@ WITH `bfcte_0` AS (
SELECT
`C_CUSTKEY` AS `bfcol_0`,
`C_NAME` AS `bfcol_1`
FROM `bigframes-dev`.`tpch`.`CUSTOMER` AS `bft_2` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`CUSTOMER` AS `bft_2`
), `bfcte_1` AS (
SELECT
`L_ORDERKEY` AS `bfcol_2`,
`L_QUANTITY` AS `bfcol_3`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_1`
), `bfcte_2` AS (
SELECT
`O_ORDERKEY` AS `bfcol_4`,
`O_CUSTKEY` AS `bfcol_5`,
`O_TOTALPRICE` AS `bfcol_6`,
`O_ORDERDATE` AS `bfcol_7`
FROM `bigframes-dev`.`tpch`.`ORDERS` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`ORDERS` AS `bft_0`
), `bfcte_3` AS (
SELECT
`bfcol_2`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,14 @@ WITH `bfcte_0` AS (
`L_DISCOUNT` AS `bfcol_4`,
`L_SHIPINSTRUCT` AS `bfcol_5`,
`L_SHIPMODE` AS `bfcol_6`
FROM `bigframes-dev`.`tpch`.`LINEITEM` AS `bft_1` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`LINEITEM` AS `bft_1`
), `bfcte_2` AS (
SELECT
`P_PARTKEY` AS `bfcol_7`,
`P_BRAND` AS `bfcol_8`,
`P_SIZE` AS `bfcol_9`,
`P_CONTAINER` AS `bfcol_10`
FROM `bigframes-dev`.`tpch`.`PART` AS `bft_0` FOR SYSTEM_TIME AS OF '2026-03-10T18:00:00'
FROM `bigframes-dev-perf`.`tpch_0001t`.`PART` AS `bft_0`
), `bfcte_3` AS (
SELECT
`bfcol_7`,
Expand Down
Loading
Loading