Update ArrayRecord Reader to allow reading from GCS directly since Reglieli

PiperOrigin-RevId: 792768387

Author: ArrayRecord Team

MODULE.bazel

@@ -21,15 +21,15 @@ module(
     repo_name = "com_google_array_record",
 )
 
-bazel_dep(name = "rules_proto", version = "7.0.2")
-bazel_dep(name = "rules_python", version = "1.0.0")
+bazel_dep(name = "rules_proto", version = "7.1.0")
+bazel_dep(name = "rules_python", version = "1.4.1")
 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 = "20250127.0")
+bazel_dep(name = "abseil-cpp", version = "20250127.1")
 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")
+bazel_dep(name = "riegeli", version = "0.0.0-20250717-5b2e77e")
 bazel_dep(name = "pybind11_bazel", version = "2.12.0")
 
 SUPPORTED_PYTHON_VERSIONS = [

python/BUILD

@@ -17,9 +17,12 @@ pybind_extension(
         "//cpp:array_record_reader",
         "//cpp:array_record_writer",
         "//cpp:thread_pool",
+        "//third_party/cloud_cpp/google/cloud/storage:storage_client",
         "@riegeli//riegeli/base:initializer",
         "@riegeli//riegeli/bytes:fd_reader",
         "@riegeli//riegeli/bytes:fd_writer",
+        "@riegeli//riegeli/gcs:gcs_object",
+        "@riegeli//riegeli/gcs:gcs_reader",
     ],
 )
 

python/array_record_module.cc

@@ -21,19 +21,25 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "file/base/options.pb.h"
 #include "absl/status/status.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "cpp/array_record_reader.h"
 #include "cpp/array_record_writer.h"
 #include "cpp/thread_pool.h"
+#include "third_party/cloud_cpp/google/cloud/storage/client.h"
 #include "pybind11/gil.h"
 #include "pybind11/pybind11.h"
 #include "pybind11/pytypes.h"
 #include "pybind11/stl.h"
 #include "riegeli/base/maker.h"
 #include "riegeli/bytes/fd_reader.h"
 #include "riegeli/bytes/fd_writer.h"
+#include "riegeli/gcs/gcs_object.h"
+#include "riegeli/gcs/gcs_reader.h"
+
 
 namespace py = pybind11;
 
@@ -50,10 +56,13 @@ PYBIND11_MODULE(array_record_module, m) {
                throw py::value_error(
                    std::string(status_or_option.status().message()));
              }
+             riegeli::FdWriterBase::Options file_writer_options;
+             file_writer_options.set_buffer_size(size_t{16} << 20);
              // Release the GIL because IO is time consuming.
              py::gil_scoped_release scoped_release;
              return new array_record::ArrayRecordWriter(
-                 riegeli::Maker<riegeli::FdWriter>(path),
+                 riegeli::Maker<riegeli::FdWriter>(
+                     path, std::move(file_writer_options)),
                  status_or_option.value());
            }),
            py::arg("path"), py::arg("options") = "")
@@ -84,18 +93,29 @@ PYBIND11_MODULE(array_record_module, m) {
                    std::string(status_or_option.status().message()));
              }
              riegeli::FdReaderBase::Options file_reader_options;
+             riegeli::GcsReader::Options gcs_reader_options;
              if (kwargs.contains("file_reader_buffer_size")) {
                auto file_reader_buffer_size =
                    kwargs["file_reader_buffer_size"].cast<int64_t>();
                file_reader_options.set_buffer_size(file_reader_buffer_size);
+               gcs_reader_options.set_buffer_size(file_reader_buffer_size);
              }
              // Release the GIL because IO is time consuming.
              py::gil_scoped_release scoped_release;
-             return new array_record::ArrayRecordReader(
-                 riegeli::Maker<riegeli::FdReader>(
-                     path, std::move(file_reader_options)),
-                 status_or_option.value(),
-                 array_record::ArrayRecordGlobalPool());
+             if (absl::StartsWith(path, "gs://")) {
+               return new array_record::ArrayRecordReader(
+                   riegeli::Maker<riegeli::GcsReader>(
+                       google::cloud::storage::Client(),
+                       riegeli::GcsObject(path), std::move(gcs_reader_options)),
+                   status_or_option.value(),
+                   array_record::ArrayRecordGlobalPool());
+             } else {
+               return new array_record::ArrayRecordReader(
+                   riegeli::Maker<riegeli::FdReader>(
+                       path, std::move(file_reader_options)),
+                   status_or_option.value(),
+                   array_record::ArrayRecordGlobalPool());
+             }
            }),
            py::arg("path"), py::arg("options") = "", R"(
            ArrayRecordReader for fast sequential or random access.