deps: update to ibis-framework 9.x and newer sqlglot (#827)

* deps: update to ibis-framework 9.x and newer sqlglot

* update sqlglot and ibis

* bump minimum pandas

* bump pyarrow

* fix bfill and ffill

* nearly implement describe

* remove remaining reference to vendored_ibis_ops.ApproximateMultiQuantile

* support ToJsonString

* partial support for quantile

* fix inmemorytable

* fixed Series.explode

* nearly fix to_datetime

* remove tests I added

* patch for python 3.9 support

* fix unit tests

* fix explode with time type

* fix array_agg

* fix array_agg for asc order

* actually fix array_agg

* fix remote function

* fix in-memory nullable integer compilation

* fix test_df_construct_pandas_default on Python 3.9

* fix ShiftOp windows

* fix inf to SQL by treating values as literal in in memory table

* fix unit tests for ibis-framework 9.2.0

* fix Python 3.10 unit tests by syncing deps

* fixing remote function after merge

* fix visit_NonNullLiteral for int types

* visit_WindowFunction to fix s.median() method

* fix lint

* fix s.diff with window

* fix mypy

* patch visit_And to fix is_monotonic methods

* fix mypy and fillna warning

* undo window changes for test_series_autocorr

* undo fill_null because it was missed at 9.0 version

* vendor more of ibis for python 3.9 compatibility

* add default arg for nulls_first for python 3.9 support

* restore integer conversion

* fix window tests: diff, duplicated, shift

* fixing ibis parenthesize_inputs bugs and related tests

* fixing lint

* disable test_query_complexity_error

* fix doctest np.int64(0) upgrades

* fix doctest np.int64(0) upgrades more

* fix groupby diff

* addressing system-3.12/doctest issues related to numpy 2.1.1

* fix test_df_apply_axis_1_complex

* address compiler errors after merge

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* fix unit-test compile errors

* remove unused ibis codes

* fix fillna deprecated warning

* add _remove_null_ordering_from_unsupported_window back to fix test_precision_score etc ml tests

* fix is_monotonic_decreasing test

* fix explode after merge

* fix numpy on remote function test

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* ml numpy sql generations

---------

Co-authored-by: Chelsea Lin <[email protected]>
Co-authored-by: Chelsea Lin <[email protected]>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

bigframes/core/block_transforms.py

@@ -387,10 +387,9 @@ def value_counts(
 
 def pct_change(block: blocks.Block, periods: int = 1) -> blocks.Block:
     column_labels = block.column_labels
-    window_spec = windows.rows(
-        preceding=periods if periods > 0 else None,
-        following=-periods if periods < 0 else None,
-    )
+
+    # Window framing clause is not allowed for analytic function lag.
+    window_spec = windows.unbound()
 
     original_columns = block.value_columns
     block, shift_columns = block.multi_apply_window_op(

bigframes/core/blocks.py

@@ -50,7 +50,7 @@
 import bigframes.core.schema as bf_schema
 import bigframes.core.sql as sql
 import bigframes.core.utils as utils
-import bigframes.core.window_spec as window_specs
+import bigframes.core.window_spec as windows
 import bigframes.dtypes
 import bigframes.exceptions
 import bigframes.features
@@ -900,7 +900,7 @@ def multi_apply_window_op(
         self,
         columns: typing.Sequence[str],
         op: agg_ops.WindowOp,
-        window_spec: window_specs.WindowSpec,
+        window_spec: windows.WindowSpec,
         *,
         skip_null_groups: bool = False,
         never_skip_nulls: bool = False,
@@ -959,7 +959,7 @@ def apply_window_op(
         self,
         column: str,
         op: agg_ops.WindowOp,
-        window_spec: window_specs.WindowSpec,
+        window_spec: windows.WindowSpec,
         *,
         result_label: Label = None,
         skip_null_groups: bool = False,
@@ -1475,7 +1475,7 @@ def grouped_head(
         value_columns: typing.Sequence[str],
         n: int,
     ):
-        window_spec = window_specs.cumulative_rows(grouping_keys=tuple(by_column_ids))
+        window_spec = windows.cumulative_rows(grouping_keys=tuple(by_column_ids))
 
         block, result_id = self.apply_window_op(
             value_columns[0],
@@ -2383,10 +2383,7 @@ def _is_monotonic(
             return self._stats_cache[column_name][op_name]
 
         period = 1
-        window = window_specs.rows(
-            preceding=period,
-            following=None,
-        )
+        window_spec = windows.rows()
 
         # any NaN value means not monotonic
         block, last_notna_id = self.apply_unary_op(column_ids[0], ops.notnull_op)
@@ -2402,7 +2399,7 @@ def _is_monotonic(
         last_result_id = None
         for column_id in column_ids[::-1]:
             block, lag_result_id = block.apply_window_op(
-                column_id, agg_ops.ShiftOp(period), window
+                column_id, agg_ops.ShiftOp(period), window_spec
             )
             block, strict_monotonic_id = block.apply_binary_op(
                 column_id, lag_result_id, ops.gt_op if increasing else ops.lt_op

bigframes/core/compile/aggregate_compiler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import functools
 import typing
-from typing import cast, Optional
+from typing import cast, List, Optional
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.ibis.expr.operations as vendored_ibis_ops
@@ -31,6 +31,17 @@
 scalar_compiler = scalar_compilers.scalar_op_compiler
 
 
+# TODO(swast): We can remove this if ibis adds general approx_quantile
+# See: https://github.com/ibis-project/ibis/issues/9541
+@ibis.udf.agg.builtin
+def approx_quantiles(expression: float, number) -> List[float]:
+    """APPROX_QUANTILES
+
+    https://cloud.google.com/bigquery/docs/reference/standard-sql/approximate_aggregate_functions#approx_quantiles
+    """
+    return []  # pragma: NO COVER
+
+
 def compile_aggregate(
     aggregate: ex.Aggregation,
     bindings: typing.Dict[str, ibis_types.Value],
@@ -176,15 +187,12 @@ def _(
     column: ibis_types.NumericColumn,
     window=None,
 ) -> ibis_types.NumericValue:
-    # PERCENTILE_CONT has very few allowed windows. For example, "window
-    # framing clause is not allowed for analytic function percentile_cont".
+    # APPROX_QUANTILES has very few allowed windows.
     if window is not None:
         raise NotImplementedError(
             f"Approx Quartiles with windowing is not supported. {constants.FEEDBACK_LINK}"
         )
-    value = vendored_ibis_ops.ApproximateMultiQuantile(
-        column, num_bins=4  # type: ignore
-    ).to_expr()[op.quartile]
+    value = approx_quantiles(column, 4)[op.quartile]  # type: ignore
     return cast(ibis_types.NumericValue, value)
 
 
@@ -513,11 +521,15 @@ def _(
     column: ibis_types.Column,
     window=None,
 ) -> ibis_types.BooleanValue:
-    # BQ will return null for empty column, result would be true in pandas.
-    result = _is_true(column).all()
+    # BQ will return null for empty column, result would be false in pandas.
+    result = _apply_window_if_present(_is_true(column).all(), window)
+    literal = ibis_types.literal(True)
+
     return cast(
         ibis_types.BooleanScalar,
-        _apply_window_if_present(result, window).fillna(ibis_types.literal(True)),
+        result.fill_null(literal)
+        if hasattr(result, "fill_null")
+        else result.fillna(literal),
     )
 
 
@@ -528,10 +540,14 @@ def _(
     window=None,
 ) -> ibis_types.BooleanValue:
     # BQ will return null for empty column, result would be false in pandas.
-    result = _is_true(column).any()
+    result = _apply_window_if_present(_is_true(column).any(), window)
+    literal = ibis_types.literal(False)
+
     return cast(
         ibis_types.BooleanScalar,
-        _apply_window_if_present(result, window).fillna(ibis_types.literal(False)),
+        result.fill_null(literal)
+        if hasattr(result, "fill_null")
+        else result.fillna(literal),
     )
 
 

bigframes/core/compile/compiled.py

@@ -19,10 +19,9 @@
 import typing
 from typing import Collection, Literal, Optional, Sequence
 
-import bigframes_vendored.ibis.expr.operations as vendored_ibis_ops
+import bigframes_vendored.ibis.backends.bigquery.backend as ibis_bigquery
 import google.cloud.bigquery
 import ibis
-import ibis.backends.bigquery as ibis_bigquery
 import ibis.backends.bigquery.datatypes
 import ibis.common.deferred  # type: ignore
 import ibis.expr.datatypes as ibis_dtypes
@@ -407,18 +406,13 @@ def explode(self, offsets: typing.Sequence[int]) -> UnorderedIR:
 
         # The offset array ensures null represents empty arrays after unnesting.
         offset_array_id = bigframes.core.guid.generate_guid("offset_array_")
-        offset_array = (
-            vendored_ibis_ops.GenerateArray(
-                ibis.greatest(
-                    0,
-                    ibis.least(
-                        *[table[column_id].length() - 1 for column_id in column_ids]
-                    ),
-                )
-            )
-            .to_expr()
-            .name(offset_array_id),
-        )
+        offset_array = ibis.range(
+            0,
+            ibis.greatest(
+                1,  # We always want at least 1 element to fill in NULLs for empty arrays.
+                ibis.least(*[table[column_id].length() for column_id in column_ids]),
+            ),
+        ).name(offset_array_id)
         table_w_offset_array = table.select(
             offset_array,
             *self._column_names,
@@ -718,21 +712,13 @@ def explode(self, offsets: typing.Sequence[int]) -> OrderedIR:
         column_ids = tuple(table.columns[offset] for offset in offsets)
 
         offset_array_id = bigframes.core.guid.generate_guid("offset_array_")
-        offset_array = (
-            vendored_ibis_ops.GenerateArray(
-                ibis.greatest(
-                    0,
-                    ibis.least(
-                        *[
-                            table[table.columns[offset]].length() - 1
-                            for offset in offsets
-                        ]
-                    ),
-                )
-            )
-            .to_expr()
-            .name(offset_array_id),
-        )
+        offset_array = ibis.range(
+            0,
+            ibis.greatest(
+                1,  # We always want at least 1 element to fill in NULLs for empty arrays.
+                ibis.least(*[table[column_id].length() for column_id in column_ids]),
+            ),
+        ).name(offset_array_id)
         table_w_offset_array = table.select(
             offset_array,
             *self._column_names,
@@ -870,7 +856,7 @@ def project_window_op(
 
         clauses = []
         if op.skips_nulls and not never_skip_nulls:
-            clauses.append((column.isnull(), ibis.NA))
+            clauses.append((column.isnull(), ibis.null()))
         if window_spec.min_periods:
             if op.skips_nulls:
                 # Most operations do not count NULL values towards min_periods
@@ -891,7 +877,7 @@ def project_window_op(
             clauses.append(
                 (
                     observation_count < ibis_types.literal(window_spec.min_periods),
-                    ibis.NA,
+                    ibis.null(),
                 )
             )
         if clauses:
@@ -1322,9 +1308,10 @@ def _ibis_window_from_spec(
                     bounds.preceding, bounds.following, how="range"
                 )
             if isinstance(bounds, RowsWindowBounds):
-                window = window.preceding_following(
-                    bounds.preceding, bounds.following, how="rows"
-                )
+                if bounds.preceding is not None or bounds.following is not None:
+                    window = window.preceding_following(
+                        bounds.preceding, bounds.following, how="rows"
+                    )
             else:
                 raise ValueError(f"unrecognized window bounds {bounds}")
         return window

bigframes/core/compile/default_ordering.py

@@ -49,7 +49,12 @@ def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringVa
         # Needed for JSON, STRUCT and ARRAY datatypes
         result = vendored_ibis_ops.ToJsonString(column).to_expr()  # type: ignore
     # Escape backslashes and use backslash as delineator
-    escaped = cast(ibis_types.StringColumn, result.fillna("")).replace("\\", "\\\\")  # type: ignore
+    escaped = cast(
+        ibis_types.StringColumn,
+        result.fill_null("") if hasattr(result, "fill_null") else result.fillna(""),
+    ).replace(
+        "\\", "\\\\"
+    )  # type: ignore
     return cast(ibis_types.StringColumn, ibis.literal("\\")).concat(escaped)
 
 

bigframes/core/compile/scalar_op_compiler.py

@@ -842,7 +842,7 @@ def isin_op_impl(x: ibis_types.Value, op: ops.IsInOp):
 @scalar_op_compiler.register_unary_op(ops.ToDatetimeOp, pass_op=True)
 def to_datetime_op_impl(x: ibis_types.Value, op: ops.ToDatetimeOp):
     if x.type() == ibis_dtypes.str:
-        return vendored_ibis_ops.SafeCastToDatetime(x).to_expr()
+        return x.try_cast(ibis_dtypes.Timestamp(None))
     else:
         # Numerical inputs.
         if op.format:
@@ -995,8 +995,14 @@ def eq_nulls_match_op(
     y: ibis_types.Value,
 ):
     """Variant of eq_op where nulls match each other. Only use where dtypes are known to be same."""
-    left = x.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$"))
-    right = y.cast(ibis_dtypes.str).fillna(ibis_types.literal("$NULL_SENTINEL$"))
+    literal = ibis_types.literal("$NULL_SENTINEL$")
+    if hasattr(x, "fill_null"):
+        left = x.cast(ibis_dtypes.str).fill_null(literal)
+        right = y.cast(ibis_dtypes.str).fill_null(literal)
+    else:
+        left = x.cast(ibis_dtypes.str).fillna(literal)
+        right = y.cast(ibis_dtypes.str).fillna(literal)
+
     return left == right
 
 
@@ -1379,7 +1385,10 @@ def fillna_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
-    return x.fillna(typing.cast(ibis_types.Scalar, y))
+    if hasattr(x, "fill_null"):
+        return x.fill_null(typing.cast(ibis_types.Scalar, y))
+    else:
+        return x.fillna(typing.cast(ibis_types.Scalar, y))
 
 
 @scalar_op_compiler.register_binary_op(ops.round_op)

bigframes/core/compile/single_column.py

@@ -179,4 +179,8 @@ def value_to_join_key(value: ibis_types.Value):
     """Converts nullable values to non-null string SQL will not match null keys together - but pandas does."""
     if not value.type().is_string():
         value = value.cast(ibis_dtypes.str)
-    return value.fillna(ibis_types.literal("$NULL_SENTINEL$"))
+    return (
+        value.fill_null(ibis_types.literal("$NULL_SENTINEL$"))
+        if hasattr(value, "fill_null")
+        else value.fillna(ibis_types.literal("$NULL_SENTINEL$"))
+    )

bigframes/core/groupby/__init__.py

@@ -255,19 +255,17 @@ def cumprod(self, *args, **kwargs) -> df.DataFrame:
 
     @validations.requires_ordering()
     def shift(self, periods=1) -> series.Series:
-        window = window_specs.rows(
+        # Window framing clause is not allowed for analytic function lag.
+        window = window_specs.unbound(
             grouping_keys=tuple(self._by_col_ids),
-            preceding=periods if periods > 0 else None,
-            following=-periods if periods < 0 else None,
         )
         return self._apply_window_op(agg_ops.ShiftOp(periods), window=window)
 
     @validations.requires_ordering()
     def diff(self, periods=1) -> series.Series:
+        # Window framing clause is not allowed for analytic function lag.
         window = window_specs.rows(
             grouping_keys=tuple(self._by_col_ids),
-            preceding=periods if periods > 0 else None,
-            following=-periods if periods < 0 else None,
         )
         return self._apply_window_op(agg_ops.DiffOp(periods), window=window)
 
@@ -685,19 +683,16 @@ def cumcount(self, *args, **kwargs) -> series.Series:
     @validations.requires_ordering()
     def shift(self, periods=1) -> series.Series:
         """Shift index by desired number of periods."""
+        # Window framing clause is not allowed for analytic function lag.
         window = window_specs.rows(
             grouping_keys=tuple(self._by_col_ids),
-            preceding=periods if periods > 0 else None,
-            following=-periods if periods < 0 else None,
         )
         return self._apply_window_op(agg_ops.ShiftOp(periods), window=window)
 
     @validations.requires_ordering()
     def diff(self, periods=1) -> series.Series:
         window = window_specs.rows(
             grouping_keys=tuple(self._by_col_ids),
-            preceding=periods if periods > 0 else None,
-            following=-periods if periods < 0 else None,
         )
         return self._apply_window_op(agg_ops.DiffOp(periods), window=window)
 

bigframes/core/window_spec.py

@@ -70,7 +70,6 @@ def rows(
     Returns:
         WindowSpec
     """
-    assert (preceding is not None) or (following is not None)
     bounds = RowsWindowBounds(preceding=preceding, following=following)
     return WindowSpec(
         grouping_keys=grouping_keys,