Merge branch 'main' into github_more_action

Author: TrevorBergeron

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/blob/_functions.py

@@ -112,6 +112,10 @@ def image_blur_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
@@ -121,7 +125,7 @@ def image_blur_func(
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
     dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -135,12 +139,13 @@ def image_blur_func(
     ext = ext_mappings.get(ext, ext)
     content_type = "image/" + ext
 
-    requests.put(
+    session.put(
         url=dst_url,
         data=bts,
         headers={
             "Content-Type": content_type,
         },
+        timeout=30,
     )
 
     return dst_obj_ref_rt
@@ -157,13 +162,17 @@ def image_blur_to_bytes_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -193,6 +202,10 @@ def image_resize_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
@@ -202,7 +215,7 @@ def image_resize_func(
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
     dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -216,12 +229,13 @@ def image_resize_func(
     ext = ext_mappings.get(ext, ext)
     content_type = "image/" + ext
 
-    requests.put(
+    session.put(
         url=dst_url,
         data=bts,
         headers={
             "Content-Type": content_type,
         },
+        timeout=30,
     )
 
     return dst_obj_ref_rt
@@ -245,13 +259,17 @@ def image_resize_to_bytes_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -280,6 +298,10 @@ def image_normalize_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
@@ -296,7 +318,7 @@ def image_normalize_func(
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
     dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -312,12 +334,13 @@ def image_normalize_func(
     ext = ext_mappings.get(ext, ext)
     content_type = "image/" + ext
 
-    requests.put(
+    session.put(
         url=dst_url,
         data=bts,
         headers={
             "Content-Type": content_type,
         },
+        timeout=30,
     )
 
     return dst_obj_ref_rt
@@ -336,6 +359,10 @@ def image_normalize_to_bytes_func(
     import cv2 as cv  # type: ignore
     import numpy as np
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     ext = ext or ".jpeg"
 
@@ -349,7 +376,7 @@ def image_normalize_to_bytes_func(
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
-    response = requests.get(src_url)
+    response = session.get(src_url, timeout=30)
     bts = response.content
 
     nparr = np.frombuffer(bts, np.uint8)
@@ -374,11 +401,15 @@ def pdf_extract_func(src_obj_ref_rt: str) -> str:
 
     from pypdf import PdfReader  # type: ignore
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
-    response = requests.get(src_url, stream=True)
+    response = session.get(src_url, timeout=30, stream=True)
     response.raise_for_status()
     pdf_bytes = response.content
 
@@ -403,11 +434,15 @@ def pdf_chunk_func(src_obj_ref_rt: str, chunk_size: int, overlap_size: int) -> s
 
     from pypdf import PdfReader  # type: ignore
     import requests
+    from requests import adapters
+
+    session = requests.Session()
+    session.mount("https://", adapters.HTTPAdapter(max_retries=3))
 
     src_obj_ref_rt_json = json.loads(src_obj_ref_rt)
     src_url = src_obj_ref_rt_json["access_urls"]["read_url"]
 
-    response = requests.get(src_url, stream=True)
+    response = session.get(src_url, timeout=30, stream=True)
     response.raise_for_status()
     pdf_bytes = response.content
 

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/bigframe_node.py

@@ -32,6 +32,8 @@
 
 COLUMN_SET = frozenset[identifiers.ColumnId]
 
+T = typing.TypeVar("T")
+
 
 @dataclasses.dataclass(frozen=True)
 class Field:
@@ -382,3 +384,14 @@ def bottom_up(
             results[node] = result
 
         return results[self]
+
+    def reduce_up(self, reduction: Callable[[BigFrameNode, Tuple[T, ...]], T]) -> T:
+        """Apply a bottom-up reduction to the tree."""
+        results: dict[BigFrameNode, T] = {}
+        for node in list(self.iter_nodes_topo()):
+            # child nodes have already been transformed
+            child_results = tuple(results[child] for child in node.child_nodes)
+            result = reduction(node, child_results)
+            results[node] = result
+
+        return results[self]