feat(bigquery): move bigquery backend back into the main repo

Author: cpcloud

.github/workflows/ibis-snowflake.yml renamed to .github/workflows/ibis-backends-cloud.yml

@@ -1,5 +1,5 @@
 # vim: filetype=yaml
-name: Snowflake Backend
+name: Cloud Backends
 
 on:
   push:
@@ -33,6 +33,8 @@ jobs:
         backend:
           - name: snowflake
             title: Snowflake
+          - name: bigquery
+            title: BigQuery
     steps:
       - name: checkout
         uses: actions/checkout@v3
@@ -46,6 +48,7 @@ jobs:
       - run: python -m pip install --upgrade pip 'poetry>=1.2'
 
       - name: set a compatible pyarrow version
+        if: ${{ matrix.backend.name == 'snowflake' }}
         run: poetry add pyarrow@'>=8.0.0,<8.1.0' --optional
 
       - name: install ibis
@@ -58,11 +61,20 @@ jobs:
       - name: download backend data
         run: just download-data
 
-      - name: "run parallel tests: ${{ matrix.backend.name }}"
-        run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup
+      - uses: google-github-actions/auth@v1
+        if: ${{ matrix.backend.name == 'bigquery' }}
+        with:
+          credentials_json: ${{ secrets.GCP_CREDENTIALS }}
+
+      - name: setup snowflake credentials
+        if: ${{ matrix.backend.name == 'snowflake' }}
+        run: echo "SNOWFLAKE_URL=${SNOWFLAKE_URL}" >> "$GITHUB_ENV"
         env:
           SNOWFLAKE_URL: ${{ secrets.SNOWFLAKE_URL }}
 
+      - name: "run parallel tests: ${{ matrix.backend.name }}"
+        run: just ci-check -m ${{ matrix.backend.name }} --numprocesses auto --dist=loadgroup
+
       - name: upload code coverage
         if: success()
         uses: codecov/codecov-action@v3

ibis/backends/base/sql/alchemy/__init__.py

@@ -142,16 +142,11 @@ def _to_geodataframe(df, schema):
         return df
 
     def fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
-
         import pandas as pd
 
-        df = pd.DataFrame.from_records(
-            cursor,
-            columns=cursor.keys(),
-            coerce_float=True,
-        )
+        df = pd.DataFrame.from_records(cursor, columns=schema.names, coerce_float=True)
         df = schema.apply_to(df)
-        if len(df) and geospatial_supported:
+        if not df.empty and geospatial_supported:
             return self._to_geodataframe(df, schema)
         return df
 

ibis/backends/base/sql/alchemy/registry.py

@@ -332,14 +332,6 @@ def _window(t, op):
 
     window_op = arg
 
-    _require_order_by = (
-        ops.DenseRank,
-        ops.MinRank,
-        ops.NTile,
-        ops.PercentRank,
-        ops.CumeDist,
-    )
-
     if isinstance(window_op, ops.CumulativeOp):
         arg = _cumulative_to_window(t, arg, window).op()
         return t.translate(arg)
@@ -352,29 +344,18 @@ def _window(t, op):
 
     # Some analytic functions need to have the expression of interest in
     # the ORDER BY part of the window clause
-    if isinstance(window_op, _require_order_by) and not window._order_by:
+    if isinstance(window_op, t._require_order_by) and not window._order_by:
         order_by = t.translate(window_op.args[0])
     else:
         order_by = [t.translate(arg) for arg in window._order_by]
 
     partition_by = [t.translate(arg) for arg in window._group_by]
 
-    frame_clause_not_allowed = (
-        ops.Lag,
-        ops.Lead,
-        ops.DenseRank,
-        ops.MinRank,
-        ops.NTile,
-        ops.PercentRank,
-        ops.CumeDist,
-        ops.RowNumber,
-    )
-
     how = {'range': 'range_'}.get(window.how, window.how)
     preceding = window.preceding
     additional_params = (
         {}
-        if isinstance(window_op, frame_clause_not_allowed)
+        if t._forbids_frame_clause and isinstance(window_op, t._forbids_frame_clause)
         else {
             how: (
                 -preceding if preceding is not None else preceding,

ibis/backends/base/sql/alchemy/translator.py

@@ -33,7 +33,6 @@ def subcontext(self):
 
 
 class AlchemyExprTranslator(ExprTranslator):
-
     _registry = sqlalchemy_operation_registry
     _rewrites = ExprTranslator._rewrites.copy()
     _type_map = ibis_type_to_sqla
@@ -47,6 +46,14 @@ class AlchemyExprTranslator(ExprTranslator):
     native_json_type = True
     _always_quote_columns = False
 
+    _require_order_by = (
+        ops.DenseRank,
+        ops.MinRank,
+        ops.NTile,
+        ops.PercentRank,
+        ops.CumeDist,
+    )
+
     def name(self, translated, name, force=True):
         return translated.label(name)
 

ibis/backends/base/sql/compiler/translator.py

@@ -162,6 +162,25 @@ class ExprTranslator:
 
     _registry = operation_registry
     _rewrites: dict[ops.Node, Callable] = {}
+    _forbids_frame_clause = (
+        ops.DenseRank,
+        ops.MinRank,
+        ops.NTile,
+        ops.PercentRank,
+        ops.CumeDist,
+        ops.RowNumber,
+    )
+    _require_order_by = (
+        ops.Lag,
+        ops.Lead,
+        ops.DenseRank,
+        ops.MinRank,
+        ops.FirstValue,
+        ops.LastValue,
+        ops.PercentRank,
+        ops.CumeDist,
+        ops.NTile,
+    )
 
     def __init__(self, node, context, named=False, permit_subquery=False):
         self.node = node

ibis/backends/base/sql/registry/window.py

@@ -1,5 +1,5 @@
 from operator import add, mul, sub
-from typing import Optional, Union
+from typing import Optional
 
 import ibis
 import ibis.common.exceptions as com
@@ -41,7 +41,7 @@
 }
 
 
-def _replace_interval_with_scalar(expr: Union[ir.Expr, dt.Interval, float]):
+def _replace_interval_with_scalar(op: ops.Node):
     """Good old Depth-First Search to identify the Interval and IntervalValue
     components of the expression and return a comparable scalar expression.
 
@@ -54,37 +54,22 @@ def _replace_interval_with_scalar(expr: Union[ir.Expr, dt.Interval, float]):
     -------
     preceding : float or ir.FloatingScalar, depending upon the expr
     """
-    if isinstance(expr, ir.Expr):
-        expr_op = expr.op()
-    else:
-        expr_op = None
-
-    # TODO: fix this to use to_interval if possible
-    if not isinstance(expr, (dt.Interval, ir.IntervalValue)):
-        # Literal expressions have op method but native types do not.
-        if isinstance(expr_op, ops.Literal):
-            return expr_op.value
-        else:
-            return expr
-    elif isinstance(expr, dt.Interval):
+    if isinstance(op, ops.Literal):
+        unit = getattr(op.output_dtype, "unit", "us")
         try:
-            microseconds = _map_interval_to_microseconds[expr.unit]
-            return microseconds
+            micros = _map_interval_to_microseconds[unit]
+            return op.value * micros
         except KeyError:
-            raise ValueError(
-                "Expected preceding values of week(), "
-                + "day(), hour(), minute(), second(), millisecond(), "
-                + f"microseconds(), nanoseconds(); got {expr}"
-            )
-    elif expr_op.args and isinstance(expr, ir.IntervalValue):
-        if len(expr_op.args) > 2:
+            raise ValueError(f"Unsupported unit {unit!r}")
+    elif op.args and isinstance(op.output_dtype, dt.Interval):
+        if len(op.args) > 2:
             raise NotImplementedError("'preceding' argument cannot be parsed.")
-        left_arg = _replace_interval_with_scalar(expr_op.args[0])
-        right_arg = _replace_interval_with_scalar(expr_op.args[1])
-        method = _map_interval_op_to_op[type(expr_op)]
+        left_arg = _replace_interval_with_scalar(op.args[0])
+        right_arg = _replace_interval_with_scalar(op.args[1])
+        method = _map_interval_op_to_op[type(op)]
         return method(left_arg, right_arg)
     else:
-        raise TypeError(f'expr has unknown type {type(expr).__name__}')
+        raise TypeError(f'input has unknown type {type(op)}')
 
 
 def cumulative_to_window(translator, op, window):
@@ -105,20 +90,20 @@ def cumulative_to_window(translator, op, window):
 
 def time_range_to_range_window(_, window):
     # Check that ORDER BY column is a single time column:
-    order_by_vars = [x.op().args[0] for x in window._order_by]
+    order_by_vars = [x.args[0] for x in window._order_by]
     if len(order_by_vars) > 1:
         raise com.IbisInputError(
             f"Expected 1 order-by variable, got {len(order_by_vars)}"
         )
 
-    order_var = window._order_by[0].op().args[0]
-    timestamp_order_var = order_var.cast('int64')
+    order_var = order_by_vars[0]
+    timestamp_order_var = ops.Cast(order_var, dt.int64).to_expr()
     window = window._replace(order_by=timestamp_order_var, how='range')
 
     # Need to change preceding interval expression to scalars
     preceding = window.preceding
     if isinstance(preceding, ir.IntervalScalar):
-        new_preceding = _replace_interval_with_scalar(preceding)
+        new_preceding = _replace_interval_with_scalar(preceding.op())
         window = window._replace(preceding=new_preceding)
 
     return window
@@ -165,22 +150,12 @@ def _foll(f: Optional[int]) -> str:
 
         return f'{prefix} FOLLOWING'
 
-    frame_clause_not_allowed = (
-        ops.Lag,
-        ops.Lead,
-        ops.DenseRank,
-        ops.MinRank,
-        ops.NTile,
-        ops.PercentRank,
-        ops.CumeDist,
-        ops.RowNumber,
-    )
-
-    if isinstance(op.expr, frame_clause_not_allowed):
+    if translator._forbids_frame_clause and isinstance(
+        op.expr, translator._forbids_frame_clause
+    ):
         frame = None
     elif p is not None and f is not None:
         frame = f'{window.how.upper()} BETWEEN {_prec(p)} AND {_foll(f)}'
-
     elif p is not None:
         if isinstance(p, tuple):
             start, end = p
@@ -200,7 +175,6 @@ def _foll(f: Optional[int]) -> str:
             kind = 'ROWS' if f > 0 else 'RANGE'
             frame = f'{kind} BETWEEN UNBOUNDED PRECEDING AND {_foll(f)}'
     else:
-        # no-op, default is full sample
         frame = None
 
     if frame is not None:
@@ -223,18 +197,6 @@ def _foll(f: Optional[int]) -> str:
 def window(translator, op):
     arg, window = op.args
 
-    _require_order_by = (
-        ops.Lag,
-        ops.Lead,
-        ops.DenseRank,
-        ops.MinRank,
-        ops.FirstValue,
-        ops.LastValue,
-        ops.PercentRank,
-        ops.CumeDist,
-        ops.NTile,
-    )
-
     _unsupported_reductions = (
         ops.ApproxMedian,
         ops.GroupConcat,
@@ -252,15 +214,18 @@ def window(translator, op):
 
     # Some analytic functions need to have the expression of interest in
     # the ORDER BY part of the window clause
-    if isinstance(arg, _require_order_by) and not window._order_by:
+    if isinstance(arg, translator._require_order_by) and not window._order_by:
         window = window.order_by(arg.args[0])
 
     # Time ranges need to be converted to microseconds.
     # FIXME(kszucs): avoid the expression roundtrip
     if window.how == 'range':
-        order_by_types = [type(x.op().args[0]) for x in window._order_by]
-        time_range_types = (ir.TimeColumn, ir.DateColumn, ir.TimestampColumn)
-        if any(col_type in time_range_types for col_type in order_by_types):
+        time_range_types = (dt.Time, dt.Date, dt.Timestamp)
+        if any(
+            isinstance(c.output_dtype, time_range_types)
+            and c.output_shape.is_columnar()
+            for c in window._order_by
+        ):
             window = time_range_to_range_window(translator, window)
 
     window_formatted = format_window(translator, op, window)