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

shuoweil · shuoweil · commit 9b4b18a31307 · 2025-11-25T20:01:19.000Z
diff --git a/bigframes/bigquery/__init__.py b/bigframes/bigquery/__init__.py
@@ -47,6 +47,7 @@
     json_extract,
     json_extract_array,
     json_extract_string_array,
+    json_keys,
     json_query,
     json_query_array,
     json_set,
@@ -138,6 +139,7 @@
     "json_extract",
     "json_extract_array",
     "json_extract_string_array",
+    "json_keys",
     "json_query",
     "json_query_array",
     "json_set",
diff --git a/bigframes/bigquery/_operations/json.py b/bigframes/bigquery/_operations/json.py
@@ -421,6 +421,35 @@ def json_value_array(
     return input._apply_unary_op(ops.JSONValueArray(json_path=json_path))
 
 
+def json_keys(
+    input: series.Series,
+    max_depth: Optional[int] = None,
+) -> series.Series:
+    """Returns all keys in the root of a JSON object as an ARRAY of STRINGs.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+
+        >>> s = bpd.Series(['{"b": {"c": 2}, "a": 1}'], dtype="json")
+        >>> bbq.json_keys(s)
+        0    ['a' 'b' 'b.c']
+        dtype: list<item: string>[pyarrow]
+
+    Args:
+        input (bigframes.series.Series):
+            The Series containing JSON data.
+        max_depth (int, optional):
+            Specifies the maximum depth of nested fields to search for keys. If not
+            provided, searched keys at all levels.
+
+    Returns:
+        bigframes.series.Series: A new Series containing arrays of keys from the input JSON.
+    """
+    return input._apply_unary_op(ops.JSONKeys(max_depth=max_depth))
+
+
 def to_json(
     input: series.Series,
 ) -> series.Series:
diff --git a/bigframes/core/compile/ibis_compiler/scalar_op_registry.py b/bigframes/core/compile/ibis_compiler/scalar_op_registry.py
@@ -1234,6 +1234,11 @@ def json_value_array_op_impl(x: ibis_types.Value, op: ops.JSONValueArray):
     return json_value_array(json_obj=x, json_path=op.json_path)
 
 
+@scalar_op_compiler.register_unary_op(ops.JSONKeys, pass_op=True)
+def json_keys_op_impl(x: ibis_types.Value, op: ops.JSONKeys):
+    return json_keys(x, op.max_depth)
+
+
 # Blob Ops
 @scalar_op_compiler.register_unary_op(ops.obj_fetch_metadata_op)
 def obj_fetch_metadata_op_impl(obj_ref: ibis_types.Value):
@@ -2059,6 +2064,14 @@ def to_json_string(value) -> ibis_dtypes.String:  # type: ignore[empty-body]
     """Convert value to JSON-formatted string."""
 
 
+@ibis_udf.scalar.builtin(name="json_keys")
+def json_keys(  # type: ignore[empty-body]
+    json_obj: ibis_dtypes.JSON,
+    max_depth: ibis_dtypes.Int64,
+) -> ibis_dtypes.Array[ibis_dtypes.String]:
+    """Extracts unique JSON keys from a JSON expression."""
+
+
 @ibis_udf.scalar.builtin(name="json_value")
 def json_value(  # type: ignore[empty-body]
     json_obj: ibis_dtypes.JSON, json_path: ibis_dtypes.String
diff --git a/bigframes/core/compile/sqlglot/aggregations/nullary_compiler.py b/bigframes/core/compile/sqlglot/aggregations/nullary_compiler.py
@@ -49,5 +49,5 @@ def _(
     result: sge.Expression = sge.func("ROW_NUMBER")
     if window is None:
         # ROW_NUMBER always needs an OVER clause.
-        return sge.Window(this=result)
-    return apply_window_if_present(result, window, include_framing_clauses=False)
+        return sge.Window(this=result) - 1
+    return apply_window_if_present(result, window, include_framing_clauses=False) - 1
diff --git a/bigframes/core/compile/sqlglot/aggregations/unary_compiler.py b/bigframes/core/compile/sqlglot/aggregations/unary_compiler.py
@@ -400,6 +400,43 @@ def _(
     return apply_window_if_present(expr, window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.QcutOp)
+def _(
+    op: agg_ops.QcutOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    percent_ranks_order_by = sge.Ordered(this=column.expr, desc=False)
+    percent_ranks = apply_window_if_present(
+        sge.func("PERCENT_RANK"),
+        window,
+        include_framing_clauses=False,
+        order_by_override=[percent_ranks_order_by],
+    )
+    if isinstance(op.quantiles, int):
+        scaled_rank = percent_ranks * sge.convert(op.quantiles)
+        # Calculate the 0-based bucket index.
+        bucket_index = sge.func("CEIL", scaled_rank) - sge.convert(1)
+        safe_bucket_index = sge.func("GREATEST", bucket_index, 0)
+
+        return sge.If(
+            this=sge.Is(this=column.expr, expression=sge.Null()),
+            true=sge.Null(),
+            false=sge.Cast(this=safe_bucket_index, to="INT64"),
+        )
+    else:
+        case = sge.Case()
+        first_quantile = sge.convert(op.quantiles[0])
+        case = case.when(
+            sge.LT(this=percent_ranks, expression=first_quantile), sge.Null()
+        )
+        for bucket_n in range(len(op.quantiles) - 1):
+            quantile = sge.convert(op.quantiles[bucket_n + 1])
+            bucket = sge.convert(bucket_n)
+            case = case.when(sge.LTE(this=percent_ranks, expression=quantile), bucket)
+        return case.else_(sge.Null())
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.QuantileOp)
 def _(
     op: agg_ops.QuantileOp,
diff --git a/bigframes/core/compile/sqlglot/aggregations/windows.py b/bigframes/core/compile/sqlglot/aggregations/windows.py
@@ -26,6 +26,7 @@ def apply_window_if_present(
     value: sge.Expression,
     window: typing.Optional[window_spec.WindowSpec] = None,
     include_framing_clauses: bool = True,
+    order_by_override: typing.Optional[typing.List[sge.Ordered]] = None,
 ) -> sge.Expression:
     if window is None:
         return value
@@ -44,7 +45,11 @@ def apply_window_if_present(
     else:
         order_by = get_window_order_by(window.ordering)
 
-    order = sge.Order(expressions=order_by) if order_by else None
+    order = None
+    if order_by_override is not None and len(order_by_override) > 0:
+        order = sge.Order(expressions=order_by_override)
+    elif order_by:
+        order = sge.Order(expressions=order_by)
 
     group_by = (
         [
diff --git a/bigframes/core/compile/sqlglot/expressions/json_ops.py b/bigframes/core/compile/sqlglot/expressions/json_ops.py
@@ -39,6 +39,11 @@ def _(expr: TypedExpr, op: ops.JSONExtractStringArray) -> sge.Expression:
     return sge.func("JSON_EXTRACT_STRING_ARRAY", expr.expr, sge.convert(op.json_path))
 
 
+@register_unary_op(ops.JSONKeys, pass_op=True)
+def _(expr: TypedExpr, op: ops.JSONKeys) -> sge.Expression:
+    return sge.func("JSON_KEYS", expr.expr, sge.convert(op.max_depth))
+
+
 @register_unary_op(ops.JSONQuery, pass_op=True)
 def _(expr: TypedExpr, op: ops.JSONQuery) -> sge.Expression:
     return sge.func("JSON_QUERY", expr.expr, sge.convert(op.json_path))
diff --git a/bigframes/core/compile/sqlglot/expressions/numeric_ops.py b/bigframes/core/compile/sqlglot/expressions/numeric_ops.py
@@ -443,6 +443,14 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     )
 
 
+@register_binary_op(ops.unsafe_pow_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    """For internal use only - where domain and overflow checks are not needed."""
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.Pow(this=left_expr, expression=right_expr)
+
+
 @register_unary_op(numeric_ops.isnan_op)
 def isnan(arg: TypedExpr) -> sge.Expression:
     return sge.IsNan(this=arg.expr)
diff --git a/bigframes/core/compile/sqlglot/sqlglot_ir.py b/bigframes/core/compile/sqlglot/sqlglot_ir.py
@@ -637,7 +637,12 @@ def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
 
 
 def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
-    sqlglot_type = sgt.from_bigframes_dtype(dtype)
+    sqlglot_type = sgt.from_bigframes_dtype(dtype) if dtype else None
+    if sqlglot_type is None:
+        if value is not None:
+            raise ValueError("Cannot infer SQLGlot type from None dtype.")
+        return sge.Null()
+
     if value is None:
         return _cast(sge.Null(), sqlglot_type)
     elif dtype == dtypes.BYTES_DTYPE:
diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py
@@ -128,6 +128,7 @@
     JSONExtract,
     JSONExtractArray,
     JSONExtractStringArray,
+    JSONKeys,
     JSONQuery,
     JSONQueryArray,
     JSONSet,
@@ -381,6 +382,7 @@
     "JSONExtract",
     "JSONExtractArray",
     "JSONExtractStringArray",
+    "JSONKeys",
     "JSONQuery",
     "JSONQueryArray",
     "JSONSet",
diff --git a/bigframes/operations/json_ops.py b/bigframes/operations/json_ops.py
@@ -199,6 +199,23 @@ def output_type(self, *input_types):
         return input_type
 
 
+@dataclasses.dataclass(frozen=True)
+class JSONKeys(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "json_keys"
+    max_depth: typing.Optional[int] = None
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if input_type != dtypes.JSON_DTYPE:
+            raise TypeError(
+                "Input type must be a valid JSON object or JSON-formatted string type."
+                + f" Received type: {input_type}"
+            )
+        return pd.ArrowDtype(
+            pa.list_(dtypes.bigframes_dtype_to_arrow_dtype(dtypes.STRING_DTYPE))
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class JSONDecode(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "json_decode"
diff --git a/tests/system/small/bigquery/test_json.py b/tests/system/small/bigquery/test_json.py
@@ -434,3 +434,53 @@ def test_to_json_string_from_struct():
     )
 
     pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_keys():
+    json_data = [
+        '{"name": "Alice", "age": 30}',
+        '{"city": "New York", "country": "USA", "active": true}',
+        "{}",
+        '{"items": [1, 2, 3]}',
+    ]
+    s = bpd.Series(json_data, dtype=dtypes.JSON_DTYPE)
+    actual = bbq.json_keys(s)
+
+    expected_data_pandas = [
+        ["age", "name"],
+        [
+            "active",
+            "city",
+            "country",
+        ],
+        [],
+        ["items"],
+    ]
+    expected = bpd.Series(
+        expected_data_pandas, dtype=pd.ArrowDtype(pa.list_(pa.string()))
+    )
+    pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_keys_with_max_depth():
+    json_data = [
+        '{"user": {"name": "Bob", "details": {"id": 123, "status": "approved"}}}',
+        '{"user": {"name": "Charlie"}}',
+    ]
+    s = bpd.Series(json_data, dtype=dtypes.JSON_DTYPE)
+    actual = bbq.json_keys(s, max_depth=2)
+
+    expected_data_pandas = [
+        ["user", "user.details", "user.name"],
+        ["user", "user.name"],
+    ]
+    expected = bpd.Series(
+        expected_data_pandas, dtype=pd.ArrowDtype(pa.list_(pa.string()))
+    )
+    pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_keys_from_string_error():
+    s = bpd.Series(['{"a": 1, "b": 2}', '{"c": 3}'])
+    with pytest.raises(TypeError):
+        bbq.json_keys(s)
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number/out.sql
@@ -19,7 +19,7 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    ROW_NUMBER() OVER () AS `bfcol_32`
+    ROW_NUMBER() OVER () - 1 AS `bfcol_32`
   FROM `bfcte_0`
 )
 SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number_with_window/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number_with_window/out.sql
@@ -5,7 +5,7 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
-    ROW_NUMBER() OVER (ORDER BY `int64_col` ASC NULLS LAST) AS `bfcol_1`
+    ROW_NUMBER() OVER (ORDER BY `int64_col` ASC NULLS LAST) - 1 AS `bfcol_1`
   FROM `bfcte_0`
 )
 SELECT
diff --git a/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_qcut/out.sql b/tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_qcut/out.sql
@@ -0,0 +1,61 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col`,
+    `rowindex`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    NOT `int64_col` IS NULL AS `bfcol_4`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    *,
+    IF(
+      `int64_col` IS NULL,
+      NULL,
+      CAST(GREATEST(
+        CEIL(PERCENT_RANK() OVER (PARTITION BY `bfcol_4` ORDER BY `int64_col` ASC) * 4) - 1,
+        0
+      ) AS INT64)
+    ) AS `bfcol_5`
+  FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    *,
+    IF(`bfcol_4`, `bfcol_5`, NULL) AS `bfcol_6`
+  FROM `bfcte_2`
+), `bfcte_4` AS (
+  SELECT
+    *,
+    NOT `int64_col` IS NULL AS `bfcol_10`
+  FROM `bfcte_3`
+), `bfcte_5` AS (
+  SELECT
+    *,
+    CASE
+      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) < 0
+      THEN NULL
+      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.25
+      THEN 0
+      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.5
+      THEN 1
+      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 0.75
+      THEN 2
+      WHEN PERCENT_RANK() OVER (PARTITION BY `bfcol_10` ORDER BY `int64_col` ASC) <= 1
+      THEN 3
+      ELSE NULL
+    END AS `bfcol_11`
+  FROM `bfcte_4`
+), `bfcte_6` AS (
+  SELECT
+    *,
+    IF(`bfcol_10`, `bfcol_11`, NULL) AS `bfcol_12`
+  FROM `bfcte_5`
+)
+SELECT
+  `rowindex`,
+  `int64_col`,
+  `bfcol_6` AS `qcut_w_int`,
+  `bfcol_12` AS `qcut_w_list`
+FROM `bfcte_6`
diff --git a/tests/unit/core/compile/sqlglot/aggregations/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/aggregations/test_unary_compiler.py
@@ -434,6 +434,27 @@ def test_pop_var(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql_window, "window_out.sql")
 
 
+def test_qcut(scalar_types_df: bpd.DataFrame, snapshot):
+    if sys.version_info < (3, 12):
+        pytest.skip(
+            "Skipping test due to inconsistent SQL formatting on Python < 3.12.",
+        )
+
+    col_name = "int64_col"
+    bf = scalar_types_df[[col_name]]
+    bf["qcut_w_int"] = bpd.qcut(bf[col_name], q=4, labels=False, duplicates="drop")
+
+    q_list = tuple([0, 0.25, 0.5, 0.75, 1])
+    bf["qcut_w_list"] = bpd.qcut(
+        scalar_types_df[col_name],
+        q=q_list,
+        labels=False,
+        duplicates="drop",
+    )
+
+    snapshot.assert_match(bf.sql, "out.sql")
+
+
 def test_quantile(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "int64_col"
     bf_df = scalar_types_df[[col_name]]
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_json_keys/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_json_keys/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_unsafe_pow_op/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_unsafe_pow_op/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/test_json_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_json_ops.py
diff --git a/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat/out.sql
diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql
