Merge main to shuowei-anywidget-nested-strcut-array

Author: shuoweil

bigframes/core/compile/sqlglot/aggregations/windows.py

@@ -44,6 +44,7 @@ def apply_window_if_present(
         order_by = None
     elif window.is_range_bounded:
         order_by = get_window_order_by((window.ordering[0],))
+        order_by = remove_null_ordering_for_range_windows(order_by)
     else:
         order_by = get_window_order_by(window.ordering)
 
@@ -150,6 +151,30 @@ def get_window_order_by(
     return tuple(order_by)
 
 
+def remove_null_ordering_for_range_windows(
+    order_by: typing.Optional[tuple[sge.Ordered, ...]],
+) -> typing.Optional[tuple[sge.Ordered, ...]]:
+    """Removes NULL FIRST/LAST from ORDER BY expressions in RANGE windows.
+    Here's the support matrix:
+    ✅ sum(x) over (order by y desc nulls last)
+    🚫 sum(x) over (order by y asc nulls last)
+    ✅ sum(x) over (order by y asc nulls first)
+    🚫 sum(x) over (order by y desc nulls first)
+    """
+    if order_by is None:
+        return None
+
+    new_order_by = []
+    for key in order_by:
+        kargs = key.args
+        if kargs.get("desc") is True and kargs.get("nulls_first", False):
+            kargs["nulls_first"] = False
+        elif kargs.get("desc") is False and not kargs.setdefault("nulls_first", True):
+            kargs["nulls_first"] = True
+        new_order_by.append(sge.Ordered(**kargs))
+    return tuple(new_order_by)
+
+
 def _get_window_bounds(
     value, is_preceding: bool
 ) -> tuple[typing.Union[str, sge.Expression], typing.Optional[str]]:

bigframes/core/compile/sqlglot/compiler.py

@@ -356,6 +356,9 @@ def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotI
                 observation_count = windows.apply_window_if_present(
                     sge.func("SUM", is_observation), window_spec
                 )
+                observation_count = sge.func(
+                    "COALESCE", observation_count, sge.convert(0)
+                )
             else:
                 # Operations like count treat even NULLs as valid observations
                 # for the sake of min_periods notnull is just used to convert

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

@@ -89,27 +89,39 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.ge_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.GTE(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.gt_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.GT(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.lt_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.LT(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.le_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.LTE(this=left_expr, expression=right_expr)

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

@@ -140,6 +140,19 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.Coalesce(this=left.expr, expressions=[right.expr])
 
 
+@register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
+def _(
+    left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp
+) -> sge.Expression:
+    routine_ref = op.function_def.routine_ref
+    # Quote project, dataset, and routine IDs to avoid keyword clashes.
+    func_name = (
+        f"`{routine_ref.project}`.`{routine_ref.dataset_id}`.`{routine_ref.routine_id}`"
+    )
+
+    return sge.func(func_name, left.expr, right.expr)
+
+
 @register_nary_op(ops.case_when_op)
 def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     # Need to upcast BOOL to INT if any output is numeric

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

@@ -388,6 +388,9 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.add_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
         # String addition
         return sge.Concat(expressions=[left.expr, right.expr])
@@ -442,6 +445,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.floordiv_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
 
@@ -525,6 +531,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.mul_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
 
@@ -548,6 +557,9 @@ def _(expr: TypedExpr, n_digits: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.sub_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype):
         left_expr = _coerce_bool_to_int(left)
         right_expr = _coerce_bool_to_int(right)

bigframes/display/anywidget.py

@@ -66,6 +66,7 @@ class TableWidget(_WIDGET_BASE):
 
     page = traitlets.Int(0).tag(sync=True)
     page_size = traitlets.Int(0).tag(sync=True)
+    max_columns = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     row_count = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     table_html = traitlets.Unicode("").tag(sync=True)
     sort_context = traitlets.List(traitlets.Dict(), default_value=[]).tag(sync=True)
@@ -103,10 +104,13 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
 
         # respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
+        initial_max_columns = bigframes.options.display.max_columns
 
         # 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
+        self.max_columns = initial_max_columns
+
         # TODO(b/469861913): Nested columns from structs (e.g., 'struct_col.name') are not currently sortable.
         # TODO(b/463754889): Support non-string column labels for sorting.
         if all(isinstance(col, str) for col in dataframe.columns):
@@ -218,6 +222,14 @@ def _validate_page_size(self, proposal: dict[str, Any]) -> int:
         max_page_size = 1000
         return min(value, max_page_size)
 
+    @traitlets.validate("max_columns")
+    def _validate_max_columns(self, proposal: dict[str, Any]) -> int:
+        """Validate max columns to ensure it's positive or 0 (for all)."""
+        value = proposal["value"]
+        if value is None:
+            return 0  # Normalize None to 0 for traitlet
+        return max(0, value)
+
     def _get_next_batch(self) -> bool:
         """
         Gets the next batch of data from the generator and appends to cache.
@@ -348,6 +360,7 @@ def _set_table_html(self) -> None:
                     dataframe=page_data,
                     table_id=f"table-{self._table_id}",
                     orderable_columns=self.orderable_columns,
+                    max_columns=self.max_columns,
                 )
 
         if new_page is not None:
@@ -382,3 +395,10 @@ def _page_size_changed(self, _change: dict[str, Any]) -> None:
 
         # Update the table display
         self._set_table_html()
+
+    @traitlets.observe("max_columns")
+    def _max_columns_changed(self, _change: dict[str, Any]) -> None:
+        """Handler for when max columns is changed from the frontend."""
+        if not self._initial_load_complete:
+            return
+        self._set_table_html()

bigframes/display/html.py

@@ -46,34 +46,64 @@ def render_html(
     dataframe: pd.DataFrame,
     table_id: str,
     orderable_columns: list[str] | None = None,
+    max_columns: int | None = None,
 ) -> str:
     """Render a pandas DataFrame to HTML with specific styling and nested data support."""
     # Flatten nested data first
     flatten_result = _flatten.flatten_nested_data(dataframe)
+    flat_df = flatten_result.dataframe
 
     orderable_columns = orderable_columns or []
     classes = "dataframe table table-striped table-hover"
     table_html_parts = [f'<table border="1" class="{classes}" id="{table_id}">']
+
+    # Handle column truncation
+    columns = list(flat_df.columns)
+    if max_columns is not None and max_columns > 0 and len(columns) > max_columns:
+        half = max_columns // 2
+        left_columns = columns[:half]
+        # Ensure we don't take more than available if half is 0 or calculation is weird,
+        # but typical case is safe.
+        right_count = max_columns - half
+        right_columns = columns[-right_count:] if right_count > 0 else []
+        show_ellipsis = True
+    else:
+        left_columns = columns
+        right_columns = []
+        show_ellipsis = False
+
     table_html_parts.append(
-        _render_table_header(flatten_result.dataframe, orderable_columns)
+        _render_table_header(
+            flat_df, orderable_columns, left_columns, right_columns, show_ellipsis
+        )
     )
     table_html_parts.append(
         _render_table_body(
-            flatten_result.dataframe,
+            flat_df,
             flatten_result.row_labels,
             flatten_result.continuation_rows,
             flatten_result.cleared_on_continuation,
             flatten_result.nested_columns,
+            left_columns,
+            right_columns,
+            show_ellipsis,
         )
     )
     table_html_parts.append("</table>")
     return "".join(table_html_parts)
 
 
-def _render_table_header(dataframe: pd.DataFrame, orderable_columns: list[str]) -> str:
+def _render_table_header(
+    dataframe: pd.DataFrame,
+    orderable_columns: list[str],
+    left_columns: list[Any],
+    right_columns: list[Any],
+    show_ellipsis: bool,
+) -> str:
     """Render the header of the HTML table."""
     header_parts = ["  <thead>", "    <tr>"]
-    for col in dataframe.columns:
+
+    def render_col_header(col):
         th_classes = []
         if col in orderable_columns:
             th_classes.append("sortable")
@@ -82,6 +112,18 @@ def _render_table_header(dataframe: pd.DataFrame, orderable_columns: list[str])
             f'      <th {class_str}><div class="bf-header-content">'
             f"{html.escape(str(col))}</div></th>"
         )
+
+    for col in left_columns:
+        render_col_header(col)
+
+    if show_ellipsis:
+        header_parts.append(
+            '      <th><div class="bf-header-content" style="cursor: default;">...</div></th>'
+        )
+
+    for col in right_columns:
+        render_col_header(col)
+
     header_parts.extend(["    </tr>", "  </thead>"])
     return "\n".join(header_parts)
 
@@ -92,6 +134,9 @@ def _render_table_body(
     continuation_rows: set[int] | None,
     clear_on_continuation: list[str],
     nested_originated_columns: set[str],
+    left_columns: list[Any],
+    right_columns: list[Any],
+    show_ellipsis: bool,
 ) -> str:
     """Render the body of the HTML table."""
     body_parts = ["  <tbody>"]
@@ -117,7 +162,9 @@ def _render_table_body(
             body_parts.append("    <tr>")
 
         row = dataframe.iloc[i]
-        for col_name, value in row.items():
+
+        def render_col_cell(col_name):
+            value = row[col_name]
             dtype = dataframe.dtypes.loc[col_name]  # type: ignore
             cell_html = _render_cell(
                 value,
@@ -129,6 +176,17 @@ def _render_table_body(
                 precision,
             )
             body_parts.append(cell_html)
+
+        for col in left_columns:
+            render_col_cell(col)
+
+        if show_ellipsis:
+            # Ellipsis cell
+            body_parts.append('      <td class="cell-align-left">...</td>')
+
+        for col in right_columns:
+            render_col_cell(col)
+
         body_parts.append("    </tr>")
     body_parts.append("  </tbody>")
     return "\n".join(body_parts)

bigframes/display/table_widget.css

@@ -120,15 +120,24 @@ body[data-theme='dark'] .bigframes-widget.bigframes-widget {
   margin: 0 8px;
 }
 
-.bigframes-widget .page-size {
+.bigframes-widget .settings {
   align-items: center;
   display: flex;
   flex-direction: row;
-  gap: 4px;
+  gap: 16px;
   justify-content: end;
 }
 
-.bigframes-widget .page-size label {
+.bigframes-widget .page-size,
+.bigframes-widget .max-columns {
+  align-items: center;
+  display: flex;
+  flex-direction: row;
+  gap: 4px;
+}
+
+.bigframes-widget .page-size label,
+.bigframes-widget .max-columns label {
   margin-right: 8px;
 }