pr comments

Author: TrevorBergeron

bigframes/functions/_function_client.py

@@ -229,6 +229,7 @@ def provision_bq_managed_function(
 
         # TODO(b/406283812): Expose the capability to pass down
         # capture_references=True in the public udf API.
+        # TODO(b/495508827): Include all config in the value hash.
         if (
             capture_references
             and (python_version := _utils.get_python_version())

bigframes/functions/_utils.py

@@ -191,7 +191,6 @@ def get_managed_function_name(
     session_id: str | None = None,
 ):
     """Get a name for the bigframes managed function for the given user defined function."""
-    # TODO: Move over to logic used by remote functions
     parts = [_BIGFRAMES_FUNCTION_PREFIX]
     if session_id:
         parts.append(session_id)

bigframes/functions/udf_def.py

@@ -15,7 +15,6 @@
 
 import dataclasses
 import functools
-import hashlib
 import inspect
 import io
 import os
@@ -25,6 +24,7 @@
 
 import cloudpickle
 from google.cloud import bigquery
+import google_crc32c
 import pandas as pd
 
 import bigframes.dtypes
@@ -84,7 +84,7 @@ def sql_type(self) -> str:
         return self.dtype.sql_type
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self.name.encode())
         hash_val.update(self.dtype.stable_hash())
         return hash_val.digest()
@@ -116,7 +116,7 @@ def sql_type(self) -> str:
         return function_typing.sdk_type_to_sql_string(sdk_type)
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self._py_type.__name__.encode())
         return hash_val.digest()
 
@@ -145,7 +145,6 @@ class VirtualListTypeV1:
     def py_type(self) -> Type[list[Any]]:
         return list[self.inner_dtype.py_type]  # type: ignore
 
-    # TODO: Specify emulating type and mapping expressions between said types
     @property
     def bf_type(self) -> bigframes.dtypes.Dtype:
         return bigframes.dtypes.list_type(self.inner_dtype.bf_type)
@@ -165,7 +164,8 @@ def out_expr(
         if self.inner_dtype.py_type is str:
             return as_str_list
         elif self.inner_dtype.py_type is bool:
-            # TODO: hack so we don't need to make ArrayMap support general expressions yet
+            # hack so we don't need to make ArrayMap support general expressions yet
+            # with b/495513753 we can map the equality operator instead
             return ops.ArrayMapOp(ops.IsInOp(values=("true",))).as_expr(as_str_list)
         else:
             return ops.ArrayMapOp(ops.AsTypeOp(self.inner_dtype.bf_type)).as_expr(
@@ -177,7 +177,7 @@ def sql_type(self) -> str:
         return f"ARRAY<{self.inner_dtype.sql_type}>"
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self._PROTOCOL_ID.encode())
         hash_val.update(self.inner_dtype.stable_hash())
         return hash_val.digest()
@@ -212,7 +212,7 @@ def emulating_type(self) -> DirectScalarType:
         return DirectScalarType(str)
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self._PROTOCOL_ID.encode())
         return hash_val.digest()
 
@@ -227,6 +227,7 @@ class UdfSignature:
     output: DirectScalarType | VirtualListTypeV1
 
     def __post_init__(self):
+        # Validate inputs and outputs are of the correct types.
         assert all(isinstance(arg, UdfArg) for arg in self.inputs)
         assert isinstance(self.output, (DirectScalarType, VirtualListTypeV1))
 
@@ -240,7 +241,6 @@ def to_sql_input_signature(self) -> str:
     def protocol_metadata(self) -> str | None:
         import bigframes.functions._utils
 
-        # TODO: The output field itself should handle this, to handle protocol versioning.
         if isinstance(self.output, VirtualListTypeV1):
             return bigframes.functions._utils.get_bigframes_metadata(
                 python_output_type=self.output.py_type
@@ -362,7 +362,7 @@ def to_remote_function_compatible(self) -> UdfSignature:
         return self
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         for input_type in self.inputs:
             hash_val.update(input_type.stable_hash())
         hash_val.update(self.output.stable_hash())
@@ -438,7 +438,7 @@ def stable_hash(self) -> bytes:
             def_copy, protocol=_pickle_protocol_version
         )
 
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(normalized_pickled_code)
 
         if self.package_requirements:
@@ -464,7 +464,7 @@ class CloudRunFunctionConfig:
     concurrency: int | None
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self.code.stable_hash())
         hash_val.update(self.signature.stable_hash())
         hash_val.update(str(self.timeout_seconds).encode())
@@ -503,7 +503,7 @@ def from_bq_routine(cls, routine: bigquery.Routine) -> RemoteFunctionConfig:
         )
 
     def stable_hash(self) -> bytes:
-        hash_val = hashlib.md5()
+        hash_val = google_crc32c.Checksum()
         hash_val.update(self.endpoint.encode())
         hash_val.update(self.signature.stable_hash())
         hash_val.update(self.connection_id.encode())

bigframes/operations/array_ops.py

@@ -93,7 +93,7 @@ def output_type(self, *input_types):
 @dataclasses.dataclass(frozen=True)
 class ArrayMapOp(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "array_map"
-    # TODO: Generalize to chained expressions
+    # TODO(b/495513753): Generalize to chained expressions
     map_op: base_ops.UnaryOp
 
     def output_type(self, *input_types):

setup.py

@@ -46,6 +46,7 @@
     "google-cloud-bigquery-connection >=1.12.0",
     "google-cloud-resource-manager >=1.10.3",
     "google-cloud-storage >=2.0.0",
+    "google-crc32c >=1.0.0,<2.0.0",
     "grpc-google-iam-v1 >= 0.14.2",
     "numpy >=1.24.0",
     "pandas >=1.5.3",