Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 7 additions & 8 deletions MODULE.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -13,25 +13,24 @@
# limitations under the License.

# TODO(fchern): automate version string alignment with setup.py
VERSION = "0.8.0"
VERSION = "0.8.1"

module(
name = "array_record",
version = VERSION,
repo_name = "com_google_array_record",
)

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 = "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 = "googletest", version = "1.15.2")
bazel_dep(name = "abseil-cpp", version = "20250127.1")
bazel_dep(name = "abseil-cpp", version = "20240722.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-20250717-5b2e77e")
bazel_dep(name = "riegeli", version = "0.0.0-20241218-3385e3c")
bazel_dep(name = "pybind11_bazel", version = "2.12.0")
bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc0")

SUPPORTED_PYTHON_VERSIONS = [
"3.10",
Expand Down
10 changes: 4 additions & 6 deletions oss/build_whl.sh
Original file line number Diff line number Diff line change
Expand Up @@ -90,12 +90,10 @@ function main() {
$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>=2.20.0 grain
$PYTHON_BIN oss/test_import_grain.py
$PYTHON_BIN oss/test_import_tensorflow.py
fi
$PYTHON_BIN -m pip install jax tensorflow>=2.20.0 grain
# Re-enable the grain import test once the new version is released.
# $PYTHON_BIN oss/test_import_grain.py
$PYTHON_BIN oss/test_import_tensorflow.py
}

main
2 changes: 0 additions & 2 deletions python/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,6 @@ pybind_extension(
"@riegeli//riegeli/base:initializer",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:fd_writer",
"@riegeli//riegeli/gcs:gcs_object",
"@riegeli//riegeli/gcs:gcs_reader",
],
)

Expand Down
29 changes: 6 additions & 23 deletions python/array_record_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ limitations under the License.
#include <vector>

#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"
Expand All @@ -35,8 +34,6 @@ limitations under the License.
#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;

Expand All @@ -53,13 +50,10 @@ 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, std::move(file_writer_options)),
riegeli::Maker<riegeli::FdWriter>(path),
status_or_option.value());
}),
py::arg("path"), py::arg("options") = "")
Expand Down Expand Up @@ -90,29 +84,18 @@ 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;
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());
}
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.
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ def has_ext_modules(self):

setup(
name='array_record',
version='0.8.0',
version='0.8.1',
description='A file format that achieves a new frontier of IO efficiency',
author='ArrayRecord team',
author_email='no-reply@google.com',
Expand Down