Upgrade protobuf to match with TF 2.20.0rc0

Author: mtsokol

MODULE.bazel

@@ -22,11 +22,11 @@ module(
 )
 
 bazel_dep(name = "rules_proto", version = "7.0.2")
-bazel_dep(name = "rules_python", version = "0.37.0")
-bazel_dep(name = "platforms", version = "0.0.10")
-bazel_dep(name = "protobuf", version = "28.3")
+bazel_dep(name = "rules_python", version = "1.0.0")
+bazel_dep(name = "platforms", version = "0.0.11")
+bazel_dep(name = "protobuf", version = "31.1")
 bazel_dep(name = "googletest", version = "1.15.2")
-bazel_dep(name = "abseil-cpp", version = "20240722.0")
+bazel_dep(name = "abseil-cpp", version = "20250127.0")
 bazel_dep(name = "abseil-py", version = "2.1.0")
 bazel_dep(name = "eigen", version = "3.4.0.bcr.3")
 bazel_dep(name = "riegeli", version = "0.0.0-20241218-3385e3c")

beam/examples/requirements.txt

@@ -1,3 +1 @@
 array-record[beam]
-google-cloud-storage==2.11.0
-tensorflow==2.14.0

cpp/array_record_writer.h

@@ -217,11 +217,11 @@ class ArrayRecordWriterBase : public riegeli::Object {
     // 31 in 64-bit build).
     //
     // Default: `absl::nullopt`.
-    Options& set_window_log(std::optional<int> window_log) {
+    Options& set_window_log(absl::optional<int> window_log) {
       compressor_options_.set_window_log(window_log);
       return *this;
     }
-    std::optional<int> window_log() const {
+    absl::optional<int> window_log() const {
       return compressor_options_.window_log();
     }
 

oss/build_whl.sh

@@ -87,12 +87,12 @@ function main() {
   echo $(date) : "=== Output wheel file is in: ${DEST}"
 
   # Install ArrayRecord from the wheel and run smoke tests.
-  $PYTHON_BIN -m pip install --find-links="${DEST}" --pre array-record
+  $PYTHON_BIN -m pip install ${OUTPUT_DIR}/all_dist/array_record*.whl
   $PYTHON_BIN -c 'import array_record'
   $PYTHON_BIN -c 'from array_record.python import array_record_data_source'
   # TF is not available on Python 3.13 and above.
   if [ "$(uname)" != "Darwin" ] && (( "${PYTHON_MINOR_VERSION}" < 13 )); then
-    $PYTHON_BIN -m pip install jax tensorflow grain
+    $PYTHON_BIN -m pip install jax tensorflow==2.20.0rc0 grain
     $PYTHON_BIN oss/test_import_grain.py
     $PYTHON_BIN oss/test_import_tensorflow.py
   fi

setup.py

@@ -9,11 +9,19 @@
     'etils[epath]',
 ]
 
+TF_PACKAGE = [
+  'tensorflow>=2.20.0rc0'
+]
+
 BEAM_EXTRAS = [
-    'apache-beam[gcp]==2.53.0',
+    'apache-beam[gcp]>=2.53.0',
     'google-cloud-storage>=2.11.0',
-    'tensorflow>=2.14.0'
-]
+] + TF_PACKAGE
+
+TEST_EXTRAS = [
+  'jax',
+  'grain',
+] + TF_PACKAGE
 
 
 class BinaryDistribution(Distribution):
@@ -34,7 +42,7 @@ def has_ext_modules(self):
     package_data={'': ['*.so']},
     python_requires='>=3.10',
     install_requires=REQUIRED_PACKAGES,
-    extras_require={'beam': BEAM_EXTRAS},
+    extras_require={'beam': BEAM_EXTRAS, 'test': TEST_EXTRAS},
     url='https://github.com/google/array_record',
     license='Apache-2.0',
     classifiers=[