Merge branch 'main' into sycai_precision_score_binary

Author: sycai

bigframes/bigquery/__init__.py

@@ -50,6 +50,7 @@
     json_value,
     json_value_array,
     parse_json,
+    to_json_string,
 )
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
 from bigframes.bigquery._operations.sql import sql_scalar
@@ -87,6 +88,7 @@
     json_value,
     json_value_array,
     parse_json,
+    to_json_string,
     # search ops
     create_vector_index,
     vector_search,

bigframes/bigquery/_operations/json.py

@@ -430,6 +430,40 @@ def json_value_array(
     return input._apply_unary_op(ops.JSONValueArray(json_path=json_path))
 
 
+def to_json_string(
+    input: series.Series,
+) -> series.Series:
+    """Converts a series to a JSON-formatted STRING value.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> bpd.options.display.progress_bar = None
+
+        >>> s = bpd.Series([1, 2, 3])
+        >>> bbq.to_json_string(s)
+        0    1
+        1    2
+        2    3
+        dtype: string
+
+        >>> s = bpd.Series([{"int": 1, "str": "pandas"}, {"int": 2, "str": "numpy"}])
+        >>> bbq.to_json_string(s)
+        0    {"int":1,"str":"pandas"}
+        1     {"int":2,"str":"numpy"}
+        dtype: string
+
+    Args:
+        input (bigframes.series.Series):
+            The Series to be converted.
+
+    Returns:
+        bigframes.series.Series: A new Series with the JSON-formatted STRING value.
+    """
+    return input._apply_unary_op(ops.ToJSONString())
+
+
 @utils.preview(name="The JSON-related API `parse_json`")
 def parse_json(
     input: series.Series,

bigframes/core/blocks.py

@@ -37,6 +37,7 @@
     Optional,
     Sequence,
     Tuple,
+    TYPE_CHECKING,
     Union,
 )
 import warnings
@@ -69,6 +70,9 @@
 from bigframes.session import dry_runs, execution_spec
 from bigframes.session import executor as executors
 
+if TYPE_CHECKING:
+    from bigframes.session.executor import ExecuteResult
+
 # Type constraint for wherever column labels are used
 Label = typing.Hashable
 
@@ -404,13 +408,15 @@ def reset_index(
         col_level: Union[str, int] = 0,
         col_fill: typing.Hashable = "",
         allow_duplicates: bool = False,
+        replacement: Optional[bigframes.enums.DefaultIndexKind] = None,
     ) -> 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
-            allow_duplicates:
+            allow_duplicates: if false, duplicate col labels will result in error
+            replacement: if not null, will override default index replacement type
 
         Returns:
             A new Block because dropping index columns can break references
@@ -425,23 +431,19 @@ def reset_index(
             level_ids = self.index_columns
 
         expr = self._expr
+        replacement_idx_type = replacement or self.session._default_index_type
         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
-        ):
+        elif replacement_idx_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:
+        elif replacement_idx_type == bigframes.enums.DefaultIndexKind.NULL:
             new_index_cols = []
             new_index_labels = []
         else:
-            raise ValueError(
-                f"Unrecognized default index kind: {self.session._default_index_type}"
-            )
+            raise ValueError(f"Unrecognized default index kind: {replacement_idx_type}")
 
         if drop:
             # Even though the index might be part of the ordering, keep that
@@ -630,15 +632,17 @@ def to_pandas(
             max_download_size, sampling_method, random_state
         )
 
-        df, query_job = self._materialize_local(
+        ex_result = self._materialize_local(
             materialize_options=MaterializationOptions(
                 downsampling=sampling,
                 allow_large_results=allow_large_results,
                 ordered=ordered,
             )
         )
+        df = ex_result.to_pandas()
+        df = self._copy_index_to_pandas(df)
         df.set_axis(self.column_labels, axis=1, copy=False)
-        return df, query_job
+        return df, ex_result.query_job
 
     def _get_sampling_option(
         self,
@@ -746,7 +750,7 @@ def _copy_index_to_pandas(self, df: pd.DataFrame) -> pd.DataFrame:
 
     def _materialize_local(
         self, materialize_options: MaterializationOptions = MaterializationOptions()
-    ) -> Tuple[pd.DataFrame, Optional[bigquery.QueryJob]]:
+    ) -> ExecuteResult:
         """Run query and download results as a pandas DataFrame. Return the total number of results as well."""
         # TODO(swast): Allow for dry run and timeout.
         under_10gb = (
@@ -815,8 +819,7 @@ def _materialize_local(
                 MaterializationOptions(ordered=materialize_options.ordered)
             )
         else:
-            df = execute_result.to_pandas()
-            return self._copy_index_to_pandas(df), execute_result.query_job
+            return execute_result
 
     def _downsample(
         self, total_rows: int, sampling_method: str, fraction: float, random_state

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -2068,9 +2068,7 @@ def json_extract_string_array(  # type: ignore[empty-body]
 
 
 @ibis_udf.scalar.builtin(name="to_json_string")
-def to_json_string(  # type: ignore[empty-body]
-    value,
-) -> ibis_dtypes.String:
+def to_json_string(value) -> ibis_dtypes.String:  # type: ignore[empty-body]
     """Convert value to JSON-formatted string."""
 
 

bigframes/core/compile/sqlglot/compiler.py

@@ -244,6 +244,28 @@ def compile_join(
             joins_nulls=node.joins_nulls,
         )
 
+    @_compile_node.register
+    def compile_isin_join(
+        self, node: nodes.InNode, left: ir.SQLGlotIR, right: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        conditions = (
+            typed_expr.TypedExpr(
+                scalar_compiler.compile_scalar_expression(node.left_col),
+                node.left_col.output_type,
+            ),
+            typed_expr.TypedExpr(
+                scalar_compiler.compile_scalar_expression(node.right_col),
+                node.right_col.output_type,
+            ),
+        )
+
+        return left.isin_join(
+            right,
+            indicator_col=node.indicator_col.sql,
+            conditions=conditions,
+            joins_nulls=node.joins_nulls,
+        )
+
     @_compile_node.register
     def compile_concat(
         self, node: nodes.ConcatNode, *children: ir.SQLGlotIR

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -336,6 +336,68 @@ def join(
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def isin_join(
+        self,
+        right: SQLGlotIR,
+        indicator_col: str,
+        conditions: tuple[typed_expr.TypedExpr, typed_expr.TypedExpr],
+        joins_nulls: bool = True,
+    ) -> SQLGlotIR:
+        """Joins the current query with another SQLGlotIR instance."""
+        left_cte_name = sge.to_identifier(
+            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
+        )
+
+        left_select = _select_to_cte(self.expr, left_cte_name)
+        # Prefer subquery over CTE for the IN clause's right side to improve SQL readability.
+        right_select = right.expr
+
+        left_ctes = left_select.args.pop("with", [])
+        right_ctes = right_select.args.pop("with", [])
+        merged_ctes = [*left_ctes, *right_ctes]
+
+        left_condition = typed_expr.TypedExpr(
+            sge.Column(this=conditions[0].expr, table=left_cte_name),
+            conditions[0].dtype,
+        )
+
+        new_column: sge.Expression
+        if joins_nulls:
+            right_table_name = sge.to_identifier(
+                next(self.uid_gen.get_uid_stream("bft_")), quoted=self.quoted
+            )
+            right_condition = typed_expr.TypedExpr(
+                sge.Column(this=conditions[1].expr, table=right_table_name),
+                conditions[1].dtype,
+            )
+            new_column = sge.Exists(
+                this=sge.Select()
+                .select(sge.convert(1))
+                .from_(sge.Alias(this=right_select.subquery(), alias=right_table_name))
+                .where(
+                    _join_condition(left_condition, right_condition, joins_nulls=True)
+                )
+            )
+        else:
+            new_column = sge.In(
+                this=left_condition.expr,
+                expressions=[right_select.subquery()],
+            )
+
+        new_column = sge.Alias(
+            this=new_column,
+            alias=sge.to_identifier(indicator_col, quoted=self.quoted),
+        )
+
+        new_expr = (
+            sge.Select()
+            .select(sge.Column(this=sge.Star(), table=left_cte_name), new_column)
+            .from_(sge.Table(this=left_cte_name))
+        )
+        new_expr.set("with", sge.With(expressions=merged_ctes))
+
+        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+
     def explode(
         self,
         column_names: tuple[str, ...],