googleapis · chelsea-lin · Mar 6, 2025 · Mar 3, 2025 · Mar 3, 2025 · Mar 5, 2025
@@ -59,7 +59,9 @@ def _get_validated_location(value: Optional[str]) -> Optional[str]:
    # -> bpd.options.bigquery.location = "us-central-1"
    # -> location.setter
    # -> _get_validated_location
-    msg = UNKNOWN_LOCATION_MESSAGE.format(location=location, possibility=possibility)
+    msg = bfe.format_message(
+        UNKNOWN_LOCATION_MESSAGE.format(location=location, possibility=possibility)
+    )
    warnings.warn(msg, stacklevel=3, category=bfe.UnknownLocationWarning)

    return value
@@ -294,7 +296,7 @@ def use_regional_endpoints(self, value: bool):
            )

        if value:
-            msg = (
+            msg = bfe.format_message(
                "Use of regional endpoints is a feature in preview and "
                "available only in selected regions and projects. "
            )
@@ -354,7 +356,7 @@ def client_endpoints_override(self) -> dict:

    @client_endpoints_override.setter
    def client_endpoints_override(self, value: dict):
-        msg = (
+        msg = bfe.format_message(
            "This is an advanced configuration option for directly setting endpoints. "
            "Incorrect use may lead to unexpected behavior or system instability. "
            "Proceed only if you fully understand its implications."

@@ -34,7 +34,7 @@ def semantic_operators(self) -> bool:
    @semantic_operators.setter
    def semantic_operators(self, value: bool):
        if value is True:
-            msg = (
+            msg = bfe.format_message(
                "Semantic operators are still under experiments, and are subject "
                "to change in the future."
            )
@@ -48,7 +48,7 @@ def blob(self) -> bool:
    @blob.setter
    def blob(self, value: bool):
        if value is True:
-            msg = (
+            msg = bfe.format_message(
                "BigFrames Blob is still under experiments. It may not work and "
                "subject to change in the future."
            )
@@ -62,7 +62,7 @@ def udf(self) -> bool:
    @udf.setter
    def udf(self, value: bool):
        if value is True:
-            msg = (
+            msg = bfe.format_message(
                "BigFrames managed function (udf) is still under experiments. "
                "It may not work and subject to change in the future."
            )

@@ -107,8 +107,8 @@ def from_table(
        if offsets_col and primary_key:
            raise ValueError("must set at most one of 'offests', 'primary_key'")
        if any(i.field_type == "JSON" for i in table.schema if i.name in schema.names):
-            msg = (
-                "Interpreting JSON column(s) as the `db_dtypes.dbjson` extension type is"
+            msg = bfe.format_message(
+                "Interpreting JSON column(s) as the `db_dtypes.dbjson` extension type is "
                "in preview; this behavior may change in future versions."
            )
            warnings.warn(msg, bfe.PreviewWarning)
@@ -232,7 +232,9 @@ def slice(
        self, start: Optional[int], stop: Optional[int], step: Optional[int]
    ) -> ArrayValue:
        if self.node.order_ambiguous and not (self.session._strictly_ordered):
-            msg = "Window ordering may be ambiguous, this can cause unstable results."
+            msg = bfe.format_message(
+                "Window ordering may be ambiguous, this can cause unstable results."
+            )
            warnings.warn(msg, bfe.AmbiguousWindowWarning)
        return ArrayValue(
            nodes.SliceNode(
@@ -254,7 +256,7 @@ def promote_offsets(self) -> Tuple[ArrayValue, str]:
                    "Generating offsets not supported in partial ordering mode"
                )
            else:
-                msg = (
+                msg = bfe.format_message(
                    "Window ordering may be ambiguous, this can cause unstable results."
                )
                warnings.warn(msg, category=bfe.AmbiguousWindowWarning)
@@ -417,7 +419,9 @@ def project_window_op(
                        "Generating offsets not supported in partial ordering mode"
                    )
                else:
-                    msg = "Window ordering may be ambiguous, this can cause unstable results."
+                    msg = bfe.format_message(
+                        "Window ordering may be ambiguous, this can cause unstable results."
+                    )
                    warnings.warn(msg, category=bfe.AmbiguousWindowWarning)

        output_name = self._gen_namespaced_uid()

@@ -64,6 +64,7 @@
 import bigframes.core.utils as utils
 import bigframes.core.window_spec as windows
 import bigframes.dtypes
+import bigframes.exceptions as bfe
 import bigframes.features
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
@@ -630,12 +631,12 @@ def _materialize_local(
            # Since we cannot acquire the table size without a query_job,
            # we skip the sampling.
            if sample_config.enable_downsampling:
-                warnings.warn(
+                msg = bfe.format_message(
                    "Sampling is disabled and there is no download size limit when 'allow_large_results' is set to "
                    "False. To prevent downloading excessive data, it is recommended to use the peek() method, or "
-                    "limit the data with methods like .head() or .sample() before proceeding with downloads.",
-                    UserWarning,
+                    "limit the data with methods like .head() or .sample() before proceeding with downloads."
                )
+                warnings.warn(msg, category=UserWarning)
            fraction = 2

        # TODO: Maybe materialize before downsampling
@@ -652,7 +653,7 @@ def _materialize_local(
                    " # Setting it to None will download all the data\n"
                    f"{constants.FEEDBACK_LINK}"
                )
-            msg = (
+            msg = bfe.format_message(
                f"The data size ({table_mb:.2f} MB) exceeds the maximum download limit of"
                f"({max_download_size} MB). It will be downsampled to {max_download_size} "
                "MB for download.\nPlease refer to the documentation for configuring "

@@ -165,7 +165,7 @@ def _(
 ) -> ibis_types.NumericValue:
    # Will be null if all inputs are null. Pandas defaults to zero sum though.
    bq_sum = _apply_window_if_present(column.sum(), window)
-    return bq_sum.fillna(ibis_types.literal(0))
+    return bq_sum.fill_null(ibis_types.literal(0))


 @compile_unary_agg.register
@@ -610,12 +610,7 @@ def _(
    result = _apply_window_if_present(_is_true(column).all(), window)
    literal = ibis_types.literal(True)

-    return cast(
-        ibis_types.BooleanScalar,
-        result.fill_null(literal)
-        if hasattr(result, "fill_null")
-        else result.fillna(literal),
-    )
+    return cast(ibis_types.BooleanScalar, result.fill_null(literal))


 @compile_unary_agg.register
@@ -628,12 +623,7 @@ def _(
    result = _apply_window_if_present(_is_true(column).any(), window)
    literal = ibis_types.literal(False)

-    return cast(
-        ibis_types.BooleanScalar,
-        result.fill_null(literal)
-        if hasattr(result, "fill_null")
-        else result.fillna(literal),
-    )
+    return cast(ibis_types.BooleanScalar, result.fill_null(literal))


 @compile_ordered_unary_agg.register

@@ -39,7 +39,7 @@ def _try_close_session(session: bigframes.session.Session):
        session_id = session.session_id
        location = session._location
        project_id = session._project
-        msg = (
+        msg = bfe.format_message(
            f"Session cleanup failed for session with id: {session_id}, "
            f"location: {location}, project: {project_id}"
        )

@@ -407,7 +407,7 @@ def _struct_accessor_check_and_warn(
        return

    if not bigframes.dtypes.is_string_like(series.index.dtype):
-        msg = (
+        msg = bfe.format_message(
            "Are you trying to access struct fields? If so, please use Series.struct.field(...) "
            "method instead."
        )

@@ -196,7 +196,7 @@ def decorator(func):

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
-            warnings.warn(msg, category=bfe.PreviewWarning)
+            warnings.warn(bfe.format_message(msg), category=bfe.PreviewWarning)
            return func(*args, **kwargs)

        return wrapper

@@ -1581,7 +1581,10 @@ def to_arrow(
        Returns:
            pyarrow.Table: A pyarrow Table with all rows and columns of this DataFrame.
        """
-        msg = "to_arrow is in preview. Types and unnamed / duplicate name columns may change in future."
+        msg = bfe.format_message(
+            "to_arrow is in preview. Types and unnamed or duplicate name columns may "
+            "change in future."
+        )
        warnings.warn(msg, category=bfe.PreviewWarning)

        pa_table, query_job = self._block.to_arrow(
@@ -4104,7 +4107,7 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
        # to the applied function should be a Series, not a scalar.

        if utils.get_axis_number(axis) == 1:
-            msg = "axis=1 scenario is in preview."
+            msg = bfe.format_message("axis=1 scenario is in preview.")
            warnings.warn(msg, category=bfe.PreviewWarning)

            # TODO(jialuo): Deprecate the "bigframes_remote_function" attribute.

@@ -14,6 +14,8 @@

 """Public exceptions and warnings used across BigQuery DataFrames."""

+import textwrap
+
 # NOTE: This module should not depend on any others in the package.


@@ -87,3 +89,27 @@ class ApiDeprecationWarning(FutureWarning):

 class BadIndexerKeyWarning(Warning):
    """The indexer key is not used correctly."""
+
+
+class ColorFormatter:
+    WARNING = "\033[93m"
+    ENDC = "\033[0m"
+
+
+def format_message(message: str, fill: bool = True):
+    """Formats a warning message with ANSI color codes for the warning color.
+
+    Args:
+        message: The warning message string.
+        fill: Whether to wrap the message text using `textwrap.fill`.
+            Defaults to True.  Set to False to prevent wrapping,
+            especially if the message already contains newlines.
+
+    Returns:
+        The formatted message string, with ANSI color codes for warning color
+        if color is supported, otherwise the original message.  If `fill` is
+        True, the message will be wrapped to fit the terminal width.
+    """
+    if fill:
+        message = textwrap.fill(message)
+    return ColorFormatter.WARNING + message + ColorFormatter.ENDC
@@ -489,7 +489,7 @@ def remote_function(

        if cloud_function_ingress_settings is None:
            cloud_function_ingress_settings = "all"
-            msg = (
+            msg = bfe.format_message(
                "The `cloud_function_ingress_settings` are set to 'all' by default, "
                "which will change to 'internal-only' for enhanced security in future version 2.0 onwards. "
                "However, you will be able to explicitly pass cloud_function_ingress_settings='all' if you need. "
@@ -549,7 +549,7 @@ def wrapper(func):
                (input_type := input_types[0]) == bf_series.Series
                or input_type == pandas.Series
            ):
-                msg = "input_types=Series is in preview."
+                msg = bfe.format_message("input_types=Series is in preview.")
                warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)

                # we will model the row as a json serialized string containing the data
@@ -836,7 +836,7 @@ def wrapper(func):
                (input_type := input_types[0]) == bf_series.Series
                or input_type == pandas.Series
            ):
-                msg = "input_types=Series is in preview."
+                msg = bfe.format_message("input_types=Series is in preview.")
                warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)

                # we will model the row as a json serialized string containing

@@ -231,7 +231,7 @@ def func(*bigframes_args, **bigframes_kwargs):
            )
        function_input_dtypes.append(input_dtype)
    if has_unknown_dtypes:
-        msg = (
+        msg = bfe.format_message(
            "The function has one or more missing input data types. BigQuery DataFrames "
            f"will assume default data type {bigframes.dtypes.DEFAULT_DTYPE} for them."
        )

@@ -27,6 +27,7 @@

 import bigframes_vendored.sklearn.base

+import bigframes.exceptions as bfe
 from bigframes.ml import core
 import bigframes.ml.utils as utils
 import bigframes.pandas as bpd
@@ -269,7 +270,7 @@ def _predict_and_retry(

            if df_succ.empty:
                if max_retries > 0:
-                    msg = "Can't make any progress, stop retrying."
+                    msg = bfe.format_message("Can't make any progress, stop retrying.")
                    warnings.warn(msg, category=RuntimeWarning)
                break

@@ -281,7 +282,7 @@ def _predict_and_retry(
                break

        if not df_fail.empty:
-            msg = (
+            msg = bfe.format_message(
                f"Some predictions failed. Check column {self._status_col} for detailed "
                "status. You may want to filter the failed rows and retry."
            )