import google.auth
import google.auth.exceptions
import ibis # noqa: F401
import pytest
from google.api_core.exceptions import NotFound
from google.cloud import bigquery
import ibis_bigquery
DEFAULT_PROJECT_ID = "ibis-gbq"
PROJECT_ID_ENV_VAR = "GOOGLE_BIGQUERY_PROJECT_ID"
DATASET_ID = "ibis_gbq_testing"
TESTING_DATA_URI = "https://raw.githubusercontent.com/ibis-project/testing-data/master"
bq = ibis_bigquery.Backend()
def pytest_addoption(parser):
parser.addoption(
"--save-dataset",
action="store_true",
default=False,
help="saves all test data in the testing dataset",
)
parser.addoption(
"--no-refresh-dataset",
action="store_true",
default=False,
help="do not refresh the test data in the testing dataset",
)
import os
import ibis
import ibis.expr.datatypes as dt
import pandas as pd
import pandas.testing as tm
import pytest
from pytest import param
import ibis_bigquery
from ibis_bigquery import udf # noqa: E402
PROJECT_ID = os.environ.get("GOOGLE_BIGQUERY_PROJECT_ID", "ibis-gbq")
DATASET_ID = "testing"
bq_backend = ibis_bigquery.Backend()
@pytest.fixture(scope="module")
def alltypes(client):
t = client.table("functional_alltypes")
expr = t[t.bigint_col.isin([10, 20])].limit(10)
return expr
@pytest.fixture(scope="module")
def df(alltypes):
return alltypes.execute()
def test_udf(client, alltypes, df):
def _array_index(translator, expr):
# SAFE_OFFSET returns NULL if out of bounds
return "{}[OFFSET({})]".format(*map(translator.translate, expr.op().args))
def _translate_pattern(translator, pattern):
# add 'r' to string literals to indicate to Cloud Spanner this is a raw string
return "r" * isinstance(pattern.op(), ops.Literal) + translator.translate(pattern)
def _regex_extract(translator, expr):
arg, pattern, index = expr.op().args
regex = _translate_pattern(translator, pattern)
result = "REGEXP_EXTRACT({}, {})".format(translator.translate(arg), regex)
return result
_operation_registry = bigquery_compiler._operation_registry.copy()
_operation_registry.update(
{ops.RegexExtract: _regex_extract, ops.ArrayIndex: _array_index,}
)
compiles = bigquery_compiler.BigQueryExprTranslator.compiles
rewrites = bigquery_compiler.BigQueryExprTranslator.rewrites
dialect = ibis_bigquery.Backend().dialect