chore: format warning message with newlines and ansi color (#1447)

* use ibis fill_null instead of fillna

* minor typo in JSON warning messages

* chore: format warning message with newlines and ansi color

Author: chelsea-lin

bigframes/_config/bigquery_options.py

@@ -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."

bigframes/_config/experiment_options.py

@@ -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."
             )

bigframes/core/array_value.py

@@ -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()

bigframes/core/blocks.py

@@ -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 "

bigframes/core/compile/aggregate_compiler.py

@@ -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

bigframes/core/global_session.py

@@ -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}"
         )

bigframes/core/indexers.py

@@ -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."
         )

bigframes/core/utils.py

@@ -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

bigframes/dataframe.py

@@ -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.

bigframes/exceptions.py

@@ -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