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