Merge branch 'main' into udf-packages

Author: jialuoo

bigframes/_config/display_options.py

@@ -26,8 +26,12 @@
 class DisplayOptions:
     __doc__ = vendored_pandas_config.display_options_doc
 
+    # Options borrowed from pandas.
     max_columns: int = 20
-    max_rows: int = 25
+    max_rows: int = 10
+    precision: int = 6
+
+    # Options unique to BigQuery DataFrames.
     progress_bar: Optional[str] = "auto"
     repr_mode: Literal["head", "deferred", "anywidget"] = "head"
 
@@ -52,6 +56,8 @@ def pandas_repr(display_options: DisplayOptions):
         display_options.max_columns,
         "display.max_rows",
         display_options.max_rows,
+        "display.precision",
+        display_options.precision,
         "display.show_dimensions",
         True,
     ) as pandas_context:

bigframes/blob/_functions.py

@@ -473,7 +473,9 @@ def pdf_extract_func(src_obj_ref_rt: str) -> str:
     return result_json
 
 
-pdf_extract_def = FunctionDef(pdf_extract_func, ["pypdf", "requests", "pypdf[crypto]"])
+pdf_extract_def = FunctionDef(
+    pdf_extract_func, ["pypdf>=5.3.1,<6.0.0", "requests", "cryptography==43.0.3"]
+)
 
 
 # Extracts text from a PDF url and chunks it simultaneously
@@ -527,4 +529,6 @@ def pdf_chunk_func(src_obj_ref_rt: str, chunk_size: int, overlap_size: int) -> s
     return result_json
 
 
-pdf_chunk_def = FunctionDef(pdf_chunk_func, ["pypdf", "requests", "pypdf[crypto]"])
+pdf_chunk_def = FunctionDef(
+    pdf_chunk_func, ["pypdf>=5.3.1,<6.0.0", "requests", "cryptography==43.0.3"]
+)

bigframes/core/blocks.py

@@ -387,25 +387,39 @@ def reversed(self) -> Block:
             index_labels=self.index.names,
         )
 
-    def reset_index(self, drop: bool = True) -> Block:
+    def reset_index(self, level: LevelsType = None, drop: bool = True) -> Block:
         """Reset the index of the block, promoting the old index to a value column.
 
         Arguments:
+            level: the label or index level of the index levels to remove.
             name: this is the column id for the new value id derived from the old index
 
         Returns:
             A new Block because dropping index columns can break references
             from Index classes that point to this block.
         """
+        if level:
+            # preserve original order, not user provided order
+            level_ids: Sequence[str] = [
+                id for id in self.index_columns if id in self.index.resolve_level(level)
+            ]
+        else:
+            level_ids = self.index_columns
+
         expr = self._expr
-        if (
+        if set(self.index_columns) > set(level_ids):
+            new_index_cols = [col for col in self.index_columns if col not in level_ids]
+            new_index_labels = [self.col_id_to_index_name[id] for id in new_index_cols]
+        elif (
             self.session._default_index_type
             == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
         ):
             expr, new_index_col_id = expr.promote_offsets()
             new_index_cols = [new_index_col_id]
+            new_index_labels = [None]
         elif self.session._default_index_type == bigframes.enums.DefaultIndexKind.NULL:
             new_index_cols = []
+            new_index_labels = []
         else:
             raise ValueError(
                 f"Unrecognized default index kind: {self.session._default_index_type}"
@@ -415,22 +429,23 @@ def reset_index(self, drop: bool = True) -> Block:
             # Even though the index might be part of the ordering, keep that
             # ordering expression as reset_index shouldn't change the row
             # order.
-            expr = expr.drop_columns(self.index_columns)
+            expr = expr.drop_columns(level_ids)
             return Block(
                 expr,
                 index_columns=new_index_cols,
+                index_labels=new_index_labels,
                 column_labels=self.column_labels,
             )
         else:
             # Add index names to column index
-            index_labels = self.index.names
             column_labels_modified = self.column_labels
-            for level, label in enumerate(index_labels):
+            for position, level_id in enumerate(level_ids):
+                label = self.col_id_to_index_name[level_id]
                 if label is None:
-                    if "index" not in self.column_labels and len(index_labels) <= 1:
+                    if "index" not in self.column_labels and self.index.nlevels <= 1:
                         label = "index"
                     else:
-                        label = f"level_{level}"
+                        label = f"level_{self.index_columns.index(level_id)}"
 
                 if label in self.column_labels:
                     raise ValueError(f"cannot insert {label}, already exists")
@@ -439,11 +454,12 @@ def reset_index(self, drop: bool = True) -> Block:
                     label = tuple(label if i == 0 else "" for i in range(nlevels))
                 # Create index copy with label inserted
                 # See: https://pandas.pydata.org/docs/reference/api/pandas.Index.insert.html
-                column_labels_modified = column_labels_modified.insert(level, label)
+                column_labels_modified = column_labels_modified.insert(position, label)
 
             return Block(
-                expr,
+                expr.select_columns((*new_index_cols, *level_ids, *self.value_columns)),
                 index_columns=new_index_cols,
+                index_labels=new_index_labels,
                 column_labels=column_labels_modified,
             )
 

bigframes/core/compile/polars/compiler.py

@@ -168,7 +168,7 @@ def compile_op(self, op: ops.ScalarOp, *args: pl.Expr) -> pl.Expr:
 
         @compile_op.register(gen_ops.InvertOp)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
-            return ~input
+            return input.not_()
 
         @compile_op.register(num_ops.AbsOp)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:

bigframes/core/compile/polars/lowering.py

@@ -14,10 +14,18 @@
 
 import dataclasses
 
+import numpy as np
+
 from bigframes import dtypes
 from bigframes.core import bigframe_node, expression
 from bigframes.core.rewrite import op_lowering
-from bigframes.operations import comparison_ops, datetime_ops, json_ops, numeric_ops
+from bigframes.operations import (
+    comparison_ops,
+    datetime_ops,
+    generic_ops,
+    json_ops,
+    numeric_ops,
+)
 import bigframes.operations as ops
 
 # TODO: Would be more precise to actually have separate op set for polars ops (where they diverge from the original ops)
@@ -288,6 +296,26 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
         return _lower_cast(expr.op, expr.inputs[0])
 
 
+def invert_bytes(byte_string):
+    inverted_bytes = ~np.frombuffer(byte_string, dtype=np.uint8)
+    return inverted_bytes.tobytes()
+
+
+class LowerInvertOp(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return generic_ops.InvertOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, generic_ops.InvertOp)
+        arg = expr.children[0]
+        if arg.output_type == dtypes.BYTES_DTYPE:
+            return generic_ops.PyUdfOp(invert_bytes, dtypes.BYTES_DTYPE).as_expr(
+                expr.inputs[0]
+            )
+        return expr
+
+
 def _coerce_comparables(
     expr1: expression.Expression,
     expr2: expression.Expression,
@@ -385,6 +413,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
     LowerFloorDivRule(),
     LowerModRule(),
     LowerAsTypeRule(),
+    LowerInvertOp(),
 )
 
 

bigframes/core/compile/polars/operations/generic_ops.py

@@ -45,3 +45,14 @@ def isnull_op_impl(
     input: pl.Expr,
 ) -> pl.Expr:
     return input.is_null()
+
+
+@polars_compiler.register_op(generic_ops.PyUdfOp)
+def py_udf_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: generic_ops.PyUdfOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.map_elements(
+        op.fn, return_dtype=polars_compiler._DTYPE_MAPPING[op._output_type]
+    )

bigframes/dataframe.py

@@ -23,6 +23,7 @@
 import re
 import sys
 import textwrap
+import traceback
 import typing
 from typing import (
     Callable,
@@ -814,7 +815,9 @@ def _repr_html_(self) -> str:
             except (AttributeError, ValueError, ImportError):
                 # Fallback if anywidget is not available
                 warnings.warn(
-                    "Anywidget mode is not available. Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. Falling back to deferred mode."
+                    "Anywidget mode is not available. "
+                    "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. "
+                    f"Falling back to deferred mode. Error: {traceback.format_exc()}"
                 )
                 return formatter.repr_query_job(self._compute_dry_run())
 
@@ -2312,9 +2315,39 @@ def _assign_series_join_on_index(
 
         return DataFrame(block.with_index_labels(self._block.index.names))
 
-    def reset_index(self, *, drop: bool = False) -> DataFrame:
-        block = self._block.reset_index(drop)
-        return DataFrame(block)
+    @overload  # type: ignore[override]
+    def reset_index(
+        self,
+        level: blocks.LevelsType = ...,
+        drop: bool = ...,
+        inplace: Literal[False] = ...,
+    ) -> DataFrame:
+        ...
+
+    @overload
+    def reset_index(
+        self,
+        level: blocks.LevelsType = ...,
+        drop: bool = ...,
+        inplace: Literal[True] = ...,
+    ) -> None:
+        ...
+
+    @overload
+    def reset_index(
+        self, level: blocks.LevelsType = None, drop: bool = False, inplace: bool = ...
+    ) -> Optional[DataFrame]:
+        ...
+
+    def reset_index(
+        self, level: blocks.LevelsType = None, drop: bool = False, inplace: bool = False
+    ) -> Optional[DataFrame]:
+        block = self._block.reset_index(level, drop)
+        if inplace:
+            self._set_block(block)
+            return None
+        else:
+            return DataFrame(block)
 
     def set_index(
         self,

bigframes/display/anywidget.py

@@ -23,6 +23,7 @@
 import pandas as pd
 
 import bigframes
+import bigframes.display.html
 
 # 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 limit the surface that
@@ -201,12 +202,9 @@ def _set_table_html(self):
         page_data = cached_data.iloc[start:end]
 
         # Generate HTML table
-        self.table_html = page_data.to_html(
-            index=False,
-            max_rows=None,
+        self.table_html = bigframes.display.html.render_html(
+            dataframe=page_data,
             table_id=f"table-{self._table_id}",
-            classes="table table-striped table-hover",
-            escape=False,
         )
 
     @traitlets.observe("page")

bigframes/display/html.py

@@ -0,0 +1,79 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HTML rendering for DataFrames and other objects."""
+
+from __future__ import annotations
+
+import html
+
+import pandas as pd
+import pandas.api.types
+
+from bigframes._config import options
+
+
+def _is_dtype_numeric(dtype) -> bool:
+    """Check if a dtype is numeric for alignment purposes."""
+    return pandas.api.types.is_numeric_dtype(dtype)
+
+
+def render_html(
+    *,
+    dataframe: pd.DataFrame,
+    table_id: str,
+) -> 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
+
+    # Render table head
+    table_html.append("  <thead>")
+    table_html.append('    <tr style="text-align: left;">')
+    for col in dataframe.columns:
+        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>'
+        )
+    table_html.append("    </tr>")
+    table_html.append("  </thead>")
+
+    # Render table body
+    table_html.append("  <tbody>")
+    for i in range(len(dataframe)):
+        table_html.append("    <tr>")
+        row = dataframe.iloc[i]
+        for col_name, value in row.items():
+            dtype = dataframe.dtypes.loc[col_name]  # type: ignore
+            align = "right" if _is_dtype_numeric(dtype) else "left"
+            table_html.append(
+                '      <td style="text-align: {}; padding: 0.5em;">'.format(align)
+            )
+
+            # TODO(b/438181139): Consider semi-exploding ARRAY/STRUCT columns
+            # into multiple rows/columns like the BQ UI does.
+            if pandas.api.types.is_scalar(value) and pd.isna(value):
+                table_html.append('        <em style="color: gray;">&lt;NA&gt;</em>')
+            else:
+                if isinstance(value, float):
+                    formatted_value = f"{value:.{precision}f}"
+                    table_html.append(f"        {html.escape(formatted_value)}")
+                else:
+                    table_html.append(f"        {html.escape(str(value))}")
+            table_html.append("      </td>")
+        table_html.append("    </tr>")
+    table_html.append("  </tbody>")
+    table_html.append("</table>")
+
+    return "\n".join(table_html)

bigframes/operations/generic_ops.py

@@ -446,3 +446,15 @@ class SqlScalarOp(base_ops.NaryOp):
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return self._output_type
+
+
+@dataclasses.dataclass(frozen=True)
+class PyUdfOp(base_ops.NaryOp):
+    """Represents a local UDF."""
+
+    name: typing.ClassVar[str] = "py_udf"
+    fn: typing.Callable
+    _output_type: dtypes.ExpressionType
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        return self._output_type