diff --git a/python/Cargo.lock b/python/Cargo.lock index eb871fb..dfe5562 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -1673,9 +1673,9 @@ dependencies = [ [[package]] name = "vec64" -version = "0.5.3" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "662d4c3c6fe8d14a8b4a0b97155a0b22dd6d9b6626f42224be6987b0e0f1ca2b" +checksum = "8af6b3f149c2d27972bdfaa5b3499cd3969efbf5a5cc04c05988147deb513c85" [[package]] name = "version_check" diff --git a/python/src/input.rs b/python/src/input.rs index 781fd61..cfb4971 100644 --- a/python/src/input.rs +++ b/python/src/input.rs @@ -32,6 +32,7 @@ use lightstream::models::readers::ipc::mmap_table::MmapTableReader; use lightstream::models::readers::http::HttpTableReader; use lightstream::models::readers::parallel::tcp::TcpParallelTableReader; use lightstream::models::readers::json::JsonReader; +use lightstream::models::readers::object_storage::StorageReader; use lightstream::models::readers::parquet::load_parquet_table; use lightstream::models::readers::quic::QuicTableReader; use lightstream::models::readers::stdio::StdinTableReader; @@ -98,7 +99,7 @@ pub enum FileIO { reader: ChunkedParquetReader, }, Csv { - reader: CsvReader>, + reader: CsvReader, }, /// CSV text streamed over the stdio wire, one batch per /// `batch_size` rows. @@ -109,7 +110,7 @@ pub enum FileIO { reader: ChunkedCsvReader, }, Json { - reader: JsonReader>, + reader: JsonReader, }, } diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 4aa3056..2fad41a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -462,6 +462,29 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "aws-lc-rs" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2b2dcc879c3bae0d371e77c99f2238400ef24ec001394befa67b6e543add9e" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.44.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09fae7be8bb3174e05c6afdb34199e6dc0c7c04ba9fa237b1967adfbde27483" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + [[package]] name = "axum" version = "0.8.9" @@ -757,12 +780,31 @@ version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "colorchoice" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "compact_str" version = "0.9.1" @@ -853,6 +895,16 @@ dependencies = [ "libc", ] +[[package]] +name = "crc-fast" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5" +dependencies = [ + "digest 0.10.7", + "spin", +] + [[package]] name = "crc32fast" version = "1.5.2" @@ -1070,6 +1122,12 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "dyn-clone" version = "1.0.20" @@ -1238,6 +1296,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.34" @@ -1801,6 +1865,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1844,6 +1917,55 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.119", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.119", +] + [[package]] name = "jobserver" version = "0.1.35" @@ -1856,9 +1978,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.105" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce57d20d1ea864ce2ac172ab472d409214f4fd359f0b2a2775abdf522e2af99e" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" dependencies = [ "cfg-if", "futures-util", @@ -1961,6 +2083,7 @@ dependencies = [ "log", "memchr", "minarrow", + "object_store 0.14.1", "polars", "prost", "quinn", @@ -1980,6 +2103,7 @@ dependencies = [ "tokio-uring", "tonic", "tracing", + "url", "vec64", "wtransport", "zstd", @@ -2043,6 +2167,16 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "memchr" version = "2.8.0" @@ -2245,9 +2379,9 @@ dependencies = [ "itertools 0.14.0", "parking_lot", "percent-encoding", - "quick-xml", + "quick-xml 0.39.4", "rand 0.10.2", - "reqwest", + "reqwest 0.12.28", "ring", "serde", "serde_json", @@ -2261,6 +2395,46 @@ dependencies = [ "web-time", ] +[[package]] +name = "object_store" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +dependencies = [ + "async-trait", + "aws-lc-rs", + "base64 0.22.1", + "bytes", + "chrono", + "crc-fast", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body-util", + "httparse", + "humantime", + "hyper", + "itertools 0.15.0", + "md-5", + "parking_lot", + "percent-encoding", + "quick-xml 0.41.0", + "rand 0.10.2", + "reqwest 0.13.4", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror", + "tokio", + "tracing", + "url", + "wasm-bindgen-futures", + "web-time", +] + [[package]] name = "octets" version = "0.3.7" @@ -2649,7 +2823,7 @@ version = "0.55.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f6995c2121c31687704fae5281517e29bbd34838fcf8dd15c8e8ae526bb6624" dependencies = [ - "object_store", + "object_store 0.13.2", "parking_lot", "polars-arrow-format", "regex", @@ -2706,7 +2880,7 @@ dependencies = [ "memchr", "memmap2", "num-traits", - "object_store", + "object_store 0.13.2", "parking_lot", "percent-encoding", "polars-arrow", @@ -2721,7 +2895,7 @@ dependencies = [ "rand 0.10.2", "rayon", "regex", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "simdutf8", @@ -3073,9 +3247,9 @@ checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" [[package]] name = "portable-atomic-util" -version = "0.2.8" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" dependencies = [ "portable-atomic", ] @@ -3165,6 +3339,16 @@ dependencies = [ "serde", ] +[[package]] +name = "quick-xml" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "quinn" version = "0.11.12" @@ -3191,6 +3375,7 @@ version = "0.11.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc" dependencies = [ + "aws-lc-rs", "bytes", "getrandom 0.4.3", "lru-slab", @@ -3466,7 +3651,46 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", + "web-sys", +] + +[[package]] +name = "reqwest" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", "web-sys", ] @@ -3546,6 +3770,7 @@ version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ + "aws-lc-rs", "once_cell", "ring", "rustls-pki-types", @@ -3576,12 +3801,40 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -3787,6 +4040,16 @@ dependencies = [ "value-trait", ] +[[package]] +name = "simd_cesu8" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" +dependencies = [ + "rustc_version", + "simdutf8", +] + [[package]] name = "simdutf8" version = "0.1.5" @@ -3846,6 +4109,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spin" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" + [[package]] name = "sqlparser" version = "0.62.0" @@ -4446,9 +4715,9 @@ dependencies = [ [[package]] name = "vec64" -version = "0.5.3" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "662d4c3c6fe8d14a8b4a0b97155a0b22dd6d9b6626f42224be6987b0e0f1ca2b" +checksum = "8af6b3f149c2d27972bdfaa5b3499cd3969efbf5a5cc04c05988147deb513c85" dependencies = [ "libc", ] @@ -4501,9 +4770,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.128" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aecb87a33d3b0c5e3b7aa46336eaf486cffafbd281b195e4c8b80d50df2351bf" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" dependencies = [ "cfg-if", "once_cell", @@ -4514,9 +4783,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.78" +version = "0.4.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ef4c5d3d2cdf5c54f4231181768f5510842e350db025faf1f7163b1030ed928" +checksum = "6b7777d5cc23d0e91404e53ce2d5e8ec7acae3026b16233dba62cd3246457950" dependencies = [ "js-sys", "wasm-bindgen", @@ -4524,9 +4793,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.128" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a690d511e3c1a8b3a55e33511e3c2c00c78415cd23650f32b808627f5696b9ed" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4534,22 +4803,22 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.128" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "411e4887f0071ef2d2164a9d5fdf2d20efbef78fccd3a78b0c10a1dc5295e48a" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn 3.0.6", + "syn 2.0.119", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.128" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81941cd78d0c92026c33e5e01312845a4cb1e9af3407f9134b100dd03144103e" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" dependencies = [ "unicode-ident", ] @@ -4567,11 +4836,24 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" -version = "0.3.105" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fbddc4a036f00ec4f18c83445bd3115cb306a91da554919a099d9222fe4a7f8" +checksum = "c435338968042f4f59a557f690a253676d47ce13ceb55d70100e7facf6620a30" dependencies = [ "js-sys", "wasm-bindgen", @@ -4587,6 +4869,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "0.26.11" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index b446ec6..5f8e094 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -42,7 +42,7 @@ libc = { version = "0.2.183", optional = true } futures-sink = "0.3.32" log = "0.4" minarrow = { version = "0.18.3", features = ["chunked", "views", "select", "size"], default-features = false } -vec64 = { version = "0.5.2" } +vec64 = { version = "0.5.4" } snap = { version = "1.0", optional = true } tokio-tungstenite = { version = "0.28", optional = true } zstd = { version = "0.13", optional = true } @@ -88,6 +88,11 @@ polars = { version = "0.55.2", optional = true, default-features = false, featur # `decoders::int_ascii` module instead of the `atoi` crate. fast-float2 = { version = "0.2", optional = true } env_logger = { version = "0.11", optional = true } +# Object storage for remote URLs in the file readers and writers. The +# `tokio` feature provides the multipart upload support used by every +# backend. +object_store = { version = "0.14.1", optional = true, default-features = false, features = ["tokio"] } +url = { version = "2", optional = true } # Tokio's filesystem and multi-threaded runtime are not available on # WebAssembly targets. @@ -200,6 +205,18 @@ io_uring = ["uds", "tokio-uring", "protocol"] # Stdin/stdout transport stdio = ["tokio/io-std"] +# Object storage URLs (`s3://`, `gs://`, `az://`, `memory://`) in place of +# file paths for the file readers and writers. The in-memory store is +# always available. Each cloud backend is enabled by its own feature. +# Not available on WebAssembly targets. +object_storage = ["dep:object_store", "dep:url", "dep:bytes"] +# Amazon S3 and S3-compatible stores +aws = ["object_storage", "object_store/aws"] +# Google Cloud Storage +gcp = ["object_storage", "object_store/gcp"] +# Azure Blob Storage +azure = ["object_storage", "object_store/azure"] + # Lightstream protocol multiplexing protocol = [] # Protobuf support for protocol messages via prost diff --git a/rust/scripts/run-object-storage-s3-tests.sh b/rust/scripts/run-object-storage-s3-tests.sh new file mode 100755 index 0000000..b7e5cc8 --- /dev/null +++ b/rust/scripts/run-object-storage-s3-tests.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Run the object storage tests against Amazon S3. +# +# Usage: +# LIGHTSTREAM_S3_URL=s3://bucket/prefix scripts/run-object-storage-s3-tests.sh [profile] [cargo test args] +# +# Credentials come from the AWS CLI, which resolves profiles and SSO sessions +# that the object_store client does not read itself. The region is taken +# from AWS_REGION when set, otherwise from the CLI's configured region, and +# must match the bucket's region. +set -euo pipefail +cd "$(dirname "$0")/.." + +: "${LIGHTSTREAM_S3_URL:?set LIGHTSTREAM_S3_URL, for example s3://bucket/prefix}" + +if [[ $# -gt 0 && "$1" != -* ]]; then + export AWS_PROFILE="$1" + shift +fi + +eval "$(aws configure export-credentials --format env)" +export AWS_REGION="${AWS_REGION:-$(aws configure get region || true)}" +: "${AWS_REGION:?set AWS_REGION to the region of the bucket}" + +echo "Running object storage tests against ${LIGHTSTREAM_S3_URL} in ${AWS_REGION}" +cargo test --features "aws parquet csv json" --test object_storage_s3 -- --ignored "$@" diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 194693a..fa33880 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -217,6 +217,10 @@ pub mod models { /// Readers for files, mmap, and async streams pub mod readers { + /// Read handles over local files and object storage, used by the + /// file readers. + pub mod object_storage; + /// Arrow IPC readers pub mod ipc { /// File-based IPC reader @@ -246,7 +250,8 @@ pub mod models { pub mod parquet; /// Chunked file readers - glob a directory of `-NNNNN.` - /// files and present them as an ordered iterator of `Table`s. + /// files, including its partition directories, and present them as + /// an ordered iterator of `Table`s. pub mod chunked { /// Chunked CSV reader. #[cfg(feature = "csv")] @@ -315,6 +320,10 @@ pub mod models { /// Writers for Arrow IPC, CSV, and optionally Parquet. pub mod writers { + /// Write handles over local files and object storage, used by the + /// file writers. + pub mod object_storage; + pub mod ipc { /// Sync IPC stream writer. pub mod table_stream; @@ -352,6 +361,9 @@ pub mod models { /// Chunked Arrow IPC writer. Holds a tokio runtime internally /// to drive the async IPC file writer. pub mod arrow; + + /// Partition paths for the chunked writers. + pub mod partition; } /// TCP table writer @@ -477,6 +489,10 @@ pub mod models { /// Stdin byte stream adapter. #[cfg(feature = "stdio")] pub mod stdio; + + /// Object storage byte stream and write adapters. + #[cfg(feature = "object_storage")] + pub mod object_store; } /// Codecs for Arrow IPC and Lightstream protocol. @@ -508,6 +524,10 @@ pub mod models { pub mod io_uring; } +/// Object storage URLs for the file readers and writers. +#[cfg(feature = "object_storage")] +pub mod object_storage; + /// FlatBuffers-compiled Arrow IPC metadata support. pub mod arrow { /// Flatbuffers Arrow file metadata diff --git a/rust/src/models/readers/chunked/arrow.rs b/rust/src/models/readers/chunked/arrow.rs index fd28f9a..d0efe97 100644 --- a/rust/src/models/readers/chunked/arrow.rs +++ b/rust/src/models/readers/chunked/arrow.rs @@ -9,8 +9,10 @@ //! Reads a directory of `-NNNNNNNNNN.arrow` files emitted by //! `ChunkedArrowWriter` and presents them as an ordered iterator of //! `Table`s. Each chunk file is a complete, independently readable Arrow -//! IPC file. The reader sorts files by their numeric index so consumers -//! see batches in write order. +//! IPC file. The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! ## Concurrency model //! @@ -28,7 +30,6 @@ //! decode) across worker threads and returns a `SuperTable` with batches //! in write order. -use std::fs; use std::io; use std::path::{Path, PathBuf}; @@ -37,6 +38,7 @@ use std::sync::Arc; use minarrow::{Concatenate, SuperTable, Table}; use crate::models::readers::ipc::file_table::FileTableReader; +use crate::models::readers::object_storage::list_files; use crate::traits::chunked_table_reader::ChunkedTableReader; /// Iterator over chunk files in a directory written by @@ -64,9 +66,7 @@ impl ChunkedTableReader for ChunkedArrowReader { fn list_paths>(dir: P, base: &str) -> io::Result> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for entry in fs::read_dir(dir.as_ref())? { - let entry = entry?; - let path = entry.path(); + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -79,7 +79,12 @@ impl ChunkedTableReader for ChunkedArrowReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } @@ -139,10 +144,26 @@ impl Iterator for ChunkedArrowReader { #[cfg(test)] mod tests { + use std::fs; + use super::*; + use std::sync::Arc; + use crate::models::writers::chunked::arrow::ChunkedArrowWriter; + use crate::models::writers::chunked::partition::PartitionType; use crate::traits::chunked_table_writer::ChunkedTableWriter; - use minarrow::{Table, fa_i32}; + use minarrow::{Array, FieldArray, IntegerArray, Table, fa_i32}; + + /// 2026-02-09T07:00:00Z and 08:00:00Z in epoch nanoseconds. + const SEVEN_AM: u64 = 1_770_620_400_000_000_000; + const EIGHT_AM: u64 = SEVEN_AM + 3_600_000_000_000; + + /// Table of `n_rows` rows, identifying each chunk by its row count. + fn rows(n_rows: i32) -> Table { + let values: Vec = (0..n_rows).collect(); + let column = FieldArray::from_arr("n", Array::from_int32(IntegerArray::from_slice(&values))); + Table::new("b", Some(vec![column])) + } #[test] fn par_load_batched_returns_batches_in_write_order() { @@ -187,4 +208,76 @@ mod tests { fs::remove_dir_all(&dir).ok(); } + + #[test] + fn auto_partitioned_chunks_read_back_by_partition_then_index() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_auto_partition"); + let _ = fs::remove_dir_all(&dir); + + let mut w = ChunkedArrowWriter::new(&dir, "part") + .unwrap() + .with_partition(PartitionType::Auto); + let p0 = w.write_chunk_for_time(&rows(1), EIGHT_AM).unwrap(); + let p1 = w.write_chunk_for_time(&rows(2), SEVEN_AM).unwrap(); + let p2 = w.write_chunk_for_time(&rows(3), EIGHT_AM).unwrap(); + assert_eq!(p0, dir.join("date=2026-02-09/hour=08/part-0000000000.arrow")); + assert_eq!(p1, dir.join("date=2026-02-09/hour=07/part-0000000001.arrow")); + assert_eq!(p2, dir.join("date=2026-02-09/hour=08/part-0000000002.arrow")); + + // The 07:00 chunk precedes both 08:00 chunks, which keep write order. + let all = ChunkedArrowReader::open(&dir, "part", ()).unwrap().load_batched().unwrap(); + let n_rows: Vec = all.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![2, 1, 3]); + + // A partition directory reads that partition alone. + let eight = ChunkedArrowReader::open(dir.join("date=2026-02-09/hour=08"), "part", ()) + .unwrap() + .load_batched() + .unwrap(); + let n_rows: Vec = eight.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![1, 3]); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn custom_partitions_come_from_the_table() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_custom_partition"); + let _ = fs::remove_dir_all(&dir); + + let by_size = PartitionType::Custom(Arc::new(|table: &Table, _time: u64| { + format!("rows={}", table.n_rows) + })); + let mut w = ChunkedArrowWriter::new(&dir, "part") + .unwrap() + .with_partition(by_size); + let path = w.write_chunk(&rows(4)).unwrap(); + assert_eq!(path, dir.join("rows=4/part-0000000000.arrow")); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn par_write_all_creates_partition_directories() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_par_partition"); + let _ = fs::remove_dir_all(&dir); + + let w = ChunkedArrowWriter::new(&dir, "part") + .unwrap() + .with_partition(PartitionType::Auto); + let tables: Vec = (1..=6).map(rows).collect(); + let refs: Vec<&Table> = tables.iter().collect(); + let paths = w.par_write_all(&refs, Some(3)).unwrap(); + + // Every batch of one call shares one time and therefore one partition. + let partition = paths[0].parent().unwrap(); + assert!(partition.starts_with(&dir) && partition != dir); + assert!(paths.iter().all(|path| path.parent() == Some(partition))); + + let read = ChunkedArrowReader::par_load_batched(&dir, "part", (), None).unwrap(); + let n_rows: Vec = read.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![1, 2, 3, 4, 5, 6]); + + fs::remove_dir_all(&dir).ok(); + } } diff --git a/rust/src/models/readers/chunked/csv.rs b/rust/src/models/readers/chunked/csv.rs index 3caa2b9..0b4c2a6 100644 --- a/rust/src/models/readers/chunked/csv.rs +++ b/rust/src/models/readers/chunked/csv.rs @@ -9,14 +9,15 @@ //! Reads a directory of `-NNNNNNNNNN.csv` files emitted by //! `ChunkedCsvWriter` and presents them as an ordered iterator of `Table`s. //! Each chunk file is a complete, independently parseable CSV (header + -//! rows). The reader sorts files by their numeric index so consumers see -//! batches in write order. +//! rows). The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! Inherits [`ChunkedTableReader::par_load_batched`] for sync parallel //! decode across chunk files. -use std::fs::{self, File}; -use std::io::{self, BufReader}; +use std::io; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -25,6 +26,7 @@ use minarrow::{ColumnSelection, Concatenate, SuperTable, Table}; use crate::models::decoders::csv::CsvDecodeOptions; use crate::models::readers::csv::CsvReader; +use crate::models::readers::object_storage::{StorageReader, list_files}; use crate::traits::chunked_table_reader::ChunkedTableReader; /// Per-format `Options` for `ChunkedCsvReader::open`. @@ -78,9 +80,7 @@ impl ChunkedTableReader for ChunkedCsvReader { fn list_paths>(dir: P, base: &str) -> io::Result> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for entry in fs::read_dir(dir.as_ref())? { - let entry = entry?; - let path = entry.path(); + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -93,13 +93,17 @@ impl ChunkedTableReader for ChunkedCsvReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } fn read_chunk(&self, path: &Path) -> io::Result
{ - let file = File::open(path)?; - let reader = BufReader::new(file); + let reader = StorageReader::open(path)?; let mut csv = CsvReader::from_reader( reader, self.options.decode.clone(), @@ -161,6 +165,8 @@ impl Iterator for ChunkedCsvReader { #[cfg(test)] mod tests { + use std::fs; + use super::*; use crate::models::encoders::csv::CsvEncodeOptions; use crate::models::writers::chunked::csv::ChunkedCsvWriter; diff --git a/rust/src/models/readers/chunked/parquet.rs b/rust/src/models/readers/chunked/parquet.rs index 1fc0f14..f911374 100644 --- a/rust/src/models/readers/chunked/parquet.rs +++ b/rust/src/models/readers/chunked/parquet.rs @@ -9,8 +9,10 @@ //! Reads a directory of `-NNNNNNNNNN.parquet` files emitted by //! `ChunkedParquetWriter` and presents them as an ordered iterator of //! `Table`s. Each chunk file is a complete, independently parseable -//! Parquet file. The reader sorts files by their numeric index so -//! consumers see batches in write order. +//! Parquet file. The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! ## Concurrency model //! @@ -22,14 +24,13 @@ //! across worker threads and returns a `SuperTable` with batches in //! write order. -use std::fs::{self, File}; -use std::io::BufReader; use std::path::{Path, PathBuf}; use minarrow::Table; use crate::error::IoError; use crate::models::readers::parquet::{load_parquet_table_cols, load_parquet_table}; +use crate::models::readers::object_storage::{StorageReader, list_files}; use crate::traits::chunked_table_reader::ChunkedTableReader; use std::sync::Arc; @@ -61,9 +62,7 @@ impl ChunkedTableReader for ChunkedParquetReader { fn list_paths>(dir: P, base: &str) -> Result, IoError> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for entry in fs::read_dir(dir.as_ref())? { - let entry = entry?; - let path = entry.path(); + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -76,18 +75,21 @@ impl ChunkedTableReader for ChunkedParquetReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } fn read_chunk(&self, path: &Path) -> Result { - let file = File::open(path).map_err(IoError::from)?; - load_parquet_table(BufReader::new(file)) + load_parquet_table(StorageReader::open(path)?) } fn read_chunk_cols(&self, path: &Path, columns: &[&str]) -> Result { - let file = File::open(path).map_err(IoError::from)?; - load_parquet_table_cols(BufReader::new(file), columns) + load_parquet_table_cols(StorageReader::open(path)?, columns) } fn load_batched_cols(self, columns: &[&str]) -> Result { @@ -119,6 +121,8 @@ impl Iterator for ChunkedParquetReader { #[cfg(test)] mod tests { + use std::fs; + use super::*; use crate::models::writers::chunked::parquet::ChunkedParquetWriter; use crate::traits::chunked_table_writer::ChunkedTableWriter; diff --git a/rust/src/models/readers/csv.rs b/rust/src/models/readers/csv.rs index f2df97a..89703ed 100644 --- a/rust/src/models/readers/csv.rs +++ b/rust/src/models/readers/csv.rs @@ -14,8 +14,8 @@ //! See `CsvDecodeOptions` for configuration. use crate::models::decoders::csv::{CsvDecodeOptions, decode_csv}; +use crate::models::readers::object_storage::StorageReader; use minarrow::{Field, SuperTable, Table}; -use std::fs::File; use std::io::{self, BufRead, BufReader}; use std::path::Path; use std::sync::Arc; @@ -32,16 +32,17 @@ pub struct CsvReader { finished: bool, } -impl CsvReader> { +impl CsvReader { /// Open a CSV file at the given path with the given options. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. pub fn from_path>( path: P, options: CsvDecodeOptions, batch_size: usize, ) -> io::Result { - let file = File::open(path)?; - let reader = BufReader::new(file); - Ok(Self::from_reader(reader, options, batch_size)) + Ok(Self::from_reader(StorageReader::open(path)?, options, batch_size)) } } diff --git a/rust/src/models/readers/ipc/file_table.rs b/rust/src/models/readers/ipc/file_table.rs index 8382760..8a3b799 100644 --- a/rust/src/models/readers/ipc/file_table.rs +++ b/rust/src/models/readers/ipc/file_table.rs @@ -22,113 +22,10 @@ use minarrow::structs::shared_buffer::SharedBuffer; /// for e.g., the MMAP version can read millions of rows in microseconds, and very large volumes in milliseconds. /// - **Flexibility**: this standard reader is more flexible as it is not tied to memory-mapped shared memory. use std::collections::HashSet; -use std::fs::File; use std::io; -#[cfg(feature = "arena")] -use std::mem::MaybeUninit; -#[cfg(unix)] -use std::os::unix::fs::FileExt; -#[cfg(windows)] -use std::os::windows::fs::FileExt; use std::path::Path; use std::sync::Arc; -/// Read exactly `buf.len()` bytes starting at `offset` without touching -/// the file's seek position. Backed by `pread(2)` on Unix and a -/// `seek_read` loop on Windows so callers can issue many positional -/// reads against the same shared `&File` without serialising through -/// `&mut self` or reopening the file. The wrapper exists only because -/// `seek_read` on Windows can return short, so callers can't share a -/// single-call shape with Unix's `read_exact_at`. -#[cfg(unix)] -fn read_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { - file.read_exact_at(buf, offset) -} - -/// Read directly into spare capacity without first claiming that its bytes -/// are initialised. Unix `pread` accepts a raw output pointer, preserving the -/// arena fast path without a full-buffer zeroing pass. -#[cfg(all(unix, feature = "arena"))] -fn read_at_uninit(file: &File, buf: &mut [MaybeUninit], offset: u64) -> io::Result<()> { - use std::os::fd::AsRawFd; - - let mut filled = 0usize; - while filled < buf.len() { - let file_offset: libc::off_t = (offset + filled as u64).try_into().map_err(|_| { - io::Error::new(io::ErrorKind::InvalidInput, "file offset exceeds off_t") - })?; - let n = unsafe { - libc::pread( - file.as_raw_fd(), - buf.as_mut_ptr().add(filled).cast(), - buf.len() - filled, - file_offset, - ) - }; - if n == 0 { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "file ended before requested offset/length", - )); - } - if n < 0 { - let error = io::Error::last_os_error(); - if error.kind() == io::ErrorKind::Interrupted { - continue; - } - return Err(error); - } - filled += n as usize; - } - Ok(()) -} - -#[cfg(all(windows, feature = "arena"))] -fn read_at_uninit(file: &File, buf: &mut [MaybeUninit], offset: u64) -> io::Result<()> { - // Windows FileExt has no spare-capacity API. Initialise once before - // passing the region through its safe `&mut [u8]` interface. - for byte in &mut *buf { - byte.write(0); - } - let initialised = - unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr().cast::(), buf.len()) }; - read_at(file, initialised, offset) -} - -#[cfg(windows)] -fn read_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { - let mut total = 0; - while total < buf.len() { - let n = file.seek_read(&mut buf[total..], offset + total as u64)?; - if n == 0 { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "file ended before requested offset/length", - )); - } - total += n; - } - Ok(()) -} - -/// Positional file reads are available on Unix and Windows. Other targets, -/// such as WebAssembly in the browser, report `ErrorKind::Unsupported`. -#[cfg(not(any(unix, windows)))] -fn read_at(_file: &File, _buf: &mut [u8], _offset: u64) -> io::Result<()> { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "positional file reads are not supported on this target", - )) -} - -#[cfg(all(not(any(unix, windows)), feature = "arena"))] -fn read_at_uninit(_file: &File, _buf: &mut [MaybeUninit], _offset: u64) -> io::Result<()> { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "positional file reads are not supported on this target", - )) -} - use flatbuffers::Vector; use minarrow::{Field, SuperTable, Table}; @@ -142,6 +39,7 @@ use crate::models::decoders::ipc::parser::{ }; use crate::models::decoders::limits::DecodeLimits; use crate::models::readers::ipc::window::window_table; +use crate::models::readers::object_storage::StorageSource; /// Footer-declared block entry (i.e., offsets/lengths) for a dictionary or record batch. #[derive(Debug, Clone)] @@ -163,11 +61,12 @@ struct IPCFileBlock { /// shared memory. #[derive(Clone)] pub struct FileTableReader { - /// Open file handle shared across the reader's lifetime. All block - /// reads go through positional reads (`pread`/`seek_read`) on this - /// handle, so opening the file once amortises across the footer - /// parse, every dictionary block, and every record batch read. - file: Arc, + /// Storage handle shared across the reader's lifetime. All block + /// reads go through positional reads on this handle: `pread`/`seek_read` + /// for a local file and one ranged request per read for an object. + /// Opening once amortises across the footer parse, every dictionary + /// block, and every record batch read. + source: Arc, /// Arrow schema fields from the file footer schema: Vec>, /// Footer-declared dictionary block table @@ -190,9 +89,13 @@ pub struct FileTableReader { impl FileTableReader { /// Open an Arrow IPC file into heap memory and parse footer/schema/block tables. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. An object is read with ranged requests for the footer and + /// for each block as it is read. pub fn open>(path: P) -> io::Result { - let file = File::open(path.as_ref())?; - let file_len = file.metadata()?.len() as usize; + let source = StorageSource::open(path)?; + let file_len = source.len()? as usize; if file_len < 12 { return Err(io::Error::new( @@ -203,7 +106,7 @@ impl FileTableReader { // Read only the trailing 10 bytes: footer_len (4) + closing magic (6). let mut tail = [0u8; 10]; - read_at(&file, &mut tail, (file_len - 10) as u64)?; + source.read_at(&mut tail, (file_len - 10) as u64)?; if &tail[4..] != ARROW_MAGIC_NUMBER { return Err(io::Error::new( @@ -230,7 +133,7 @@ impl FileTableReader { // Read just the footer let mut footer_buf = vec![0u8; footer_len]; - read_at(&file, &mut footer_buf, footer_start as u64)?; + source.read_at(&mut footer_buf, footer_start as u64)?; let footer_msg = flatbuffers::root::(&footer_buf) .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("bad footer: {e}")))?; @@ -301,7 +204,7 @@ impl FileTableReader { .unwrap_or(0); let mut rdr = Self { - file: Arc::new(file), + source: Arc::new(source), schema: fields, dict_blocks, record_blocks, @@ -474,8 +377,9 @@ impl FileTableReader { /// shared `StreamArena` and a window over the just-written region /// is returned. Without the feature a fresh 64-byte aligned `Vec64` /// is allocated per call. Both paths use positional I/O on the - /// reader's shared `Arc` so concurrent block reads do not - /// need `&mut self` and the file is opened only once per reader. + /// reader's shared `Arc` so concurrent block reads do + /// not need `&mut self` and the file or object is opened only once per + /// reader. #[cfg(feature = "arena")] fn read_block(&self, blk: &IPCFileBlock) -> io::Result { let total = blk.meta_bytes + blk.body_bytes; @@ -487,7 +391,7 @@ impl FileTableReader { arena.ensure_capacity(total); let start = arena.write_pos(); let spare = arena.spare_uninit(); - read_at_uninit(&self.file, &mut spare[..total], blk.offset as u64)?; + self.source.read_at_uninit(&mut spare[..total], blk.offset as u64)?; // SAFETY: read_at_uninit filled the entire requested region. unsafe { arena.advance(total) }; let shared = arena.window(start, total); @@ -507,7 +411,7 @@ impl FileTableReader { unsafe { buf.set_len(total); } - read_at(&self.file, &mut buf, blk.offset as u64)?; + self.source.read_at(&mut buf, blk.offset as u64)?; Ok(SharedBuffer::from_vec64(buf)) } diff --git a/rust/src/models/readers/json.rs b/rust/src/models/readers/json.rs index 653049c..36f42a3 100644 --- a/rust/src/models/readers/json.rs +++ b/rust/src/models/readers/json.rs @@ -15,7 +15,6 @@ //! //! See [`JsonDecodeOptions`] for schema handling and type control. -use std::fs::File; use std::io::{self, BufRead, BufReader}; use std::path::Path; use std::sync::Arc; @@ -29,6 +28,7 @@ use crate::models::decoders::json::{ JsonDecodeOptions, append_ndjson_line, decode_json, finish_table, make_builders, make_field_map, }; use crate::models::encoders::json::JsonFormat; +use crate::models::readers::object_storage::StorageReader; /// Reads JSON data into Minarrow Tables. /// @@ -58,17 +58,19 @@ pub struct JsonReader { line: Vec, } -impl JsonReader> { +impl JsonReader { /// Open a JSON file at the given path. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. pub fn from_path>( path: P, format: JsonFormat, options: JsonDecodeOptions, batch_size: usize, ) -> io::Result { - let file = File::open(path)?; Ok(Self::from_reader( - BufReader::new(file), + StorageReader::open(path)?, format, options, batch_size, diff --git a/rust/src/models/readers/object_storage.rs b/rust/src/models/readers/object_storage.rs new file mode 100644 index 0000000..e818176 --- /dev/null +++ b/rust/src/models/readers/object_storage.rs @@ -0,0 +1,378 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Storage readers +//! +//! Read handles for the file readers, over a local file or, with the +//! `object_storage` feature, an object in object storage. Each reader opens +//! its handle from the path it is given, and the decoders work against the +//! standard `Read`, `BufRead` and `Seek` traits the handles implement. +//! +//! ## Handles +//! - [`StorageReader`]: sequential reads with seeking, for the CSV, JSON and +//! Parquet readers. An object is fetched whole when the reader opens. +//! - [`StorageSource`]: positional reads from a shared handle, for the +//! Arrow IPC file reader. An object is read with one ranged request per +//! read, which fetches the footer and then each record batch as it is +//! read. +//! - [`list_dir`]: directory listing for the chunked readers. Object storage +//! has no directories, and listing an object prefix returns the objects +//! immediately beneath it. +//! +//! Object storage requests block the calling thread on the caller's Tokio +//! runtime. See [`crate::object_storage`] for the runtime requirement. + +use std::fs::{self, File}; +use std::io::{self, BufRead, BufReader, Read, Seek, SeekFrom}; +#[cfg(feature = "arena")] +use std::mem::MaybeUninit; +#[cfg(unix)] +use std::os::unix::fs::FileExt; +#[cfg(windows)] +use std::os::windows::fs::FileExt; +use std::path::{Path, PathBuf}; + +#[cfg(feature = "object_storage")] +use bytes::Bytes; +#[cfg(feature = "object_storage")] +use futures_util::TryStreamExt; +#[cfg(feature = "object_storage")] +use object_store::{ObjectMeta, ObjectStoreExt}; + +#[cfg(feature = "object_storage")] +use crate::object_storage::{ObjectLocation, current_handle, io_error, object_url, resolve}; + +/// Sequential reader over a local file or a fetched object. +/// +/// Implements `Read`, `BufRead` and `Seek`. +pub enum StorageReader { + /// A buffered local file. + File(BufReader), + /// The contents of an object, fetched when the reader opened, with the + /// current read position. + #[cfg(feature = "object_storage")] + Object { bytes: Bytes, position: usize }, +} + +impl StorageReader { + /// Open the file or object URL at `path` for reading. + /// + /// An object URL fetches the whole object before returning. + pub fn open>(path: P) -> io::Result { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(path.as_ref()) { + let location = resolve(&url)?; + let bytes = current_handle()? + .block_on(async { location.store.get(&location.path).await?.bytes().await }) + .map_err(io_error)?; + return Ok(Self::Object { bytes, position: 0 }); + } + Ok(Self::File(BufReader::new(File::open(path)?))) + } +} + +impl Read for StorageReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match self { + Self::File(file) => file.read(buf), + #[cfg(feature = "object_storage")] + Self::Object { .. } => { + let available = self.fill_buf()?; + let n = available.len().min(buf.len()); + buf[..n].copy_from_slice(&available[..n]); + self.consume(n); + Ok(n) + } + } + } +} + +impl BufRead for StorageReader { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + match self { + Self::File(file) => file.fill_buf(), + #[cfg(feature = "object_storage")] + Self::Object { bytes, position } => Ok(&bytes[(*position).min(bytes.len())..]), + } + } + + fn consume(&mut self, amount: usize) { + match self { + Self::File(file) => file.consume(amount), + #[cfg(feature = "object_storage")] + Self::Object { position, .. } => *position += amount, + } + } +} + +impl Seek for StorageReader { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + match self { + Self::File(file) => file.seek(pos), + #[cfg(feature = "object_storage")] + Self::Object { bytes, position } => { + let target = match pos { + SeekFrom::Start(offset) => Some(offset as i128), + SeekFrom::End(delta) => Some(bytes.len() as i128 + delta as i128), + SeekFrom::Current(delta) => Some(*position as i128 + delta as i128), + } + .filter(|target| *target >= 0) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "seek before the start of the object", + ) + })?; + *position = target as usize; + Ok(target as u64) + } + } + } +} + +/// Positional reader over a local file or an object. +/// +/// Positional reads take `&self` and leave no read position behind, which +/// lets many reads share one handle across threads. Local files use +/// `pread(2)` on Unix and `seek_read` on Windows. Objects issue one ranged +/// request per read. +pub enum StorageSource { + /// A local file. + File(File), + /// An object and its size, recorded when the source opened. + #[cfg(feature = "object_storage")] + Object { location: ObjectLocation, len: u64 }, +} + +impl StorageSource { + /// Open the file or object URL at `path` for positional reads. + /// + /// An object URL requests the object's metadata to record its size. + pub fn open>(path: P) -> io::Result { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(path.as_ref()) { + let location = resolve(&url)?; + let meta = current_handle()? + .block_on(location.store.head(&location.path)) + .map_err(io_error)?; + return Ok(Self::Object { + location, + len: meta.size, + }); + } + Ok(Self::File(File::open(path)?)) + } + + /// Size of the file or object in bytes. + pub fn len(&self) -> io::Result { + match self { + Self::File(file) => Ok(file.metadata()?.len()), + #[cfg(feature = "object_storage")] + Self::Object { len, .. } => Ok(*len), + } + } + + /// Whether the file or object is empty. + pub fn is_empty(&self) -> io::Result { + Ok(self.len()? == 0) + } + + /// Read exactly `buf.len()` bytes starting at `offset`. + pub fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<()> { + match self { + Self::File(file) => read_file_at(file, buf, offset), + #[cfg(feature = "object_storage")] + Self::Object { location, .. } => { + let range = offset..offset + buf.len() as u64; + let bytes = current_handle()? + .block_on(location.store.get_range(&location.path, range)) + .map_err(io_error)?; + if bytes.len() != buf.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "object ended before requested offset/length", + )); + } + buf.copy_from_slice(&bytes); + Ok(()) + } + } + } + + /// Read exactly `buf.len()` bytes starting at `offset` into spare + /// capacity, without first initialising it where the platform allows. + #[cfg(feature = "arena")] + pub fn read_at_uninit(&self, buf: &mut [MaybeUninit], offset: u64) -> io::Result<()> { + match self { + Self::File(file) => read_file_at_uninit(file, buf, offset), + #[cfg(feature = "object_storage")] + Self::Object { .. } => self.read_at(initialise(buf), offset), + } + } +} + +/// Paths of the entries immediately inside the directory or object prefix +/// `dir`. +/// +/// Object entries are returned as object URLs beneath `dir`, which open +/// with the readers in the same way as local paths. +pub fn list_dir>(dir: P) -> io::Result> { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(dir.as_ref()) { + let location = resolve(&url)?; + let listing = current_handle()? + .block_on(location.store.list_with_delimiter(Some(&location.path))) + .map_err(io_error)?; + let base = url.trim_end_matches('/'); + return Ok(listing + .objects + .iter() + .filter_map(|meta| meta.location.filename()) + .map(|name| PathBuf::from(format!("{base}/{name}"))) + .collect()); + } + fs::read_dir(dir)? + .map(|entry| entry.map(|entry| entry.path())) + .collect() +} + +/// Paths of every file beneath the directory or object prefix `dir`, +/// including files in nested directories, in no particular order. +/// +/// Object entries are returned as object URLs beneath `dir`, as for +/// [`list_dir`]. On object storage this is one listing of every object under +/// the prefix. +pub fn list_files>(dir: P) -> io::Result> { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(dir.as_ref()) { + let location = resolve(&url)?; + let objects: Vec = current_handle()? + .block_on(location.store.list(Some(&location.path)).try_collect()) + .map_err(io_error)?; + let base = url.trim_end_matches('/'); + return Ok(objects + .iter() + .filter_map(|meta| meta.location.prefix_match(&location.path)) + .map(|parts| { + let relative: Vec = parts.map(|part| part.as_ref().to_string()).collect(); + PathBuf::from(format!("{base}/{}", relative.join("/"))) + }) + .collect()); + } + let mut files = Vec::new(); + let mut pending = vec![dir.as_ref().to_path_buf()]; + while let Some(dir) = pending.pop() { + for entry in fs::read_dir(&dir)? { + let path = entry?.path(); + if path.is_dir() { + pending.push(path); + } else { + files.push(path); + } + } + } + Ok(files) +} + +/// Read exactly `buf.len()` bytes starting at `offset` without touching +/// the file's seek position. Backed by `pread(2)` on Unix and a +/// `seek_read` loop on Windows so callers can issue many positional +/// reads against the same shared `&File` without serialising through +/// `&mut self` or reopening the file. The wrapper exists only because +/// `seek_read` on Windows can return short, so callers can't share a +/// single-call shape with Unix's `read_exact_at`. +#[cfg(unix)] +fn read_file_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { + file.read_exact_at(buf, offset) +} + +/// Read directly into spare capacity without first claiming that its bytes +/// are initialised. Unix `pread` accepts a raw output pointer, preserving the +/// arena fast path without a full-buffer zeroing pass. +#[cfg(all(unix, feature = "arena"))] +fn read_file_at_uninit(file: &File, buf: &mut [MaybeUninit], offset: u64) -> io::Result<()> { + use std::os::fd::AsRawFd; + + let mut filled = 0usize; + while filled < buf.len() { + let file_offset: libc::off_t = (offset + filled as u64).try_into().map_err(|_| { + io::Error::new(io::ErrorKind::InvalidInput, "file offset exceeds off_t") + })?; + let n = unsafe { + libc::pread( + file.as_raw_fd(), + buf.as_mut_ptr().add(filled).cast(), + buf.len() - filled, + file_offset, + ) + }; + if n == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "file ended before requested offset/length", + )); + } + if n < 0 { + let error = io::Error::last_os_error(); + if error.kind() == io::ErrorKind::Interrupted { + continue; + } + return Err(error); + } + filled += n as usize; + } + Ok(()) +} + +#[cfg(all(windows, feature = "arena"))] +fn read_file_at_uninit(file: &File, buf: &mut [MaybeUninit], offset: u64) -> io::Result<()> { + // Windows FileExt has no spare-capacity API. Initialise once before + // passing the region through its safe `&mut [u8]` interface. + read_file_at(file, initialise(buf), offset) +} + +#[cfg(windows)] +fn read_file_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { + let mut total = 0; + while total < buf.len() { + let n = file.seek_read(&mut buf[total..], offset + total as u64)?; + if n == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "file ended before requested offset/length", + )); + } + total += n; + } + Ok(()) +} + +/// Positional file reads are available on Unix and Windows. Other targets, +/// such as WebAssembly in the browser, report `ErrorKind::Unsupported`. +#[cfg(not(any(unix, windows)))] +fn read_file_at(_file: &File, _buf: &mut [u8], _offset: u64) -> io::Result<()> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "positional file reads are not supported on this target", + )) +} + +#[cfg(all(not(any(unix, windows)), feature = "arena"))] +fn read_file_at_uninit(_file: &File, _buf: &mut [MaybeUninit], _offset: u64) -> io::Result<()> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "positional file reads are not supported on this target", + )) +} + +/// Zero `buf` and return it as an initialised byte slice. +#[cfg(all(feature = "arena", any(windows, feature = "object_storage")))] +fn initialise(buf: &mut [MaybeUninit]) -> &mut [u8] { + for byte in &mut *buf { + byte.write(0); + } + unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr().cast::(), buf.len()) } +} diff --git a/rust/src/models/sinks/table_sink.rs b/rust/src/models/sinks/table_sink.rs index 8851ddb..7f2d792 100644 --- a/rust/src/models/sinks/table_sink.rs +++ b/rust/src/models/sinks/table_sink.rs @@ -53,7 +53,7 @@ pub type TableSink64 = GTableSink>; /// written into the wider ecosystem. pub struct GTableSink where - W: AsyncWrite + Unpin + Send + Sync + 'static, + W: AsyncWrite + Unpin + Send + 'static, B: StreamBuffer + Unpin + 'static, { pub(crate) schema: Vec, @@ -71,7 +71,7 @@ where impl GTableSink where - W: AsyncWrite + Unpin + Send + Sync + 'static, + W: AsyncWrite + Unpin + Send + 'static, B: StreamBuffer + std::fmt::Debug + Unpin + 'static, { /// Create a new generic Arrow Table writer. Pass `None` for @@ -136,7 +136,7 @@ where impl Sink for GTableSink where - W: AsyncWrite + Unpin + Send + Sync + 'static, + W: AsyncWrite + Unpin + Send + 'static, B: StreamBuffer + std::fmt::Debug + Unpin + 'static, { type Error = io::Error; diff --git a/rust/src/models/streams/object_store.rs b/rust/src/models/streams/object_store.rs new file mode 100644 index 0000000..d0b2c80 --- /dev/null +++ b/rust/src/models/streams/object_store.rs @@ -0,0 +1,108 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object storage byte stream adapters +//! +//! Async read and write adapters over objects in object storage, for the +//! async Arrow IPC readers and writers. +//! +//! - [`ObjectByteStream`] reads an object as a chunked byte stream and an +//! `AsyncRead`, ready for `TableReader`. +//! - [`ObjectWrite`] writes an object as an `AsyncWrite`, ready for +//! `TableWriter`. Objects up to one part are uploaded with a single +//! request, and larger objects stream as a multipart upload. The object +//! is completed by `poll_shutdown`. +//! +//! Both adapters use the object storage HTTP client, which is separate +//! from the HTTP/2 transport in `streams::http`. See +//! [`crate::object_storage`] for URL forms and credentials. + +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use object_store::ObjectStoreExt; +use object_store::buffered::{BufReader, BufWriter}; +use tokio::io::AsyncWrite; + +use crate::enums::BufferChunkSize; +use crate::models::streams::async_read::AsyncReadByteStream; +use crate::object_storage::{io_error, resolve}; + +/// Chunked byte stream over one object. +pub type ObjectByteStream = AsyncReadByteStream; + +impl ObjectByteStream { + /// Open the object at `url` for streaming reads. + /// + /// Requests the object's metadata, then reads its contents in ranged + /// requests as the stream is polled. + pub async fn open(url: &str, size: BufferChunkSize) -> io::Result { + let location = resolve(url)?; + let meta = location + .store + .head(&location.path) + .await + .map_err(io_error)?; + Ok(AsyncReadByteStream::new( + BufReader::new(location.store, &meta), + size, + )) + } +} + +/// Async writer for one object. +pub struct ObjectWrite { + inner: BufWriter, + /// Set once `poll_shutdown` has been called. The buffered writer completes + /// the object during shutdown and does not accept flushes after it + /// begins. + shutting_down: bool, +} + +impl ObjectWrite { + /// Create the object at `url`, replacing any existing one. + /// + /// No request is made until the first part fills or the writer shuts + /// down. + pub fn create(url: &str) -> io::Result { + let location = resolve(url)?; + let inner = BufWriter::with_capacity(location.store, location.path, location.part_size) + .with_max_concurrency(location.max_concurrent_parts); + Ok(Self { + inner, + shutting_down: false, + }) + } +} + +impl AsyncWrite for ObjectWrite { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + /// Flush buffered bytes into the upload. + /// + /// Once shutdown has begun, the remaining bytes are uploaded by shutdown + /// itself and flushing does nothing. Writers that flush before each + /// shutdown poll, such as `GTableSink`, depend on this. + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.shutting_down { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.inner).poll_flush(cx) + } + + /// Complete the object, uploading any buffered bytes. + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.shutting_down = true; + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} diff --git a/rust/src/models/writers/chunked/arrow.rs b/rust/src/models/writers/chunked/arrow.rs index ed07ac4..a6e1e92 100644 --- a/rust/src/models/writers/chunked/arrow.rs +++ b/rust/src/models/writers/chunked/arrow.rs @@ -25,7 +25,6 @@ //! created for chunk writing; callers can drive this directly from any //! sync context (e.g. the engine's per-Block invocation path). -use std::fs::{self, File}; use std::io; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicU64; @@ -33,6 +32,8 @@ use std::sync::atomic::AtomicU64; use minarrow::{Field, Table}; use crate::enums::IPCMessageProtocol; +use crate::models::writers::chunked::partition::PartitionType; +use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::ipc::sync_table::SyncTableWriter; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -46,6 +47,7 @@ pub struct ChunkedArrowWriter { dir: PathBuf, base: String, next_index: AtomicU64, + partition: Option, } impl ChunkedArrowWriter { @@ -57,13 +59,21 @@ impl ChunkedArrowWriter { /// open it). pub fn new>(dir: P, base: &str) -> io::Result { let dir = dir.as_ref().to_path_buf(); - fs::create_dir_all(&dir)?; + create_dir_all(&dir)?; Ok(Self { dir, base: base.to_string(), next_index: AtomicU64::new(0), + partition: None, }) } + + /// Place each chunk in a partition directory beneath `dir`, for + /// example `/date=2026-10-14/hour=13/-NNNNNNNNNN.arrow`. + pub fn with_partition(mut self, partition: PartitionType) -> Self { + self.partition = Some(partition); + self + } } impl ChunkedTableWriter for ChunkedArrowWriter { @@ -90,12 +100,17 @@ impl ChunkedTableWriter for ChunkedArrowWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete Arrow IPC File at `path`. Shared /// per-format encode path used by both `write_chunk` and /// `par_write_all`. fn write_chunk_at(&self, path: &Path, table: &Table) -> io::Result<()> { let schema: Vec = table.cols.iter().map(|col| (*col.field).clone()).collect(); - let file = File::create(path)?; + let file = StorageWriter::create(path)?; let mut writer: SyncTableWriter<_> = SyncTableWriter::new(file, schema, IPCMessageProtocol::File, None); writer.write_table(table.clone())?; @@ -106,6 +121,8 @@ impl ChunkedTableWriter for ChunkedArrowWriter { #[cfg(test)] mod tests { + use std::fs; + use super::*; use crate::models::readers::chunked::arrow::ChunkedArrowReader; use crate::traits::chunked_table_reader::ChunkedTableReader; diff --git a/rust/src/models/writers/chunked/csv.rs b/rust/src/models/writers/chunked/csv.rs index 9acc8b4..12f5682 100644 --- a/rust/src/models/writers/chunked/csv.rs +++ b/rust/src/models/writers/chunked/csv.rs @@ -19,7 +19,6 @@ //! `ChunkedCsvReader` globs the directory and presents the chunks back //! as an ordered iterator of `Table`s. -use std::fs::{self, File}; use std::io; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicU64; @@ -27,6 +26,8 @@ use std::sync::atomic::AtomicU64; use minarrow::Table; use crate::models::encoders::csv::CsvEncodeOptions; +use crate::models::writers::chunked::partition::PartitionType; +use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::csv::CsvWriter; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -42,6 +43,7 @@ pub struct ChunkedCsvWriter { base: String, options: CsvEncodeOptions, next_index: AtomicU64, + partition: Option, } impl ChunkedCsvWriter { @@ -50,14 +52,22 @@ impl ChunkedCsvWriter { /// index. The directory is created if it does not exist. pub fn new>(dir: P, base: &str, options: CsvEncodeOptions) -> io::Result { let dir = dir.as_ref().to_path_buf(); - fs::create_dir_all(&dir)?; + create_dir_all(&dir)?; Ok(Self { dir, base: base.to_string(), options, next_index: AtomicU64::new(0), + partition: None, }) } + + /// Place each chunk in a partition directory beneath `dir`, for + /// example `/date=2026-10-14/hour=13/-NNNNNNNNNN.csv`. + pub fn with_partition(mut self, partition: PartitionType) -> Self { + self.partition = Some(partition); + self + } } impl ChunkedTableWriter for ChunkedCsvWriter { @@ -84,11 +94,16 @@ impl ChunkedTableWriter for ChunkedCsvWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete CSV file at `path` using the writer's /// configured encode options. Shared per-format encode path used by /// both `write_chunk` and `par_write_all`. fn write_chunk_at(&self, path: &Path, table: &Table) -> io::Result<()> { - let file = File::create(path)?; + let file = StorageWriter::create(path)?; let mut writer = CsvWriter::with_options(file, self.options.clone()); writer.write_table(table)?; writer.flush()?; @@ -98,6 +113,8 @@ impl ChunkedTableWriter for ChunkedCsvWriter { #[cfg(test)] mod tests { + use std::fs; + use super::*; use crate::models::readers::chunked::csv::{ChunkedCsvReadOptions, ChunkedCsvReader}; use crate::traits::chunked_table_reader::ChunkedTableReader; diff --git a/rust/src/models/writers/chunked/parquet.rs b/rust/src/models/writers/chunked/parquet.rs index fb60fb8..2e721bb 100644 --- a/rust/src/models/writers/chunked/parquet.rs +++ b/rust/src/models/writers/chunked/parquet.rs @@ -18,7 +18,7 @@ //! across batches. The matching `ChunkedParquetReader` globs the directory //! and presents the chunks back as an ordered iterator of `Table`s. -use std::fs::{self, File}; +use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicU64; @@ -26,6 +26,8 @@ use minarrow::Table; use crate::compression::Compression; use crate::error::IoError; +use crate::models::writers::chunked::partition::PartitionType; +use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::parquet::write_parquet_table; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -36,6 +38,7 @@ pub struct ChunkedParquetWriter { base: String, compression: Option, next_index: AtomicU64, + partition: Option, } impl ChunkedParquetWriter { @@ -48,14 +51,22 @@ impl ChunkedParquetWriter { compression: Option, ) -> Result { let dir = dir.as_ref().to_path_buf(); - fs::create_dir_all(&dir)?; + create_dir_all(&dir)?; Ok(Self { dir, base: base.to_string(), compression, next_index: AtomicU64::new(0), + partition: None, }) } + + /// Place each chunk in a partition directory beneath `dir`, for + /// example `/date=2026-10-14/hour=13/-NNNNNNNNNN.parquet`. + pub fn with_partition(mut self, partition: PartitionType) -> Self { + self.partition = Some(partition); + self + } } impl ChunkedTableWriter for ChunkedParquetWriter { @@ -82,18 +93,27 @@ impl ChunkedTableWriter for ChunkedParquetWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete Parquet file at `path` using the /// writer's configured compression. Shared per-format encode path /// used by both `write_chunk` and `par_write_all`. fn write_chunk_at(&self, path: &Path, table: &Table) -> Result<(), IoError> { - let file = File::create(path)?; - write_parquet_table(table, file, self.compression)?; + let mut file = StorageWriter::create(path)?; + write_parquet_table(table, &mut file, self.compression)?; + // Flushing completes the object when `path` is an object URL. + file.flush()?; Ok(()) } } #[cfg(test)] mod tests { + use std::fs; + use super::*; use crate::models::readers::chunked::parquet::ChunkedParquetReader; use crate::traits::chunked_table_reader::ChunkedTableReader; diff --git a/rust/src/models/writers/chunked/partition.rs b/rust/src/models/writers/chunked/partition.rs new file mode 100644 index 0000000..997d10e --- /dev/null +++ b/rust/src/models/writers/chunked/partition.rs @@ -0,0 +1,137 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Chunk partitioning +//! +//! Partition paths for the chunked writers. A partitioned writer places each +//! chunk in a subdirectory of its root, for example +//! `/date=2026-10-14/hour=13/-NNNNNNNNNN.parquet`. +//! +//! ## Context +//! Readers locate chunks by listing a directory or object prefix. The +//! partition path determines how much a reader needs to list to find the +//! data for one period of time. This matters most on object storage, where +//! each page of a listing is a separate request. +//! The chunked readers include chunks in nested partition directories, +//! and opening a reader on a partition path such as `/date=2026-10-14` +//! reads that partition alone. +//! +//! ## Partition types +//! - [`PartitionType::Auto`] derives the path from the time of each write in +//! the `key=value` directory convention. Query engines such as DuckDB, +//! Polars and Spark recognise this convention, using it to prune by path +//! and exposing the values as columns. +//! - [`PartitionType::Custom`] takes the path from a caller-supplied +//! function, for schemes keyed on other columns such as venue or symbol. +//! +//! The time is supplied by the caller through +//! [`ChunkedTableWriter::write_chunk_for_time`](crate::traits::chunked_table_writer::ChunkedTableWriter::write_chunk_for_time), +//! for example from the event-time window that produced the chunk. The +//! other write methods use the current system time. + +use std::fmt; +use std::sync::Arc; + +use minarrow::Table; + +const NANOS_PER_SECOND: u64 = 1_000_000_000; +const SECONDS_PER_DAY: u64 = 86_400; +const SECONDS_PER_HOUR: u64 = 3_600; + +/// How a chunked writer groups its chunks into partition directories. +#[derive(Clone)] +pub enum PartitionType { + /// `date=YYYY-MM-DD/hour=HH` in UTC, from the time of each write. + /// + /// Hourly partitions bound the listing a reader needs for any period. + Auto, + /// Partition path returned by the function for each write, given the + /// Table and its time in epoch nanoseconds. + /// + /// The path is `/`-separated and relative to the writer's directory, + /// for example `venue=XNAS/date=2026-10-14`. + Custom(Arc String + Send + Sync>), +} + +impl PartitionType { + /// Partition path for a write of `table` at `time_ns` epoch nanoseconds. + pub fn partition(&self, table: &Table, time_ns: u64) -> String { + match self { + Self::Auto => auto_partition(time_ns), + Self::Custom(partition) => partition(table, time_ns), + } + } +} + +impl fmt::Debug for PartitionType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Auto => f.write_str("Auto"), + Self::Custom(_) => f.write_str("Custom"), + } + } +} + +/// Format an epoch-nanosecond time as `date=YYYY-MM-DD/hour=HH` in UTC. +fn auto_partition(time_ns: u64) -> String { + let epoch_secs = time_ns / NANOS_PER_SECOND; + let (year, month, day) = civil_from_days((epoch_secs / SECONDS_PER_DAY) as i64); + let hour = epoch_secs % SECONDS_PER_DAY / SECONDS_PER_HOUR; + format!("date={year:04}-{month:02}-{day:02}/hour={hour:02}") +} + +/// The civil date `(year, month, day)` for a day count from the Unix epoch, +/// using Howard Hinnant's `civil_from_days` algorithm. +fn civil_from_days(days: i64) -> (i64, u32, u32) { + let z = days + 719_468; + let era = (if z >= 0 { z } else { z - 146_096 }) / 146_097; + let doe = z - era * 146_097; + let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365; + let year = yoe + era * 400; + let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); + let mp = (5 * doy + 2) / 153; + let day = (doy - (153 * mp + 2) / 5 + 1) as u32; + let month = (if mp < 10 { mp + 3 } else { mp - 9 }) as u32; + (if month <= 2 { year + 1 } else { year }, month, day) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 2026-02-09T07:48:36.123456789Z + const SAMPLE: u64 = 1_770_623_316_123_456_789; + + #[test] + fn auto_partition_formats_date_and_hour_in_utc() { + assert_eq!(auto_partition(SAMPLE), "date=2026-02-09/hour=07"); + } + + #[test] + fn auto_partition_at_the_epoch_is_the_first_hour_of_1970() { + assert_eq!(auto_partition(0), "date=1970-01-01/hour=00"); + } + + #[test] + fn auto_partition_handles_the_leap_day() { + // 2024-02-29T23:59:59Z + let time = 1_709_251_199 * NANOS_PER_SECOND; + assert_eq!(auto_partition(time), "date=2024-02-29/hour=23"); + assert_eq!( + auto_partition(time + NANOS_PER_SECOND), + "date=2024-03-01/hour=00" + ); + } + + #[test] + fn custom_partition_receives_the_table_and_time() { + let partition = PartitionType::Custom(Arc::new(|table: &Table, time: u64| { + format!("table={}/time={time}", table.name) + })); + let table = Table::new("trades".to_string(), None); + assert_eq!(partition.partition(&table, 42), "table=trades/time=42"); + } +} diff --git a/rust/src/models/writers/csv.rs b/rust/src/models/writers/csv.rs index 05cf9dd..d497483 100644 --- a/rust/src/models/writers/csv.rs +++ b/rust/src/models/writers/csv.rs @@ -41,13 +41,13 @@ //! # Ok::<(), std::io::Error>(()) //! ``` -use std::fs::File; use std::io::{self, Write}; use std::path::Path; use minarrow::{SuperTable, Table}; use crate::models::encoders::csv::{CsvEncodeOptions, encode_supertable_csv, encode_table_csv}; +use crate::models::writers::object_storage::StorageWriter; /// A streaming CSV writer for [`minarrow::Table`] and [`minarrow::SuperTable`]. /// @@ -109,22 +109,26 @@ impl CsvWriter { } } -impl CsvWriter { - /// Open the given file path and return a `CsvWriter` - /// using default options. +impl CsvWriter { + /// Create the given file path and return a `CsvWriter` using default + /// options. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. An object is completed by [`CsvWriter::flush`]. pub fn to_path>(path: P) -> io::Result { - let file = File::create(path)?; - Ok(Self::new(file)) + Ok(Self::new(StorageWriter::create(path)?)) } - /// Open the given file path and return a `CsvWriter` - /// with custom options. + /// Create the given file path and return a `CsvWriter` with custom + /// options. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. An object is completed by [`CsvWriter::flush`]. pub fn to_path_with_options>( path: P, options: CsvEncodeOptions, ) -> io::Result { - let file = File::create(path)?; - Ok(Self::with_options(file, options)) + Ok(Self::with_options(StorageWriter::create(path)?, options)) } } diff --git a/rust/src/models/writers/ipc/sync_table.rs b/rust/src/models/writers/ipc/sync_table.rs index e3f467e..7876780 100644 --- a/rust/src/models/writers/ipc/sync_table.rs +++ b/rust/src/models/writers/ipc/sync_table.rs @@ -35,6 +35,7 @@ use minarrow::{Field, Table, TableV, Vec64}; use crate::compression::Compression; use crate::enums::IPCMessageProtocol; +use crate::models::writers::object_storage::StorageWriter; use crate::models::writers::ipc::table_stream::TableStreamWriter; use crate::traits::stream_buffer::StreamBuffer; use crate::utils::dict_values; @@ -152,7 +153,7 @@ pub fn write_tables_to_file_sync( tables: &[Table], schema: Vec, ) -> io::Result<()> { - let file = std::fs::File::create(file_path)?; + let file = StorageWriter::create(file_path)?; let mut writer = SyncTableWriter::<_, Vec64>::new(file, schema, IPCMessageProtocol::File, None); for table in tables { writer.write_table(table.clone())?; diff --git a/rust/src/models/writers/ipc/table.rs b/rust/src/models/writers/ipc/table.rs index ca0ff87..853ff5c 100644 --- a/rust/src/models/writers/ipc/table.rs +++ b/rust/src/models/writers/ipc/table.rs @@ -31,7 +31,7 @@ use crate::utils::dict_values; use futures_util::sink::SinkExt; use minarrow::{Field, Table, TableV}; #[cfg(not(target_family = "wasm"))] -use tokio::fs::File; +use crate::models::writers::object_storage::AsyncStorageWriter; use tokio::io::AsyncWrite; use minarrow::Vec64; @@ -44,14 +44,14 @@ use std::pin::Pin; /// Internally delegates to a wrapped [`GTableSink`] instance. pub struct TableWriter where - W: AsyncWrite + Unpin + Send + Sync + 'static, + W: AsyncWrite + Unpin + Send + 'static, { sink: GTableSink>, } impl TableWriter where - W: AsyncWrite + Unpin + Send + Sync + 'static, + W: AsyncWrite + Unpin + Send + 'static, { /// Create a new Arrow Table writer. Pass `None` for `compression` to /// write uncompressed batches; `Some(codec)` compresses every @@ -136,7 +136,7 @@ pub async fn write_tables_to_file( tables: &[Table], schema: Vec, ) -> io::Result<()> { - let file = File::create(file_path).await?; + let file = AsyncStorageWriter::create(file_path).await?; let mut writer = TableWriter::new(file, schema, IPCMessageProtocol::File, None)?; // Automatically register any Categorical dictionaries found in the tables. for table in tables { @@ -158,7 +158,7 @@ pub async fn write_table_to_file( table: &Table, schema: Vec, ) -> io::Result<()> { - let file = File::create(file_path).await?; + let file = AsyncStorageWriter::create(file_path).await?; let mut writer = TableWriter::new(file, schema, IPCMessageProtocol::File, None)?; // Automatically register any Categorical dictionaries found in the table. for (col_idx, col) in table.cols.iter().enumerate() { @@ -167,6 +167,8 @@ pub async fn write_table_to_file( } } writer.write_table(table.clone()).await?; + // Finishing writes the file footer and completes an object upload. + writer.finish().await?; Ok(()) } diff --git a/rust/src/models/writers/ipc/table_stream.rs b/rust/src/models/writers/ipc/table_stream.rs index cfd8e2f..ef26d2b 100644 --- a/rust/src/models/writers/ipc/table_stream.rs +++ b/rust/src/models/writers/ipc/table_stream.rs @@ -301,7 +301,7 @@ pub async fn write_tables_to_stream( protocol: IPCMessageProtocol, ) -> io::Result<()> where - W: AsyncWrite + Unpin + Send + Sync, + W: AsyncWrite + Unpin + Send, B: StreamBuffer + Unpin, { let mut writer = TableStreamWriter::::new(schema, protocol, None); @@ -337,7 +337,7 @@ pub async fn write_table_to_stream( protocol: IPCMessageProtocol, ) -> io::Result<()> where - W: AsyncWrite + Unpin + Send + Sync, + W: AsyncWrite + Unpin + Send, B: StreamBuffer + Unpin, { let mut writer = TableStreamWriter::::new(schema, protocol, None); diff --git a/rust/src/models/writers/json.rs b/rust/src/models/writers/json.rs index e230697..4144c0f 100644 --- a/rust/src/models/writers/json.rs +++ b/rust/src/models/writers/json.rs @@ -49,13 +49,13 @@ //! # Ok::<(), std::io::Error>(()) //! ``` -use std::fs::File; use std::io::{self, Write}; use std::path::Path; use minarrow::{SuperTable, Table}; use crate::models::encoders::json::{JsonEncodeOptions, encode_supertable_json, encode_table_json}; +use crate::models::writers::object_storage::StorageWriter; /// A streaming JSON writer for [`minarrow::Table`] and [`minarrow::SuperTable`]. /// @@ -105,11 +105,13 @@ impl JsonWriter { } } -impl JsonWriter { - /// Open the given file path and return a `JsonWriter`. +impl JsonWriter { + /// Create the given file path and return a `JsonWriter`. + /// + /// `path` is a local file path or, with the `object_storage` feature, an + /// object URL. An object is completed by [`JsonWriter::flush`]. pub fn to_path>(path: P, options: JsonEncodeOptions) -> io::Result { - let file = File::create(path)?; - Ok(Self::new(file, options)) + Ok(Self::new(StorageWriter::create(path)?, options)) } } diff --git a/rust/src/models/writers/object_storage.rs b/rust/src/models/writers/object_storage.rs new file mode 100644 index 0000000..c674cb1 --- /dev/null +++ b/rust/src/models/writers/object_storage.rs @@ -0,0 +1,187 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Storage writers +//! +//! Write handles for the file writers, over a local file or, with the +//! `object_storage` feature, an object in object storage. Each writer +//! creates its handle from the path it is given, and the encoders work +//! against the standard `Write` and `Seek` traits, or `AsyncWrite` for the +//! async Arrow IPC writers. +//! +//! ## Handles +//! - [`StorageWriter`]: sequential writes, for the CSV, JSON, Parquet and +//! Arrow IPC writers. An object uploads as it is written and completes on +//! `flush`. Object storage writes are append-only and complete once, +//! matching how the writers use their sink. +//! - [`AsyncStorageWriter`]: the async equivalent, for the async Arrow IPC +//! writers. An object completes on shutdown. +//! - [`create_dir_all`]: directory creation for the chunked writers. Object +//! storage has no directories, and creating an object prefix does nothing. +//! +//! Object storage requests from `StorageWriter` block the calling thread on +//! the caller's Tokio runtime. See [`crate::object_storage`] for the runtime +//! requirement. + +use std::fs::{self, File}; +use std::io::{self, Seek, SeekFrom, Write}; +use std::path::Path; +#[cfg(not(target_family = "wasm"))] +use std::pin::Pin; +#[cfg(not(target_family = "wasm"))] +use std::task::{Context, Poll}; + +#[cfg(not(target_family = "wasm"))] +use tokio::io::AsyncWrite; + +#[cfg(feature = "object_storage")] +use crate::models::streams::object_store::ObjectWrite; +#[cfg(feature = "object_storage")] +use crate::object_storage::{ObjectUpload, object_url, resolve}; + +/// Sequential writer to a local file or an object. +/// +/// Implements `Write` and `Seek`. For an object, `flush` completes the +/// object, and `Seek` reports the current position without moving it, +/// which is the form the Parquet writer uses to record offsets. +pub enum StorageWriter { + /// A local file. + File(File), + /// A streaming upload to an object. + #[cfg(feature = "object_storage")] + Object(ObjectUpload), +} + +impl StorageWriter { + /// Create the file or object at `path`, replacing any existing one. + /// + /// An object URL makes no request until the first part fills or the + /// writer is flushed. + pub fn create>(path: P) -> io::Result { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(path.as_ref()) { + return Ok(Self::Object(ObjectUpload::new(resolve(&url)?)?)); + } + Ok(Self::File(File::create(path)?)) + } +} + +impl Write for StorageWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + match self { + Self::File(file) => file.write(buf), + #[cfg(feature = "object_storage")] + Self::Object(upload) => { + upload.write(buf)?; + Ok(buf.len()) + } + } + } + + /// Flush a local file, or complete an object. + /// + /// Completing an object uploads its remaining bytes and makes it + /// visible in the store. Later flushes do nothing, and later writes + /// return an error. + fn flush(&mut self) -> io::Result<()> { + match self { + Self::File(file) => file.flush(), + #[cfg(feature = "object_storage")] + Self::Object(upload) => upload.complete().map(|_| ()), + } + } +} + +impl Seek for StorageWriter { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + match self { + Self::File(file) => file.seek(pos), + #[cfg(feature = "object_storage")] + Self::Object(upload) => { + let written = upload.written(); + let target = match pos { + SeekFrom::Start(offset) => offset as i128, + SeekFrom::End(delta) | SeekFrom::Current(delta) => { + written as i128 + delta as i128 + } + }; + match target == written as i128 && !upload.is_closed() { + true => Ok(written), + false => Err(io::Error::new( + io::ErrorKind::Unsupported, + "object storage writes are append-only", + )), + } + } + } + } +} + +/// Async writer to a local file or an object. +/// +/// Implements `AsyncWrite`. For an object, `poll_shutdown` completes the +/// object. +#[cfg(not(target_family = "wasm"))] +pub enum AsyncStorageWriter { + /// A local file. + File(tokio::fs::File), + /// A buffered upload to an object. + #[cfg(feature = "object_storage")] + Object(ObjectWrite), +} + +#[cfg(not(target_family = "wasm"))] +impl AsyncStorageWriter { + /// Create the file or object at `path`, replacing any existing one. + pub async fn create>(path: P) -> io::Result { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(path.as_ref()) { + return Ok(Self::Object(ObjectWrite::create(&url)?)); + } + Ok(Self::File(tokio::fs::File::create(path).await?)) + } +} + +#[cfg(not(target_family = "wasm"))] +impl AsyncWrite for AsyncStorageWriter { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + match self.get_mut() { + Self::File(file) => Pin::new(file).poll_write(cx, buf), + #[cfg(feature = "object_storage")] + Self::Object(object) => Pin::new(object).poll_write(cx, buf), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + Self::File(file) => Pin::new(file).poll_flush(cx), + #[cfg(feature = "object_storage")] + Self::Object(object) => Pin::new(object).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + Self::File(file) => Pin::new(file).poll_shutdown(cx), + #[cfg(feature = "object_storage")] + Self::Object(object) => Pin::new(object).poll_shutdown(cx), + } + } +} + +/// Create the directory `dir` and its parents. Does nothing for an object +/// prefix, since object storage has no directories. +pub fn create_dir_all>(dir: P) -> io::Result<()> { + #[cfg(feature = "object_storage")] + if object_url(dir.as_ref()).is_some() { + return Ok(()); + } + fs::create_dir_all(dir) +} diff --git a/rust/src/object_storage/connection.rs b/rust/src/object_storage/connection.rs new file mode 100644 index 0000000..fdf3e50 --- /dev/null +++ b/rust/src/object_storage/connection.rs @@ -0,0 +1,474 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object storage connections +//! +//! Resolves object URLs to a store and an object path. Stores are built +//! once per bucket, option set and Tokio runtime, and cached for the life of +//! the process. Readers and writers that address the same bucket from the +//! same runtime share one HTTP connection pool and one set of credentials. +//! +//! The runtime forms part of the cache key because an HTTP client's pooled +//! connections are driven by tasks on the runtime that opened them. A store +//! reused from another runtime fails its requests once the original runtime +//! has shut down. +//! +//! Options registered for a URL prefix apply to every URL beneath it. The +//! longest registered prefix that matches a URL supplies its options. +//! Registering options clears the store cache, which applies the new +//! options to later URLs. + +use std::collections::HashMap; +#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))] +use std::env; +use std::io; +use std::path::Path; +#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))] +use std::str::FromStr; +use std::sync::{Arc, LazyLock, Mutex, RwLock}; + +use object_store::memory::InMemory; +use object_store::path::Path as ObjectPath; +#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))] +use object_store::{ClientConfigKey, ClientOptions}; +use object_store::{ObjectStore, ObjectStoreScheme}; +use tokio::runtime::Handle; +use url::Url; + +use crate::object_storage::error::io_error; + +/// Default size of each multipart upload part, in bytes. S3 requires every +/// part except the last to be at least this size, and it is the minimum +/// accepted for `part_size`. +pub const DEFAULT_PART_SIZE: usize = 5 * 1024 * 1024; + +/// Default number of parts uploaded concurrently per object. +pub const DEFAULT_MAX_CONCURRENT_PARTS: usize = 8; + +/// Registered option key for the multipart upload part size, in bytes. +const PART_SIZE_KEY: &str = "part_size"; + +/// Registered option key for the number of parts uploaded concurrently per +/// object. +const MAX_CONCURRENT_PARTS_KEY: &str = "max_concurrent_parts"; + +/// URL schemes routed to object storage. Every other path is a local path. +const OBJECT_SCHEMES: &[&str] = &[ + "s3", "s3a", "gs", "az", "adl", "azure", "abfs", "abfss", "memory", +]; + +/// Options registered per URL prefix, as `(prefix, options)`. +static OPTIONS: LazyLock)>>> = + LazyLock::new(|| RwLock::new(Vec::new())); + +/// Stores built so far, keyed by scheme, bucket, matched option prefix and +/// Tokio runtime. +static STORES: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// The process-wide store behind `memory://` URLs. It is held outside the +/// store cache, which keeps its contents when options are registered. +static MEMORY: LazyLock> = LazyLock::new(|| Arc::new(InMemory::new())); + +/// One object or prefix within a store, resolved from an object URL, with +/// the upload settings registered for it. +#[derive(Debug, Clone)] +pub struct ObjectLocation { + pub(crate) store: Arc, + pub(crate) path: ObjectPath, + /// Size of each multipart upload part, in bytes. + pub(crate) part_size: usize, + /// Number of parts uploaded concurrently per object. + pub(crate) max_concurrent_parts: usize, +} + +/// Set the options used for every object URL beneath `prefix`. +/// +/// `prefix` is an object URL such as `s3://telemetry` or +/// `s3://telemetry/lightning`. The options replace any registered earlier +/// for the same prefix, and are of two kinds: +/// +/// - Upload settings, applied by lightstream to every backend: +/// - `part_size`: size of each multipart upload part in bytes. Defaults to +/// [`DEFAULT_PART_SIZE`], which is also the minimum. +/// - `max_concurrent_parts`: parts uploaded concurrently per object. +/// Defaults to [`DEFAULT_MAX_CONCURRENT_PARTS`]. Lower values suit links +/// with little upload bandwidth. +/// - Backend configuration keys, for example `aws_region`, +/// `aws_access_key_id`, `aws_endpoint` or `google_service_account`, and the +/// HTTP client keys such as `timeout`, `connect_timeout` and +/// `http1_only`. Options not given here are read by the backend from its +/// standard environment variables. +/// +/// Requests have no limit on their total duration unless `timeout` is set, +/// which lets large parts upload over slow links. Connections that cannot +/// be established fail after the client's `connect_timeout`. +/// +/// ## Errors +/// Returns `ErrorKind::InvalidInput` when `prefix` is not an object URL or an +/// upload setting is invalid. Unknown backend keys are reported when a URL +/// beneath the prefix is first opened. +pub fn register_options(prefix: &str, options: &[(&str, &str)]) -> io::Result<()> { + if object_url(Path::new(prefix)).is_none() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("'{prefix}' is not an object storage URL"), + )); + } + let prefix = prefix.trim_end_matches('/').to_string(); + let options: Vec<(String, String)> = options + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect(); + upload_settings(&options)?; + { + let mut registered = OPTIONS.write().expect("object storage options lock"); + registered.retain(|(existing, _)| *existing != prefix); + registered.push((prefix, options)); + } + STORES + .lock() + .expect("object storage store cache lock") + .clear(); + Ok(()) +} + +/// The object URL held by `path`, when its scheme is an object storage +/// scheme. +/// +/// On Windows, `Path::join` inserts backslashes into joined URLs, and these +/// are normalised to `/`. +pub(crate) fn object_url(path: &Path) -> Option { + let text = path.to_str()?; + let (scheme, _) = text.split_once("://")?; + if !OBJECT_SCHEMES.contains(&scheme.to_ascii_lowercase().as_str()) { + return None; + } + #[cfg(windows)] + let text = text.replace('\\', "/"); + Some(text.to_string()) +} + +/// Resolve an object URL to its store and object path. +/// +/// ## Errors +/// - `ErrorKind::InvalidInput` for a malformed URL or an unknown option key. +/// - `ErrorKind::Unsupported` for a scheme whose backend feature is not +/// enabled. +pub(crate) fn resolve(url: &str) -> io::Result { + let parsed = Url::parse(url).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("invalid object storage URL '{url}': {e}"), + ) + })?; + let (scheme, path) = ObjectStoreScheme::parse(&parsed).map_err(|e| io_error(e.into()))?; + let (prefix, options) = registered_options(url); + let (part_size, max_concurrent_parts) = upload_settings(&options)?; + if scheme == ObjectStoreScheme::Memory { + return Ok(ObjectLocation { + store: Arc::clone(&MEMORY), + path, + part_size, + max_concurrent_parts, + }); + } + + let runtime = Handle::try_current() + .map(|handle| handle.id().to_string()) + .unwrap_or_default(); + let key = format!( + "{}://{}|{prefix}|{runtime}", + parsed.scheme(), + parsed.host_str().unwrap_or("") + ); + + let mut stores = STORES.lock().expect("object storage store cache lock"); + let store = match stores.get(&key) { + Some(store) => Arc::clone(store), + None => { + let backend_options: Vec<(String, String)> = options + .into_iter() + .filter(|(key, _)| key != PART_SIZE_KEY && key != MAX_CONCURRENT_PARTS_KEY) + .collect(); + let store = build(scheme, &parsed, &backend_options)?; + stores.insert(key, Arc::clone(&store)); + store + } + }; + Ok(ObjectLocation { + store, + path, + part_size, + max_concurrent_parts, + }) +} + +/// The `(part_size, max_concurrent_parts)` upload settings in `options`, +/// with the defaults for any not given. +/// +/// ## Errors +/// Returns `ErrorKind::InvalidInput` for a value that is not a whole number, +/// a part size below [`DEFAULT_PART_SIZE`], or zero concurrent parts. +fn upload_settings(options: &[(String, String)]) -> io::Result<(usize, usize)> { + let mut part_size = DEFAULT_PART_SIZE; + let mut max_concurrent_parts = DEFAULT_MAX_CONCURRENT_PARTS; + for (key, value) in options { + let target = match key.as_str() { + PART_SIZE_KEY => &mut part_size, + MAX_CONCURRENT_PARTS_KEY => &mut max_concurrent_parts, + _ => continue, + }; + *target = value.parse().map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("object storage option '{key}' must be a whole number, got '{value}'"), + ) + })?; + } + if part_size < DEFAULT_PART_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "object storage option '{PART_SIZE_KEY}' must be at least {DEFAULT_PART_SIZE} bytes" + ), + )); + } + if max_concurrent_parts == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("object storage option '{MAX_CONCURRENT_PARTS_KEY}' must be at least 1"), + )); + } + Ok((part_size, max_concurrent_parts)) +} + +/// Options of the longest registered prefix that `url` falls under, with +/// that prefix. Returns an empty prefix and no options when none match. +fn registered_options(url: &str) -> (String, Vec<(String, String)>) { + let registered = OPTIONS.read().expect("object storage options lock"); + registered + .iter() + .filter(|(prefix, _)| { + url == prefix + || url + .strip_prefix(prefix.as_str()) + .is_some_and(|rest| rest.starts_with('/')) + }) + .max_by_key(|(prefix, _)| prefix.len()) + .map(|(prefix, options)| (prefix.clone(), options.clone())) + .unwrap_or_default() +} + +/// Build the store for `scheme`, applying `options` over the backend's +/// environment configuration. +/// +/// The HTTP client starts with no limit on total request time. The +/// backend's environment variables and then `options` are applied over that +/// default, and either can set a `timeout`. +#[cfg_attr( + not(any(feature = "aws", feature = "gcp", feature = "azure")), + allow(unused_variables) +)] +fn build( + scheme: ObjectStoreScheme, + url: &Url, + options: &[(String, String)], +) -> io::Result> { + match scheme { + #[cfg(feature = "aws")] + ObjectStoreScheme::AmazonS3 => { + use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; + let mut builder = AmazonS3Builder::from_env() + .with_url(url.as_str()) + .with_client_options(client_options::( + "AWS_", + |key| match key { + AmazonS3ConfigKey::Client(key) => Some(key), + _ => None, + }, + )); + for (key, value) in options { + let key: AmazonS3ConfigKey = key.parse().map_err(io_error)?; + builder = builder.with_config(key, value); + } + let store = builder.build().map_err(io_error)?; + Ok(Arc::new(store)) + } + #[cfg(feature = "gcp")] + ObjectStoreScheme::GoogleCloudStorage => { + use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey}; + let mut builder = GoogleCloudStorageBuilder::from_env() + .with_url(url.as_str()) + .with_client_options(client_options::( + "GOOGLE_", + |key| match key { + GoogleConfigKey::Client(key) => Some(key), + _ => None, + }, + )); + for (key, value) in options { + let key: GoogleConfigKey = key.parse().map_err(io_error)?; + builder = builder.with_config(key, value); + } + let store = builder.build().map_err(io_error)?; + Ok(Arc::new(store)) + } + #[cfg(feature = "azure")] + ObjectStoreScheme::MicrosoftAzure => { + use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder}; + let mut builder = MicrosoftAzureBuilder::from_env() + .with_url(url.as_str()) + .with_client_options(client_options::( + "AZURE_", + |key| match key { + AzureConfigKey::Client(key) => Some(key), + _ => None, + }, + )); + for (key, value) in options { + let key: AzureConfigKey = key.parse().map_err(io_error)?; + builder = builder.with_config(key, value); + } + let store = builder.build().map_err(io_error)?; + Ok(Arc::new(store)) + } + _ => Err(io::Error::new( + io::ErrorKind::Unsupported, + format!( + "object storage URL '{url}' requires the lightstream `aws`, `gcp` or `azure` feature for its scheme" + ), + )), + } +} + +/// HTTP client options with no limit on total request time, followed by any +/// client settings in the environment variables beginning with `prefix`. +/// +/// The backend builders' `from_env` reads the same variables. Setting the +/// client options afterwards replaces those, so the client settings among +/// them are applied again here. +#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))] +fn client_options( + prefix: &str, + client_key: impl Fn(K) -> Option, +) -> ClientOptions { + let mut options = ClientOptions::new().with_timeout_disabled(); + for (key, value) in env::vars() { + if !key.starts_with(prefix) { + continue; + } + if let Some(key) = key.to_ascii_lowercase().parse().ok().and_then(&client_key) { + options = options.with_config(key, value); + } + } + options +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn upload_settings_default_and_validate() { + assert_eq!( + upload_settings(&[]).unwrap(), + (DEFAULT_PART_SIZE, DEFAULT_MAX_CONCURRENT_PARTS) + ); + let custom = [ + (PART_SIZE_KEY.to_string(), (8 * 1024 * 1024).to_string()), + (MAX_CONCURRENT_PARTS_KEY.to_string(), "2".to_string()), + ]; + assert_eq!(upload_settings(&custom).unwrap(), (8 * 1024 * 1024, 2)); + + for (key, value) in [ + (PART_SIZE_KEY, "1024"), + (PART_SIZE_KEY, "large"), + (MAX_CONCURRENT_PARTS_KEY, "0"), + ] { + let err = upload_settings(&[(key.to_string(), value.to_string())]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput, "{key}={value}"); + } + } + + #[test] + fn registering_an_invalid_upload_setting_is_rejected() { + let err = register_options( + "memory:///invalid-settings", + &[("max_concurrent_parts", "0")], + ) + .unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + } + + #[test] + fn registered_upload_settings_reach_the_location() { + register_options( + "memory:///upload-settings", + &[("max_concurrent_parts", "2"), ("part_size", "6291456")], + ) + .unwrap(); + let location = resolve("memory:///upload-settings/a.bin").unwrap(); + assert_eq!(location.max_concurrent_parts, 2); + assert_eq!(location.part_size, 6 * 1024 * 1024); + } + + #[test] + fn object_url_detects_object_schemes_only() { + assert!(object_url(Path::new("s3://bucket/key.parquet")).is_some()); + assert!(object_url(Path::new("memory:///data/a.csv")).is_some()); + assert!(object_url(Path::new("/tmp/data/a.csv")).is_none()); + assert!(object_url(Path::new("relative/a.csv")).is_none()); + assert!(object_url(Path::new("C:\\data\\a.csv")).is_none()); + assert!(object_url(Path::new("https://host/a.csv")).is_none()); + } + + #[test] + fn object_urls_joined_as_paths_keep_their_form() { + let joined = Path::new("memory:///chunks").join("trades-0000000001.parquet"); + assert_eq!( + object_url(&joined).as_deref(), + Some("memory:///chunks/trades-0000000001.parquet") + ); + } + + #[test] + fn the_memory_store_is_shared_across_resolutions() { + let first = resolve("memory:///shared/a").unwrap(); + let second = resolve("memory:///shared/b").unwrap(); + assert!(Arc::ptr_eq(&first.store, &second.store)); + assert_eq!(first.path.as_ref(), "shared/a"); + } + + #[test] + fn the_longest_registered_prefix_supplies_the_options() { + register_options("memory:///prefix-test", &[("a", "1")]).unwrap(); + register_options("memory:///prefix-test/inner", &[("b", "2")]).unwrap(); + + let (prefix, options) = registered_options("memory:///prefix-test/inner/x"); + assert_eq!(prefix, "memory:///prefix-test/inner"); + assert_eq!(options, vec![("b".to_string(), "2".to_string())]); + + let (prefix, _) = registered_options("memory:///prefix-test/other"); + assert_eq!(prefix, "memory:///prefix-test"); + + let (prefix, _) = registered_options("memory:///prefix-testing"); + assert_eq!(prefix, ""); + } + + #[test] + fn registering_a_local_path_is_rejected() { + let err = register_options("/tmp/data", &[]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + } + + #[cfg(not(feature = "aws"))] + #[test] + fn a_backend_without_its_feature_is_unsupported() { + let err = resolve("s3://bucket/key").unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::Unsupported); + } +} diff --git a/rust/src/object_storage/error.rs b/rust/src/object_storage/error.rs new file mode 100644 index 0000000..cd02ad5 --- /dev/null +++ b/rust/src/object_storage/error.rs @@ -0,0 +1,33 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object storage errors +//! +//! The file readers and writers return `io::Result`, and object storage +//! failures are reported through the same type. The `object_store` error +//! remains available as the source of the `io::Error`, and its category is +//! mapped onto the closest `io::ErrorKind`, keeping a missing object and a +//! missing file distinguishable from other failures in the same way. + +use std::io; + +/// Convert an `object_store` error into an `io::Error` with the matching +/// kind. +pub(crate) fn io_error(error: object_store::Error) -> io::Error { + let kind = match &error { + object_store::Error::NotFound { .. } => io::ErrorKind::NotFound, + object_store::Error::AlreadyExists { .. } => io::ErrorKind::AlreadyExists, + object_store::Error::PermissionDenied { .. } + | object_store::Error::Unauthenticated { .. } => io::ErrorKind::PermissionDenied, + object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. } => { + io::ErrorKind::Unsupported + } + object_store::Error::InvalidPath { .. } + | object_store::Error::UnknownConfigurationKey { .. } => io::ErrorKind::InvalidInput, + _ => io::ErrorKind::Other, + }; + io::Error::new(kind, error) +} diff --git a/rust/src/object_storage/mod.rs b/rust/src/object_storage/mod.rs new file mode 100644 index 0000000..80ce5b7 --- /dev/null +++ b/rust/src/object_storage/mod.rs @@ -0,0 +1,83 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object storage +//! +//! Remote object storage for the file readers and writers. With the +//! `object_storage` feature enabled, every reader and writer that takes a +//! file path also accepts an object URL in its place: +//! +//! ```rust,ignore +//! let mut writer = CsvWriter::to_path("s3://telemetry/trades/2026-10-14.csv")?; +//! writer.write_table(&table)?; +//! writer.flush()?; +//! +//! let reader = FileTableReader::open("gs://archive/trades.arrow")?; +//! ``` +//! +//! ## Structure +//! - `connection.rs`: URL detection, store resolution and caching, and the +//! per-prefix option registry. +//! - `runtime.rs`: access to the caller's Tokio runtime for object storage +//! requests. +//! - `upload.rs`: the streaming upload behind `StorageWriter`. +//! - `error.rs`: conversion of store errors into `io::Error`. +//! +//! ## Backends +//! +//! | URL | Backend | Feature | +//! |---|---|---| +//! | `memory:///path` | In-process memory, shared across the process | `object_storage` | +//! | `s3://bucket/path`, `s3a://` | Amazon S3 and S3-compatible stores | `aws` | +//! | `gs://bucket/path` | Google Cloud Storage | `gcp` | +//! | `az://container/path`, `abfs://`, `abfss://`, `adl://`, `azure://` | Azure Blob Storage | `azure` | +//! +//! Paths without one of these schemes, including Windows drive paths, are +//! local file paths as before. +//! +//! ## Credentials +//! Each backend reads its standard environment variables (`AWS_*`, +//! `GOOGLE_*`, `AZURE_*`), falling back to instance or workload identity +//! where the platform provides one. Options for a bucket or prefix are set +//! with [`register_options`], and apply to every URL beneath that prefix. +//! +//! ## Uploads and timeouts +//! Objects larger than one part upload as multipart uploads, with +//! `part_size` bytes per part and up to `max_concurrent_parts` parts in +//! flight. Both are registered options, defaulting to +//! [`DEFAULT_PART_SIZE`] and [`DEFAULT_MAX_CONCURRENT_PARTS`]. Requests have +//! no limit on their total duration unless a `timeout` option is set, and +//! large parts therefore complete over slow links. Connection attempts are +//! still bounded by the client's `connect_timeout`. +//! +//! ## Runtime +//! Object storage I/O runs on the caller's Tokio runtime. The async adapters +//! in `models::streams::object_store` await their requests. The synchronous +//! readers and writers wait on their requests with `Handle::block_on`, and +//! are called with object URLs from outside the async executor: from a +//! `spawn_blocking` task, or from a thread that has entered the runtime with +//! `Handle::enter`. The parallel chunked reads and writes enter the caller's +//! runtime on each of their worker threads. +//! +//! ## HTTP versions +//! Requests use the `object_store` HTTP client, which is separate from the +//! HTTP/2 transport behind the `http` feature. The protocol is negotiated +//! per backend: S3 and Azure Blob serve HTTP/1.1, and GCS supports HTTP/2. +//! The `http1_only` and `http2_only` options force one protocol. + +mod connection; +mod error; +mod runtime; +mod upload; + +pub use connection::{ + DEFAULT_MAX_CONCURRENT_PARTS, DEFAULT_PART_SIZE, ObjectLocation, register_options, +}; +pub use upload::ObjectUpload; + +pub(crate) use connection::{object_url, resolve}; +pub(crate) use error::io_error; +pub(crate) use runtime::current_handle; diff --git a/rust/src/object_storage/runtime.rs b/rust/src/object_storage/runtime.rs new file mode 100644 index 0000000..8655d42 --- /dev/null +++ b/rust/src/object_storage/runtime.rs @@ -0,0 +1,37 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object storage on the caller's runtime +//! +//! Object storage requests run on the caller's Tokio runtime. The +//! synchronous readers and writers wait for their requests with +//! `Handle::block_on`, and multipart uploads spawn their part uploads on the +//! same handle. +//! +//! `Handle::block_on` blocks the calling thread, and it cannot be called +//! from inside an async task. The synchronous readers and writers are +//! therefore called with object URLs from outside the async executor: from a +//! `spawn_blocking` task, or from a thread that has entered the runtime with +//! `Handle::enter`. Called from inside an async task, they panic with Tokio's +//! "Cannot start a runtime from within a runtime" message. + +use std::io; + +use tokio::runtime::Handle; + +/// Handle to the caller's Tokio runtime. +/// +/// ## Errors +/// Returns `ErrorKind::Other` when the calling thread is outside a Tokio +/// runtime context. +pub(crate) fn current_handle() -> io::Result { + Handle::try_current().map_err(|_| { + io::Error::new( + io::ErrorKind::Other, + "object storage URLs require a Tokio runtime context on the calling thread", + ) + }) +} diff --git a/rust/src/object_storage/upload.rs b/rust/src/object_storage/upload.rs new file mode 100644 index 0000000..574b8ba --- /dev/null +++ b/rust/src/object_storage/upload.rs @@ -0,0 +1,215 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Object uploads +//! +//! Streaming upload of one object from synchronous writes. Written bytes +//! are buffered up to one part, of the size registered for the object's URL +//! (see [`crate::object_storage::register_options`]). An object that never fills a part is +//! uploaded with a single request when it completes. Larger objects start a +//! multipart upload when the first part fills, and each part uploads on the +//! caller's Tokio runtime while writing continues, which bounds memory to +//! the parts in flight regardless of the object size. +//! +//! An object appears in the store only when [`ObjectUpload::complete`] +//! succeeds. A failed or abandoned upload is aborted, leaving no object. + +use std::collections::VecDeque; +use std::io; +use std::mem; + +use object_store::{MultipartUpload, ObjectStoreExt, PutPayload, PutResult}; +use tokio::runtime::Handle; +use tokio::task::JoinHandle; + +use crate::object_storage::connection::ObjectLocation; +use crate::object_storage::error::io_error; +use crate::object_storage::runtime::current_handle; + +/// Upload of one object, driven by synchronous writes. +/// +/// Held by `StorageWriter` for object URLs. See the module documentation +/// for the upload behaviour. +pub struct ObjectUpload { + location: ObjectLocation, + /// The caller's runtime, captured when the upload was created. + runtime: Handle, + /// Bytes not yet handed to a part upload. + pending: Vec, + /// Multipart upload, started when the first part fills. + multipart: Option>, + /// Part uploads in flight, oldest first. + parts: VecDeque>>, + /// Bytes written to the object so far. + written: u64, + /// Set once the object has been completed or the upload aborted. + closed: bool, +} + +impl ObjectUpload { + /// Start an upload to `location` on the caller's runtime. + /// + /// ## Errors + /// Returns an error when the calling thread is outside a Tokio runtime + /// context. + pub(crate) fn new(location: ObjectLocation) -> io::Result { + Ok(Self { + location, + runtime: current_handle()?, + pending: Vec::new(), + multipart: None, + parts: VecDeque::new(), + written: 0, + closed: false, + }) + } + + /// Bytes written to the object so far. + pub(crate) fn written(&self) -> u64 { + self.written + } + + /// Whether the object has been completed or the upload aborted. + pub(crate) fn is_closed(&self) -> bool { + self.closed + } + + /// Append `buf` to the object, uploading every part that fills. + pub(crate) fn write(&mut self, buf: &[u8]) -> io::Result<()> { + if self.closed { + return Err(io::Error::new( + io::ErrorKind::Other, + format!( + "object '{}' is already complete. Object storage writes complete on the first flush", + self.location.path + ), + )); + } + self.pending.extend_from_slice(buf); + self.written += buf.len() as u64; + while self.pending.len() >= self.location.part_size { + let rest = self.pending.split_off(self.location.part_size); + let part = mem::replace(&mut self.pending, rest); + if let Err(error) = self.send_part(part) { + self.abort(); + return Err(error); + } + } + Ok(()) + } + + /// Complete the object, uploading any remaining bytes. + /// + /// Completing an object that is already complete does nothing and + /// returns `None`. + pub(crate) fn complete(&mut self) -> io::Result> { + if self.closed { + return Ok(None); + } + let result = match self.multipart.is_some() { + false => { + let payload = PutPayload::from(mem::take(&mut self.pending)); + self.runtime + .block_on(self.location.store.put(&self.location.path, payload)) + .map_err(io_error) + } + true => self.complete_multipart(), + }; + match result { + Ok(result) => { + self.closed = true; + Ok(Some(result)) + } + Err(error) => { + self.abort(); + Err(error) + } + } + } + + /// Upload the remaining bytes as the last part, wait for every part, + /// then complete the multipart upload. + fn complete_multipart(&mut self) -> io::Result { + if !self.pending.is_empty() { + let last = mem::take(&mut self.pending); + self.send_part(last)?; + } + while let Some(part) = self.parts.pop_front() { + await_part(&self.runtime, part)?; + } + let mut upload = self.multipart.take().expect("multipart upload started"); + self.runtime + .block_on(async move { + let result = upload.complete().await; + // A failed completion leaves the uploaded parts in the store + // until they are aborted. + if result.is_err() { + let _ = upload.abort().await; + } + result + }) + .map_err(io_error) + } + + /// Hand `part` to the multipart upload, starting the upload first when + /// this is the first part. + fn send_part(&mut self, part: Vec) -> io::Result<()> { + if self.multipart.is_none() { + let upload = self + .runtime + .block_on(self.location.store.put_multipart(&self.location.path)) + .map_err(io_error)?; + self.multipart = Some(upload); + } + // Wait for the oldest part once the registered number of parts is in + // flight, which bounds memory and shares the upload bandwidth across + // fewer requests. + if self.parts.len() >= self.location.max_concurrent_parts { + let oldest = self.parts.pop_front().expect("parts in flight"); + await_part(&self.runtime, oldest)?; + } + let upload = self + .multipart + .as_mut() + .expect("multipart upload started") + .put_part(PutPayload::from(part)); + self.parts.push_back(self.runtime.spawn(upload)); + Ok(()) + } + + /// Abort the upload, discarding any uploaded parts. + /// + /// The abort runs on the caller's runtime without blocking the caller, + /// which is already reporting the failure that caused it. + fn abort(&mut self) { + self.closed = true; + self.pending = Vec::new(); + for part in self.parts.drain(..) { + part.abort(); + } + if let Some(mut upload) = self.multipart.take() { + self.runtime.spawn(async move { + let _ = upload.abort().await; + }); + } + } +} + +impl Drop for ObjectUpload { + fn drop(&mut self) { + if !self.closed { + self.abort(); + } + } +} + +/// Wait for one part upload to finish. +fn await_part(runtime: &Handle, part: JoinHandle>) -> io::Result<()> { + runtime + .block_on(part) + .map_err(|e| io::Error::new(io::ErrorKind::Other, e))? + .map_err(io_error) +} diff --git a/rust/src/traits/chunked_table_reader.rs b/rust/src/traits/chunked_table_reader.rs index c19d756..db52a98 100644 --- a/rust/src/traits/chunked_table_reader.rs +++ b/rust/src/traits/chunked_table_reader.rs @@ -14,6 +14,11 @@ //! and a `load_table` that further consolidates the chunks into a //! single `Table`. The `*_cols` variants project to the named columns. //! +//! Chunks in nested partition directories, as written by a partitioned +//! chunked writer, are included and ordered by partition directory and then +//! by index. Opening a reader on one partition directory, such as +//! `/date=2026-10-14`, reads that partition alone. +//! //! `load_batched` deliberately does not consolidate; the per-chunk batches //! are preserved in the returned `SuperTable`. Callers wanting a single //! `Table` use `load_table` (which itself calls @@ -31,11 +36,13 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use minarrow::{Consolidate, SuperTable, Table}; +use tokio::runtime::Handle; /// Open / load_batched / iterate / par_load_batched shape shared by the /// per-format chunked readers. /// -/// Implementors yield one `Table` per chunk file in ascending index order. +/// Implementors yield one `Table` per chunk file, ordered by partition +/// directory and then by ascending index. /// Per-format extras (CSV decode options, batch size, etc.) flow through /// the associated `Options` type; formats with no extras use `()`. pub trait ChunkedTableReader: Iterator> + Sized { @@ -44,18 +51,19 @@ pub trait ChunkedTableReader: Iterator> + Size /// Per-format constructor options. Use `()` if no extras are needed. type Options; - /// Open every `-*.` file inside `dir`, sorted by numeric - /// suffix so chunks come out in write order. Files whose names don't - /// match the pattern are ignored. + /// Open every `-*.` file beneath `dir`, including files in + /// nested partition directories, sorted by partition directory and then + /// by numeric suffix. Chunks of an unpartitioned directory come out in + /// write order. Files whose names don't match the pattern are ignored. fn open>( dir: P, base: &str, options: Self::Options, ) -> Result; - /// List chunk file paths in `dir` matching this format's naming - /// pattern, sorted by ascending numeric index. Used by `open` to - /// populate the reader's path list. + /// List chunk file paths beneath `dir` matching this format's naming + /// pattern, sorted by partition directory and then by ascending numeric + /// index. Used by `open` to populate the reader's path list. fn list_paths>(dir: P, base: &str) -> Result, Self::Error>; /// Return the chunk paths held by this reader, in write order. @@ -169,13 +177,18 @@ pub trait ChunkedTableReader: Iterator> + Size // perfectly balanced work and a one-slot overshoot worst-case. // Avoids the Vec doubling growth as each worker pushes results. let per_worker_cap = n_files.div_ceil(n_threads); + // Worker threads enter the caller's Tokio runtime, when there is + // one, which lets object storage paths make their requests on it. + let runtime = Handle::try_current().ok(); let collected: Result, Self::Error> = std::thread::scope(|s| { let mut handles = Vec::with_capacity(n_threads); for _ in 0..n_threads { let next_index = &next_index; let this = &reader; + let runtime = &runtime; handles.push( s.spawn(move || -> Result, Self::Error> { + let _runtime = runtime.as_ref().map(Handle::enter); let mut local: Vec<(usize, Table)> = Vec::with_capacity(per_worker_cap); let paths = this.paths(); loop { @@ -233,13 +246,18 @@ pub trait ChunkedTableReader: Iterator> + Size let next_index = AtomicUsize::new(0); let per_worker_cap = n_files.div_ceil(n_threads); + // Worker threads enter the caller's Tokio runtime, when there is + // one, which lets object storage paths make their requests on it. + let runtime = Handle::try_current().ok(); let collected: Result, Self::Error> = std::thread::scope(|s| { let mut handles = Vec::with_capacity(n_threads); for _ in 0..n_threads { let next_index = &next_index; let this = &reader; + let runtime = &runtime; handles.push( s.spawn(move || -> Result, Self::Error> { + let _runtime = runtime.as_ref().map(Handle::enter); let mut local: Vec<(usize, Table)> = Vec::with_capacity(per_worker_cap); let paths = this.paths(); loop { diff --git a/rust/src/traits/chunked_table_writer.rs b/rust/src/traits/chunked_table_writer.rs index b6718b8..5bc802b 100644 --- a/rust/src/traits/chunked_table_writer.rs +++ b/rust/src/traits/chunked_table_writer.rs @@ -25,12 +25,29 @@ //! write each batch end-to-end via the writer's `write_chunk_at`. Index //! reservation goes through the same counter as `write_chunk`, so the //! two paths compose without producing duplicate filenames. +//! +//! ## Partitioning +//! +//! A writer configured with a [`PartitionType`] places each chunk in a +//! partition directory beneath its root, for example +//! `/date=2026-10-14/hour=13/-NNNNNNNNNN.`. The partition +//! path comes from the chunk's time, given through +//! [`ChunkedTableWriter::write_chunk_for_time`]. `write_chunk`, `write_all` +//! and `par_write_all` use the current system time, and the two `_all` +//! methods use one time for every batch in the call. Partition directories +//! are created as they are first written to. Chunk indices come from the +//! writer's single counter whatever the partition. use std::io; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; use minarrow::{SuperTable, Table}; +use tokio::runtime::Handle; + +use crate::models::writers::chunked::partition::PartitionType; +use crate::models::writers::object_storage::create_dir_all; /// extension / dir / base / write_chunk / write_chunk_at / par_write_all /// shape shared by the per-format chunked writers. @@ -63,6 +80,10 @@ pub trait ChunkedTableWriter: Sized { /// reserve indices off this atom so the two paths compose. fn counter(&self) -> &AtomicU64; + /// Partitioning applied to chunk paths, or `None` for a single flat + /// directory. + fn partition(&self) -> Option<&PartitionType>; + /// Total number of chunks written so far. fn batches_written(&self) -> u64 { self.counter().load(Ordering::Relaxed) @@ -80,6 +101,29 @@ pub trait ChunkedTableWriter: Sized { )) } + /// Build the path for chunk `index` of `table` at `time_ns` epoch + /// nanoseconds, creating its partition directory if needed. + /// + /// Without partitioning this is [`Self::chunk_path_for`]. + fn chunk_path_for_time( + &self, + table: &Table, + time_ns: u64, + index: u64, + ) -> io::Result { + let Some(partition) = self.partition() else { + return Ok(self.chunk_path_for(index)); + }; + let dir = self.dir().join(partition.partition(table, time_ns)); + create_dir_all(&dir)?; + Ok(dir.join(format!( + "{}-{:010}.{}", + self.base(), + index, + Self::extension() + ))) + } + /// Encode `table` into a complete chunk file at `path`, using any /// per-format options stored on the writer. Shared per-format encode /// path used both by `write_chunk` (one chunk at a time) and by @@ -88,9 +132,22 @@ pub trait ChunkedTableWriter: Sized { /// Reserve the next chunk index off the shared counter and write /// `table` to that chunk file. Returns the path of the written file. + /// + /// A partitioned writer takes the partition from the current system + /// time. See [`Self::write_chunk_for_time`] to supply the time. fn write_chunk(&mut self, table: &Table) -> Result { + self.write_chunk_for_time(table, current_time_ns()) + } + + /// Write `table` as the next chunk, with `time_ns` epoch nanoseconds as + /// its time for partitioning. Returns the path of the written file. + /// + /// The time is typically the start of the event-time window that + /// produced `table`, which places replayed or late data in the same + /// partition as the original. Writers without partitioning ignore it. + fn write_chunk_for_time(&mut self, table: &Table, time_ns: u64) -> Result { let idx = self.counter().fetch_add(1, Ordering::Relaxed); - let path = self.chunk_path_for(idx); + let path = self.chunk_path_for_time(table, time_ns, idx)?; self.write_chunk_at(&path, table)?; Ok(path) } @@ -99,9 +156,10 @@ pub trait ChunkedTableWriter: Sized { /// file, in order. Returns the paths of the written files. Advances /// the shared counter by `supertable.batches.len()`. fn write_all(&mut self, supertable: &SuperTable) -> Result, Self::Error> { + let time_ns = current_time_ns(); let mut paths = Vec::with_capacity(supertable.batches.len()); for batch in supertable.batches.iter() { - paths.push(self.write_chunk(batch.as_ref())?); + paths.push(self.write_chunk_for_time(batch.as_ref(), time_ns)?); } Ok(paths) } @@ -150,6 +208,7 @@ pub trait ChunkedTableWriter: Sized { // any concurrent `write_chunk` callers or with a later resumed // streaming write. let start = self.counter().fetch_add(n_files as u64, Ordering::Relaxed); + let time_ns = current_time_ns(); let cursor = AtomicU64::new(0); // Per-worker capacity hint: ceil(n_files / n_threads). Workers @@ -157,14 +216,19 @@ pub trait ChunkedTableWriter: Sized { // perfectly balanced work and a one-slot overshoot worst-case. // Avoids the Vec doubling growth as each worker pushes results. let per_worker_cap = n_files.div_ceil(n_threads); + // Worker threads enter the caller's Tokio runtime, when there is + // one, which lets object storage paths make their requests on it. + let runtime = Handle::try_current().ok(); let collected: Result, Self::Error> = std::thread::scope(|s| { let mut handles = Vec::with_capacity(n_threads); for _ in 0..n_threads { let cursor = &cursor; let tables = &tables; let me = &*self; + let runtime = &runtime; handles.push( s.spawn(move || -> Result, Self::Error> { + let _runtime = runtime.as_ref().map(Handle::enter); let mut local: Vec<(u64, PathBuf)> = Vec::with_capacity(per_worker_cap); loop { let i = cursor.fetch_add(1, Ordering::Relaxed); @@ -172,7 +236,7 @@ pub trait ChunkedTableWriter: Sized { break; } let idx = start + i; - let path = me.chunk_path_for(idx); + let path = me.chunk_path_for_time(tables[i as usize], time_ns, idx)?; me.write_chunk_at(&path, tables[i as usize])?; local.push((i, path)); } @@ -194,3 +258,10 @@ pub trait ChunkedTableWriter: Sized { Ok(all.into_iter().map(|(_, p)| p).collect()) } } + +/// Current system time in epoch nanoseconds. +fn current_time_ns() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| elapsed.as_nanos() as u64) +} diff --git a/rust/tests/object_storage_checks/mod.rs b/rust/tests/object_storage_checks/mod.rs new file mode 100644 index 0000000..6a58dc0 --- /dev/null +++ b/rust/tests/object_storage_checks/mod.rs @@ -0,0 +1,385 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! Object storage checks shared by the per-backend test files. +//! +//! Each check writes beneath its own path under a base object URL and reads +//! the result back through the standard file readers and writers, which +//! lets the same checks run against the in-memory store and each cloud +//! backend. +//! +//! The synchronous checks run on a thread that has entered a Tokio runtime +//! outside the async executor. The async checks run inside a +//! multi-threaded runtime and call the synchronous readers through +//! `spawn_blocking`. + +#![allow(dead_code)] + +use std::io::{ErrorKind, Read, Write}; +use std::path::Path; + +use lightstream::enums::{BufferChunkSize, IPCMessageProtocol}; +use lightstream::models::readers::ipc::file_table::FileTableReader; +use lightstream::models::readers::ipc::table::TableReader; +use lightstream::models::readers::object_storage::{StorageReader, list_dir}; +use lightstream::models::streams::object_store::{ObjectByteStream, ObjectWrite}; +use lightstream::models::writers::ipc::sync_table::write_tables_to_file_sync; +use lightstream::models::writers::ipc::table::{TableWriter, write_tables_to_file}; +use lightstream::models::writers::object_storage::StorageWriter; +use minarrow::{Array, Field, FieldArray, IntegerArray, NumericArray, StringArray, Table}; +use tokio::runtime::Runtime; +use tokio::task::spawn_blocking; + +/// Multi-threaded runtime for the synchronous checks. Each test enters it on +/// the test thread, which is outside the async executor. +pub fn runtime() -> Runtime { + Runtime::new().expect("test runtime starts") +} + +/// Table of `n` rows whose `id` column starts at `first`. +pub fn sample_table(first: i64, n: usize) -> Table { + let ids: Vec = (first..first + n as i64).collect(); + let venues: Vec<&str> = ids + .iter() + .map(|id| if id % 2 == 0 { "XNAS" } else { "XNYS" }) + .collect(); + Table::new( + "trades".to_string(), + vec![ + FieldArray::from_arr("id", Array::from_int64(IntegerArray::from_slice(&ids))), + FieldArray::from_arr( + "venue", + Array::from_string32(StringArray::from_slice(&venues)), + ), + ] + .into(), + ) +} + +pub fn schema_of(table: &Table) -> Vec { + table.cols.iter().map(|col| (*col.field).clone()).collect() +} + +/// Concatenated `id` column values of `tables`. +pub fn ids(tables: &[Table]) -> Vec { + tables + .iter() + .flat_map(|table| match &table.cols[0].array { + Array::NumericArray(NumericArray::Int64(ids)) => ids.data.as_slice().to_vec(), + // CSV type inference narrows small integers to Int32. + Array::NumericArray(NumericArray::Int32(ids)) => { + ids.data.as_slice().iter().map(|id| *id as i64).collect() + } + other => panic!("id column decoded as {other:?}"), + }) + .collect() +} + +pub fn ipc_file_reader_reads_an_object_written_by_the_sync_writer(base: &str) { + let url = format!("{base}/ipc-sync/trades.arrow"); + let first = sample_table(0, 10); + let schema = schema_of(&first); + write_tables_to_file_sync(Path::new(&url), &[first, sample_table(10, 5)], schema).unwrap(); + + let reader = FileTableReader::open(&url).unwrap(); + assert_eq!(reader.num_batches(), 2); + let tables: Vec
= (0..2).map(|i| reader.read_batch(i).unwrap()).collect(); + assert_eq!(ids(&tables), (0..15).collect::>()); +} + +pub async fn ipc_file_reader_reads_an_object_written_by_the_async_writer(base: &str) { + let url = format!("{base}/ipc-async/trades.arrow"); + let first = sample_table(0, 8); + let schema = schema_of(&first); + write_tables_to_file(&url, &[first, sample_table(8, 8)], schema) + .await + .unwrap(); + + // The file reader is synchronous and runs outside the async executor. + let tables = spawn_blocking(move || { + let reader = FileTableReader::open(&url).unwrap(); + (0..reader.num_batches()) + .map(|i| reader.read_batch(i).unwrap()) + .collect::>() + }) + .await + .unwrap(); + assert_eq!(ids(&tables), (0..16).collect::>()); +} + +pub async fn ipc_stream_round_trips_through_the_async_adapters(base: &str) { + let url = format!("{base}/ipc-stream/trades.arrows"); + let first = sample_table(0, 20); + let schema = schema_of(&first); + let mut writer = TableWriter::new( + ObjectWrite::create(&url).unwrap(), + schema, + IPCMessageProtocol::Stream, + None, + ) + .unwrap(); + writer + .write_all_tables(vec![first, sample_table(20, 20)]) + .await + .unwrap(); + + let stream = ObjectByteStream::open(&url, BufferChunkSize::Http) + .await + .unwrap(); + let reader: TableReader = TableReader::new(stream, 64 * 1024, IPCMessageProtocol::Stream, None); + let tables = reader.read_all_tables().await.unwrap(); + assert_eq!(ids(&tables), (0..40).collect::>()); +} + +pub fn large_objects_upload_in_parts_and_read_back_whole(base: &str) { + let url = format!("{base}/large/blob.bin"); + // Three and a half parts, written in uneven slices. + let data: Vec = (0..(17 * 1024 * 1024 + 512 * 1024)) + .map(|i| (i % 251) as u8) + .collect(); + let mut writer = StorageWriter::create(&url).unwrap(); + for slice in data.chunks(1_000_003) { + writer.write_all(slice).unwrap(); + } + writer.flush().unwrap(); + + let mut read = Vec::new(); + StorageReader::open(&url) + .unwrap() + .read_to_end(&mut read) + .unwrap(); + assert_eq!(read.len(), data.len()); + assert!(read == data, "object contents match what was written"); +} + +pub fn writes_after_an_object_completes_are_rejected(base: &str) { + let url = format!("{base}/complete/once.bin"); + let mut writer = StorageWriter::create(&url).unwrap(); + writer.write_all(b"first").unwrap(); + writer.flush().unwrap(); + // A second flush with nothing new written does nothing. + writer.flush().unwrap(); + assert!(writer.write_all(b"second").is_err()); +} + +pub fn a_writer_dropped_before_flushing_leaves_no_object(base: &str) { + let url = format!("{base}/dropped/never.bin"); + let mut writer = StorageWriter::create(&url).unwrap(); + writer.write_all(b"unfinished").unwrap(); + drop(writer); + + let err = StorageReader::open(&url) + .err() + .expect("no object was created"); + assert_eq!(err.kind(), ErrorKind::NotFound); +} + +pub fn a_missing_object_reports_not_found(base: &str) { + let err = FileTableReader::open(format!("{base}/missing/none.arrow")) + .err() + .expect("missing object"); + assert_eq!(err.kind(), ErrorKind::NotFound); +} + +pub fn listing_returns_object_urls_immediately_beneath_the_prefix(base: &str) { + for name in ["a.bin", "b.bin", "nested/c.bin"] { + let mut writer = StorageWriter::create(format!("{base}/listing/{name}")).unwrap(); + writer.write_all(name.as_bytes()).unwrap(); + writer.flush().unwrap(); + } + let mut listed = list_dir(format!("{base}/listing")).unwrap(); + listed.sort(); + assert_eq!( + listed, + vec![ + Path::new(&format!("{base}/listing/a.bin")).to_path_buf(), + Path::new(&format!("{base}/listing/b.bin")).to_path_buf(), + ] + ); +} + +#[cfg(feature = "parquet")] +pub mod parquet { + use super::*; + use lightstream::models::readers::chunked::parquet::ChunkedParquetReader; + use lightstream::models::readers::parquet::load_parquet_table; + use lightstream::models::writers::chunked::parquet::ChunkedParquetWriter; + use lightstream::models::writers::parquet::write_parquet_table; + use lightstream::traits::chunked_table_reader::ChunkedTableReader; + use lightstream::traits::chunked_table_writer::ChunkedTableWriter; + + pub fn parquet_round_trips_through_the_storage_handles(base: &str) { + let url = format!("{base}/parquet/trades.parquet"); + let mut writer = StorageWriter::create(&url).unwrap(); + write_parquet_table(&sample_table(0, 100), &mut writer, None).unwrap(); + writer.flush().unwrap(); + + let table = load_parquet_table(StorageReader::open(&url).unwrap()).unwrap(); + assert_eq!(ids(&[table]), (0..100).collect::>()); + } + + pub fn chunked_parquet_round_trips_beneath_an_object_prefix(base: &str) { + let dir = format!("{base}/chunked-parquet/run"); + let mut writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + for first in [0, 10, 20] { + writer.write_chunk(&sample_table(first, 10)).unwrap(); + } + + let reader = ChunkedParquetReader::open(&dir, "trades", ()).unwrap(); + let tables: Vec
= reader.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..30).collect::>()); + } + + pub fn parallel_chunk_writes_reach_the_object_prefix(base: &str) { + let dir = format!("{base}/chunked-parquet/parallel"); + let writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + let tables: Vec
= (0..6).map(|i| sample_table(i * 5, 5)).collect(); + let refs: Vec<&Table> = tables.iter().collect(); + writer.par_write_all(&refs, Some(3)).unwrap(); + + let reader = ChunkedParquetReader::open(&dir, "trades", ()).unwrap(); + let tables: Vec
= reader.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..30).collect::>()); + } + + pub fn chunked_parquet_loads_as_a_super_table(base: &str) { + let dir = format!("{base}/chunked-parquet/super-table"); + let mut writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + for first in [0, 10, 20, 30] { + writer.write_chunk(&sample_table(first, 10)).unwrap(); + } + + let sequential = ChunkedParquetReader::open(&dir, "trades", ()) + .unwrap() + .load_batched() + .unwrap(); + let parallel = ChunkedParquetReader::par_load_batched(&dir, "trades", (), Some(3)).unwrap(); + for super_table in [sequential, parallel] { + assert_eq!(super_table.batches.len(), 4); + let batches: Vec
= super_table + .batches + .iter() + .map(|batch| (**batch).clone()) + .collect(); + assert_eq!(ids(&batches), (0..40).collect::>()); + } + } +} + +pub mod arrow_chunks { + use super::*; + use lightstream::models::readers::chunked::arrow::ChunkedArrowReader; + use lightstream::models::writers::chunked::arrow::ChunkedArrowWriter; + use lightstream::models::writers::chunked::partition::PartitionType; + use lightstream::traits::chunked_table_reader::ChunkedTableReader; + use lightstream::traits::chunked_table_writer::ChunkedTableWriter; + + /// 2026-02-09T07:00:00Z in epoch nanoseconds. + const SEVEN_AM: u64 = 1_770_620_400_000_000_000; + const ONE_HOUR: u64 = 3_600_000_000_000; + + pub fn chunked_arrow_round_trips_beneath_an_object_prefix(base: &str) { + let dir = format!("{base}/chunked-arrow/run"); + let mut writer = ChunkedArrowWriter::new(&dir, "trades").unwrap(); + for first in [0, 7, 14] { + writer.write_chunk(&sample_table(first, 7)).unwrap(); + } + + let reader = ChunkedArrowReader::open(&dir, "trades", ()).unwrap(); + let tables: Vec
= reader.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..21).collect::>()); + } + + pub fn auto_partitioned_chunks_round_trip_beneath_an_object_prefix(base: &str) { + let dir = format!("{base}/chunked-arrow/partitioned"); + let mut writer = ChunkedArrowWriter::new(&dir, "trades") + .unwrap() + .with_partition(PartitionType::Auto); + // Written out of time order: ids 10..20 at 08:00, then 0..10 at 07:00. + let eight = writer + .write_chunk_for_time(&sample_table(10, 10), SEVEN_AM + ONE_HOUR) + .unwrap(); + writer.write_chunk_for_time(&sample_table(0, 10), SEVEN_AM).unwrap(); + assert_eq!( + eight.to_str().unwrap(), + format!("{dir}/date=2026-02-09/hour=08/trades-0000000000.arrow") + ); + + let all = ChunkedArrowReader::open(&dir, "trades", ()).unwrap(); + let tables: Vec
= all.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..20).collect::>()); + + let hour = format!("{dir}/date=2026-02-09/hour=07"); + let seven = ChunkedArrowReader::open(&hour, "trades", ()).unwrap(); + let tables: Vec
= seven.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..10).collect::>()); + } +} + +#[cfg(feature = "csv")] +pub mod csv { + use super::*; + use lightstream::models::readers::chunked::csv::{ChunkedCsvReadOptions, ChunkedCsvReader}; + use lightstream::models::readers::csv::CsvReader; + use lightstream::models::writers::chunked::csv::ChunkedCsvWriter; + use lightstream::models::writers::csv::CsvWriter; + use lightstream::traits::chunked_table_reader::ChunkedTableReader; + use lightstream::traits::chunked_table_writer::ChunkedTableWriter; + + pub fn csv_round_trips_through_the_path_entry_points(base: &str) { + let url = format!("{base}/csv/trades.csv"); + let mut writer = CsvWriter::to_path(&url).unwrap(); + writer.write_table(&sample_table(0, 25)).unwrap(); + writer.flush().unwrap(); + + let mut reader = CsvReader::from_path(&url, Default::default(), 1_000).unwrap(); + let table = reader.next_batch().unwrap().expect("one batch"); + assert_eq!(table.n_rows, 25); + assert_eq!(ids(&[table]), (0..25).collect::>()); + } + + pub fn chunked_csv_round_trips_beneath_an_object_prefix(base: &str) { + let dir = format!("{base}/chunked-csv/run"); + let mut writer = ChunkedCsvWriter::new(&dir, "trades", Default::default()).unwrap(); + for first in [0, 4] { + writer.write_chunk(&sample_table(first, 4)).unwrap(); + } + + let reader = + ChunkedCsvReader::open(&dir, "trades", ChunkedCsvReadOptions::default()).unwrap(); + let tables: Vec
= reader.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..8).collect::>()); + } +} + +#[cfg(feature = "json")] +pub mod json { + use super::*; + use lightstream::models::decoders::json::JsonDecodeOptions; + use lightstream::models::encoders::json::{JsonEncodeOptions, JsonFormat}; + use lightstream::models::readers::json::JsonReader; + use lightstream::models::writers::json::JsonWriter; + + pub fn json_round_trips_through_the_path_entry_points(base: &str) { + let url = format!("{base}/json/trades.ndjson"); + let options = JsonEncodeOptions { + format: JsonFormat::Ndjson, + ..Default::default() + }; + let table = sample_table(0, 12); + let decode = JsonDecodeOptions { + schema: Some(schema_of(&table)), + ..Default::default() + }; + let mut writer = JsonWriter::to_path(&url, options).unwrap(); + writer.write_table(&table).unwrap(); + writer.flush().unwrap(); + + let mut reader = JsonReader::from_path(&url, JsonFormat::Ndjson, decode, 1_000).unwrap(); + let table = reader.next_batch().unwrap().expect("one batch"); + assert_eq!(ids(&[table]), (0..12).collect::>()); + } +} diff --git a/rust/tests/object_storage_roundtrip.rs b/rust/tests/object_storage_roundtrip.rs new file mode 100644 index 0000000..874b766 --- /dev/null +++ b/rust/tests/object_storage_roundtrip.rs @@ -0,0 +1,144 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! Object storage roundtrip integration tests against the in-memory store. +//! +//! Runs the shared object storage checks with `memory://` object URLs in +//! place of file paths. The in-memory store is shared across the process, +//! and each check writes beneath its own path to keep its objects separate. + +#![cfg(feature = "object_storage")] + +mod object_storage_checks; + +use lightstream::models::readers::object_storage::StorageReader; +use lightstream::object_storage::register_options; +use object_storage_checks as checks; + +const BASE: &str = "memory:///checks"; + +/// Run a synchronous check on the test thread, inside an entered runtime. +fn run(check: fn(&str)) { + let runtime = checks::runtime(); + let _entered = runtime.enter(); + check(BASE); +} + +#[test] +fn object_urls_outside_a_runtime_context_report_an_error() { + let err = StorageReader::open(format!("{BASE}/no-runtime/a.bin")) + .err() + .expect("no runtime context"); + assert!( + err.to_string().contains("Tokio runtime context"), + "got {err}" + ); +} + +#[test] +fn ipc_file_reader_reads_an_object_written_by_the_sync_writer() { + run(checks::ipc_file_reader_reads_an_object_written_by_the_sync_writer); +} + +#[tokio::test(flavor = "multi_thread")] +async fn ipc_file_reader_reads_an_object_written_by_the_async_writer() { + checks::ipc_file_reader_reads_an_object_written_by_the_async_writer(BASE).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn ipc_stream_round_trips_through_the_async_adapters() { + checks::ipc_stream_round_trips_through_the_async_adapters(BASE).await; +} + +#[test] +fn large_objects_upload_in_parts_and_read_back_whole() { + run(checks::large_objects_upload_in_parts_and_read_back_whole); +} + +#[test] +fn registered_upload_settings_apply_to_large_uploads() { + let base = "memory:///checks-one-part-at-a-time"; + register_options( + base, + &[("max_concurrent_parts", "1"), ("part_size", "6291456")], + ) + .unwrap(); + let runtime = checks::runtime(); + let _entered = runtime.enter(); + checks::large_objects_upload_in_parts_and_read_back_whole(base); +} + +#[test] +fn writes_after_an_object_completes_are_rejected() { + run(checks::writes_after_an_object_completes_are_rejected); +} + +#[test] +fn a_writer_dropped_before_flushing_leaves_no_object() { + run(checks::a_writer_dropped_before_flushing_leaves_no_object); +} + +#[test] +fn a_missing_object_reports_not_found() { + run(checks::a_missing_object_reports_not_found); +} + +#[test] +fn listing_returns_object_urls_immediately_beneath_the_prefix() { + run(checks::listing_returns_object_urls_immediately_beneath_the_prefix); +} + +#[cfg(feature = "parquet")] +#[test] +fn parquet_round_trips_through_the_storage_handles() { + run(checks::parquet::parquet_round_trips_through_the_storage_handles); +} + +#[cfg(feature = "parquet")] +#[test] +fn chunked_parquet_round_trips_beneath_an_object_prefix() { + run(checks::parquet::chunked_parquet_round_trips_beneath_an_object_prefix); +} + +#[cfg(feature = "parquet")] +#[test] +fn parallel_chunk_writes_reach_the_object_prefix() { + run(checks::parquet::parallel_chunk_writes_reach_the_object_prefix); +} + +#[cfg(feature = "parquet")] +#[test] +fn chunked_parquet_loads_as_a_super_table() { + run(checks::parquet::chunked_parquet_loads_as_a_super_table); +} + +#[test] +fn chunked_arrow_round_trips_beneath_an_object_prefix() { + run(checks::arrow_chunks::chunked_arrow_round_trips_beneath_an_object_prefix); +} + +#[test] +fn auto_partitioned_chunks_round_trip_beneath_an_object_prefix() { + run(checks::arrow_chunks::auto_partitioned_chunks_round_trip_beneath_an_object_prefix); +} + +#[cfg(feature = "csv")] +#[test] +fn csv_round_trips_through_the_path_entry_points() { + run(checks::csv::csv_round_trips_through_the_path_entry_points); +} + +#[cfg(feature = "csv")] +#[test] +fn chunked_csv_round_trips_beneath_an_object_prefix() { + run(checks::csv::chunked_csv_round_trips_beneath_an_object_prefix); +} + +#[cfg(feature = "json")] +#[test] +fn json_round_trips_through_the_path_entry_points() { + run(checks::json::json_round_trips_through_the_path_entry_points); +} diff --git a/rust/tests/object_storage_s3.rs b/rust/tests/object_storage_s3.rs new file mode 100644 index 0000000..dc1b13d --- /dev/null +++ b/rust/tests/object_storage_s3.rs @@ -0,0 +1,282 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! Object storage roundtrip integration tests against Amazon S3. +//! +//! Runs the shared object storage checks against a real bucket. The tests +//! are ignored by default because they need a bucket and credentials. Run +//! them with `scripts/run-object-storage-s3-tests.sh`, which exports the +//! AWS CLI's credentials and runs the ignored tests. +//! +//! ## Settings +//! +//! | Variable | Purpose | +//! |---|---| +//! | `LIGHTSTREAM_S3_URL` | Bucket and optional prefix, e.g. `s3://bucket/ci` | +//! | `AWS_REGION` | Bucket region | +//! | `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN` | Credentials | +//! | `AWS_ENDPOINT` | Optional endpoint for S3-compatible stores | +//! | `LIGHTSTREAM_S3_MAX_CONCURRENT_PARTS` | Optional `max_concurrent_parts` upload setting | +//! | `LIGHTSTREAM_S3_PART_SIZE` | Optional `part_size` upload setting, in bytes | +//! +//! Each check writes beneath its own run prefix, +//! `/lightstream-test--