Merge branch 'main' into shuowei-anywidget-html-repr

Author: shuoweil

.github/workflows/js-tests.yml

@@ -0,0 +1,20 @@
+name: js-tests
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+    - name: Install modules
+      working-directory: ./tests/js
+      run: npm install
+    - name: Run tests
+      working-directory: ./tests/js
+      run: npm test

.gitignore

@@ -58,6 +58,7 @@ coverage.xml
 
 # System test environment variables.
 system_tests/local_test_setup
+tests/js/node_modules/
 
 # Make sure a generated file isn't accidentally committed.
 pylintrc

bigframes/core/compile/sqlglot/expressions/constants.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
 import sqlglot.expressions as sge
 
 _ZERO = sge.Cast(this=sge.convert(0), to="INT64")
@@ -23,3 +25,13 @@
 # FLOAT64 has 11 exponent bits, so max values is about 2**(2**10)
 # ln(2**(2**10)) == (2**10)*ln(2) ~= 709.78, so EXP(x) for x>709.78 will overflow.
 _FLOAT64_EXP_BOUND = sge.convert(709.78)
+
+# The natural logarithm of the maximum value for a signed 64-bit integer.
+# This is used to check for potential overflows in power operations involving integers
+# by checking if `exponent * log(base)` exceeds this value.
+_INT64_LOG_BOUND = math.log(2**63 - 1)
+
+# Represents the largest integer N where all integers from -N to N can be
+# represented exactly as a float64. Float64 types have a 53-bit significand precision,
+# so integers beyond this value may lose precision.
+_FLOAT64_MAX_INT_PRECISION = 2**53

bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -210,6 +210,141 @@ def _(expr: TypedExpr) -> sge.Expression:
     return expr.expr
 
 
+@register_binary_op(ops.pow_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    if left.dtype == dtypes.INT_DTYPE and right.dtype == dtypes.INT_DTYPE:
+        return _int_pow_op(left_expr, right_expr)
+    else:
+        return _float_pow_op(left_expr, right_expr)
+
+
+def _int_pow_op(
+    left_expr: sge.Expression, right_expr: sge.Expression
+) -> sge.Expression:
+    overflow_cond = sge.and_(
+        sge.NEQ(this=left_expr, expression=sge.convert(0)),
+        sge.GT(
+            this=sge.Mul(
+                this=right_expr, expression=sge.Ln(this=sge.Abs(this=left_expr))
+            ),
+            expression=sge.convert(constants._INT64_LOG_BOUND),
+        ),
+    )
+
+    return sge.Case(
+        ifs=[
+            sge.If(
+                this=overflow_cond,
+                true=sge.Null(),
+            )
+        ],
+        default=sge.Cast(
+            this=sge.Pow(
+                this=sge.Cast(
+                    this=left_expr, to=sge.DataType(this=sge.DataType.Type.DECIMAL)
+                ),
+                expression=right_expr,
+            ),
+            to="INT64",
+        ),
+    )
+
+
+def _float_pow_op(
+    left_expr: sge.Expression, right_expr: sge.Expression
+) -> sge.Expression:
+    # Most conditions here seek to prevent calling BQ POW with inputs that would generate errors.
+    # See: https://cloud.google.com/bigquery/docs/reference/standard-sql/mathematical_functions#pow
+    overflow_cond = sge.and_(
+        sge.NEQ(this=left_expr, expression=constants._ZERO),
+        sge.GT(
+            this=sge.Mul(
+                this=right_expr, expression=sge.Ln(this=sge.Abs(this=left_expr))
+            ),
+            expression=constants._FLOAT64_EXP_BOUND,
+        ),
+    )
+
+    # Float64 lose integer precision beyond 2**53, beyond this insufficient precision to get parity
+    exp_too_big = sge.GT(
+        this=sge.Abs(this=right_expr),
+        expression=sge.convert(constants._FLOAT64_MAX_INT_PRECISION),
+    )
+    # Treat very large exponents as +=INF
+    norm_exp = sge.Case(
+        ifs=[
+            sge.If(
+                this=exp_too_big,
+                true=sge.Mul(this=constants._INF, expression=sge.Sign(this=right_expr)),
+            )
+        ],
+        default=right_expr,
+    )
+
+    pow_result = sge.Pow(this=left_expr, expression=norm_exp)
+
+    # This cast is dangerous, need to only excuted where y_val has been bounds-checked
+    # Ibis needs try_cast binding to bq safe_cast
+    exponent_is_whole = sge.EQ(
+        this=sge.Cast(this=right_expr, to="INT64"), expression=right_expr
+    )
+    odd_exponent = sge.and_(
+        sge.LT(this=left_expr, expression=constants._ZERO),
+        sge.EQ(
+            this=sge.Mod(
+                this=sge.Cast(this=right_expr, to="INT64"), expression=sge.convert(2)
+            ),
+            expression=sge.convert(1),
+        ),
+    )
+    infinite_base = sge.EQ(this=sge.Abs(this=left_expr), expression=constants._INF)
+
+    return sge.Case(
+        ifs=[
+            # Might be able to do something more clever with x_val==0 case
+            sge.If(
+                this=sge.EQ(this=right_expr, expression=constants._ZERO),
+                true=sge.convert(1),
+            ),
+            sge.If(
+                this=sge.EQ(this=left_expr, expression=sge.convert(1)),
+                true=sge.convert(1),
+            ),  # Need to ignore exponent, even if it is NA
+            sge.If(
+                this=sge.and_(
+                    sge.EQ(this=left_expr, expression=constants._ZERO),
+                    sge.LT(this=right_expr, expression=constants._ZERO),
+                ),
+                true=constants._INF,
+            ),  # This case would error POW function in BQ
+            sge.If(this=infinite_base, true=pow_result),
+            sge.If(
+                this=exp_too_big, true=pow_result
+            ),  # Bigquery can actually handle the +-inf cases gracefully
+            sge.If(
+                this=sge.and_(
+                    sge.LT(this=left_expr, expression=constants._ZERO),
+                    sge.Not(this=exponent_is_whole),
+                ),
+                true=constants._NAN,
+            ),
+            sge.If(
+                this=overflow_cond,
+                true=sge.Mul(
+                    this=constants._INF,
+                    expression=sge.Case(
+                        ifs=[sge.If(this=odd_exponent, true=sge.convert(-1))],
+                        default=sge.convert(1),
+                    ),
+                ),
+            ),  # finite overflows would cause bq to error
+        ],
+        default=pow_result,
+    )
+
+
 @register_unary_op(ops.sqrt_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.Case(

bigframes/display/anywidget.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 from importlib import resources
 import functools
 import math
@@ -28,6 +29,7 @@
 from bigframes.core import blocks
 import bigframes.dataframe
 import bigframes.display.html
+import bigframes.dtypes as dtypes
 
 # anywidget and traitlets are optional dependencies. We don't want the import of
 # this module to fail if they aren't installed, though. Instead, we try to
@@ -48,6 +50,12 @@
     WIDGET_BASE = object
 
 
+@dataclasses.dataclass(frozen=True)
+class _SortState:
+    column: str
+    ascending: bool
+
+
 class TableWidget(WIDGET_BASE):
     """An interactive, paginated table widget for BigFrames DataFrames.
 
@@ -63,6 +71,9 @@ class TableWidget(WIDGET_BASE):
         allow_none=True,
     ).tag(sync=True)
     table_html = traitlets.Unicode().tag(sync=True)
+    sort_column = traitlets.Unicode("").tag(sync=True)
+    sort_ascending = traitlets.Bool(True).tag(sync=True)
+    orderable_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
     _initial_load_complete = traitlets.Bool(False).tag(sync=True)
     _batches: Optional[blocks.PandasBatches] = None
     _error_message = traitlets.Unicode(allow_none=True, default_value=None).tag(
@@ -89,15 +100,25 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         self._all_data_loaded = False
         self._batch_iter: Optional[Iterator[pd.DataFrame]] = None
         self._cached_batches: List[pd.DataFrame] = []
+        self._last_sort_state: Optional[_SortState] = None
 
         # respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
 
         # set traitlets properties that trigger observers
+        # TODO(b/462525985): Investigate and improve TableWidget UX for DataFrames with a large number of columns.
         self.page_size = initial_page_size
+        # TODO(b/463754889): Support non-string column labels for sorting.
+        if all(isinstance(col, str) for col in dataframe.columns):
+            self.orderable_columns = [
+                str(col_name)
+                for col_name, dtype in dataframe.dtypes.items()
+                if dtypes.is_orderable(dtype)
+            ]
+        else:
+            self.orderable_columns = []
 
-        # len(dataframe) is expensive, since it will trigger a
-        # SELECT COUNT(*) query. It is a must have however.
+        # obtain the row counts
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
         # before we get here so that the count might already be cached.
         self._reset_batches_for_new_page_size()
@@ -121,6 +142,11 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # Also used as a guard to prevent observers from firing during initialization.
         self._initial_load_complete = True
 
+    @traitlets.observe("_initial_load_complete")
+    def _on_initial_load_complete(self, change: Dict[str, Any]):
+        if change["new"]:
+            self._set_table_html()
+
     @functools.cached_property
     def _esm(self):
         """Load JavaScript code from external file."""
@@ -221,13 +247,17 @@ def _cached_data(self) -> pd.DataFrame:
             return pd.DataFrame(columns=self._dataframe.columns)
         return pd.concat(self._cached_batches, ignore_index=True)
 
+    def _reset_batch_cache(self) -> None:
+        """Resets batch caching attributes."""
+        self._cached_batches = []
+        self._batch_iter = None
+        self._all_data_loaded = False
+
     def _reset_batches_for_new_page_size(self) -> None:
         """Reset the batch iterator when page size changes."""
         self._batches = self._dataframe._to_pandas_batches(page_size=self.page_size)
 
-        self._cached_batches = []
-        self._batch_iter = None
-        self._all_data_loaded = False
+        self._reset_batch_cache()
 
     def _set_table_html(self) -> None:
         """Sets the current html data based on the current page and page size."""
@@ -237,6 +267,21 @@ def _set_table_html(self) -> None:
             )
             return
 
+        # Apply sorting if a column is selected
+        df_to_display = self._dataframe
+        if self.sort_column:
+            # TODO(b/463715504): Support sorting by index columns.
+            df_to_display = df_to_display.sort_values(
+                by=self.sort_column, ascending=self.sort_ascending
+            )
+
+        # Reset batches when sorting changes
+        if self._last_sort_state != _SortState(self.sort_column, self.sort_ascending):
+            self._batches = df_to_display._to_pandas_batches(page_size=self.page_size)
+            self._reset_batch_cache()
+            self._last_sort_state = _SortState(self.sort_column, self.sort_ascending)
+            self.page = 0  # Reset to first page
+
         start = self.page * self.page_size
         end = start + self.page_size
 
@@ -272,8 +317,14 @@ def _set_table_html(self) -> None:
         self.table_html = bigframes.display.html.render_html(
             dataframe=page_data,
             table_id=f"table-{self._table_id}",
+            orderable_columns=self.orderable_columns,
         )
 
+    @traitlets.observe("sort_column", "sort_ascending")
+    def _sort_changed(self, _change: Dict[str, Any]):
+        """Handler for when sorting parameters change from the frontend."""
+        self._set_table_html()
+
     @traitlets.observe("page")
     def _page_changed(self, _change: Dict[str, Any]) -> None:
         """Handler for when the page number is changed from the frontend."""
@@ -288,6 +339,9 @@ def _page_size_changed(self, _change: Dict[str, Any]) -> None:
             return
         # Reset the page to 0 when page size changes to avoid invalid page states
         self.page = 0
+        # Reset the sort state to default (no sort)
+        self.sort_column = ""
+        self.sort_ascending = True
 
         # Reset batches to use new page size for future data fetching
         self._reset_batches_for_new_page_size()

bigframes/display/html.py

@@ -17,14 +17,15 @@
 from __future__ import annotations
 
 import html
+from typing import Any
 
 import pandas as pd
 import pandas.api.types
 
 from bigframes._config import options
 
 
-def _is_dtype_numeric(dtype) -> bool:
+def _is_dtype_numeric(dtype: Any) -> bool:
     """Check if a dtype is numeric for alignment purposes."""
     return pandas.api.types.is_numeric_dtype(dtype)
 
@@ -33,18 +34,31 @@ def render_html(
     *,
     dataframe: pd.DataFrame,
     table_id: str,
+    orderable_columns: list[str] | None = None,
 ) -> str:
     """Render a pandas DataFrame to HTML with specific styling."""
     classes = "dataframe table table-striped table-hover"
     table_html = [f'<table border="1" class="{classes}" id="{table_id}">']
     precision = options.display.precision
+    orderable_columns = orderable_columns or []
 
     # Render table head
     table_html.append("  <thead>")
     table_html.append('    <tr style="text-align: left;">')
     for col in dataframe.columns:
+        th_classes = []
+        if col in orderable_columns:
+            th_classes.append("sortable")
+        class_str = f'class="{" ".join(th_classes)}"' if th_classes else ""
+        header_div = (
+            '<div style="resize: horizontal; overflow: auto; '
+            "box-sizing: border-box; width: 100%; height: 100%; "
+            'padding: 0.5em;">'
+            f"{html.escape(str(col))}"
+            "</div>"
+        )
         table_html.append(
-            f'      <th style="text-align: left;"><div style="resize: horizontal; overflow: auto; box-sizing: border-box; width: 100%; height: 100%; padding: 0.5em;">{html.escape(str(col))}</div></th>'
+            f'      <th style="text-align: left;" {class_str}>{header_div}</th>'
         )
     table_html.append("    </tr>")
     table_html.append("  </thead>")